[FFmpeg-devel] [PATCH V4] lavc/golomb: Fix UE golomb overwrite issue.

[Jun Zhao]

V4: Fix rang check error in assert base on Mark's review
V3: Clean the code logic base on Michael's review.
V2: Add set_ue_golomb_long() to support 32bits UE golomb and update the unit test.
-------------- next part --------------
From fa3f59e5fcb2cddcc44b0e895bfa02caa491fee5 Mon Sep 17 00:00:00 2001
From: Jun Zhao <jun.zhao at intel.com>
Date: Fri, 2 Jun 2017 15:05:49 +0800
Subject: [PATCH V4] lavc/golomb: Fix UE golomb overwrite issue.

put_bits just support write up to 31 bits, when write 32 bit in
put_bits, it's will overwrite the bit buffer, because the default
assert level is 0, the av_assert2(n <= 31 && value < (1U << n))
in put_bits can not be trigger runtime. Add set_ue_golomb_long()
to support 32bits UE golomb.

Signed-off-by: Jun Zhao <jun.zhao at intel.com>
---
 libavcodec/golomb.h       | 17 ++++++++++++++++-
 libavcodec/put_bits.h     | 35 +++++++++++++++++++++++++++++++++++
 libavcodec/tests/golomb.c | 19 +++++++++++++++++++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/libavcodec/golomb.h b/libavcodec/golomb.h
index 0833aff468..52bb88daa7 100644
--- a/libavcodec/golomb.h
+++ b/libavcodec/golomb.h
@@ -458,7 +458,7 @@ static inline int get_te(GetBitContext *s, int r, char *file, const char *func,
 #endif /* TRACE */
 
 /**
- * write unsigned exp golomb code.
+ * write unsigned exp golomb code. 2^16-2 at most.
  */
 static inline void set_ue_golomb(PutBitContext *pb, int i)
 {
@@ -473,6 +473,21 @@ static inline void set_ue_golomb(PutBitContext *pb, int i)
 }
 
 /**
+ * write unsigned exp golomb code. 2^32-2 at most.
+ */
+static inline void set_ue_golomb_long(PutBitContext *pb, uint32_t i)
+{
+    av_assert2(i <= (UINT32_MAX - 1));
+
+    if (i < 256)
+        put_bits(pb, ff_ue_golomb_len[i], i + 1);
+    else {
+        int e = av_log2(i + 1);
+        put_bits64(pb, 2 * e + 1, i + 1);
+    }
+}
+
+/**
  * write truncated unsigned exp golomb code.
  */
 static inline void set_te_golomb(PutBitContext *pb, int i, int range)
diff --git a/libavcodec/put_bits.h b/libavcodec/put_bits.h
index 68ed391195..132d07558e 100644
--- a/libavcodec/put_bits.h
+++ b/libavcodec/put_bits.h
@@ -221,6 +221,41 @@ static void av_unused put_bits32(PutBitContext *s, uint32_t value)
 }
 
 /**
+ * Write up to 64 bits into a bitstream.
+ */
+static inline void put_bits64(PutBitContext *s, int n, uint64_t value)
+{
+    av_assert2((n == 64) || (n < 64 && value < (UINT64_C(1) << n)));
+
+    if (n < 32)
+        put_bits(s, n, value);
+    else if (n == 32)
+        put_bits32(s, value);
+    else if (n < 64) {
+        uint32_t lo = value & 0xffffffff;
+        uint32_t hi = value >> 32;
+#ifdef BITSTREAM_WRITER_LE
+        put_bits32(s, lo);
+        put_bits(s, n - 32, hi);
+#else
+        put_bits(s, n - 32, hi);
+        put_bits32(s, lo);
+#endif
+    } else {
+        uint32_t lo = value & 0xffffffff;
+        uint32_t hi = value >> 32;
+#ifdef BITSTREAM_WRITER_LE
+        put_bits32(s, lo);
+        put_bits32(s, hi);
+#else
+        put_bits32(s, hi);
+        put_bits32(s, lo);
+#endif
+
+    }
+}
+
+/**
  * Return the pointer to the byte where the bitstream writer will put
  * the next bit.
  */
diff --git a/libavcodec/tests/golomb.c b/libavcodec/tests/golomb.c
index 965367b7be..85b8a9390b 100644
--- a/libavcodec/tests/golomb.c
+++ b/libavcodec/tests/golomb.c
@@ -21,6 +21,7 @@
 #include <stdint.h>
 #include <stdio.h>
 
+#include "libavutil/internal.h"
 #include "libavutil/mem.h"
 
 #include "libavcodec/get_bits.h"
@@ -76,6 +77,24 @@ int main(void)
         }
     }
 
+#define EXTEND_L(i) ((i) << 4 | (i) & 15)
+    init_put_bits(&pb, temp, SIZE);
+    for (i = 0; i < COUNT; i++)
+        set_ue_golomb_long(&pb, EXTEND_L(i));
+    flush_put_bits(&pb);
+
+    init_get_bits(&gb, temp, 8 * SIZE);
+    for (i = 0; i < COUNT; i++) {
+        int j, s = show_bits_long(&gb, 32);
+
+        j = get_ue_golomb_long(&gb);
+        if (j != EXTEND_L(i)) {
+            fprintf(stderr, "get_ue_golomb_long: expected %d, got %d. "
+                    "bits: %8x\n", EXTEND_L(i), j, s);
+            ret = 1;
+        }
+    }
+
     init_put_bits(&pb, temp, SIZE);
     for (i = 0; i < COUNT; i++)
         set_se_golomb(&pb, i - COUNT / 2);
-- 
2.11.0