[FFmpeg-devel] [PATCH 2/3] avcodec/lagarith: Optimize case with singleton probability distribution

[Michael Niedermayer]

Fixes: Timeout
Fixes: 10554/clusterfuzz-testcase-minimized-ffmpeg_AV_CODEC_ID_LAGARITH_fuzzer-5739938067251200

Found-by: continuous fuzzing process https://github.com/google/oss-fuzz/tree/master/projects/ffmpeg
Signed-off-by: Michael Niedermayer <michael at niedermayer.cc>
---
 libavcodec/lagarith.c    | 36 ++++++++++++++++++++++++++++++++++++
 libavcodec/lagarithrac.h |  1 +
 2 files changed, 37 insertions(+)

diff --git a/libavcodec/lagarith.c b/libavcodec/lagarith.c
index 59169be5de..0b222e1acd 100644
--- a/libavcodec/lagarith.c
+++ b/libavcodec/lagarith.c
@@ -175,6 +175,7 @@ static int lag_read_prob_header(lag_rac *rac, GetBitContext *gb)
     if (nnz == 1 && (show_bits_long(gb, 32) & 0xFFFFFF)) {
         return AVERROR_INVALIDDATA;
     }
+    rac->nnz = nnz;
 
     /* Scale probabilities so cumulative probability is an even power of 2. */
     scale_factor = av_log2(cumul_prob);
@@ -332,6 +333,41 @@ static int lag_decode_line(LagarithContext *l, lag_rac *rac,
     if (!esc_count)
         esc_count = -1;
 
+    if (rac->nnz == 1) {
+        int v = -1;;
+handle_zeros1:
+        if (l->zeros_rem) {
+            int count = FFMIN(l->zeros_rem, width - i);
+            memset(dst + i, 0, count);
+            i += count;
+            l->zeros_rem -= count;
+        }
+
+        while (i < width) {
+            if (v < 0)
+                v =  lag_get_rac(rac);
+
+            dst[i] =v;
+            ret++;
+
+            if (v)
+                l->zeros = 0;
+            else
+                l->zeros++;
+
+            i++;
+            if (l->zeros == esc_count) {
+                ret++;
+
+                l->zeros = 0;
+
+                l->zeros_rem = lag_calc_zero_run(v);
+                goto handle_zeros1;
+            }
+        }
+        return ret;
+    }
+
     /* Output any zeros remaining from the previous run */
 handle_zeros:
     if (l->zeros_rem) {
diff --git a/libavcodec/lagarithrac.h b/libavcodec/lagarithrac.h
index ee836d01db..9f37f3939c 100644
--- a/libavcodec/lagarithrac.h
+++ b/libavcodec/lagarithrac.h
@@ -47,6 +47,7 @@ typedef struct lag_rac {
     const uint8_t *bytestream;        /**< Current position in input bytestream. */
     const uint8_t *bytestream_end;    /**< End position of input bytestream. */
 
+    int nnz;
     int overread;
 #define MAX_OVERREAD 4
 
-- 
2.20.1