X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=libavformat%2Fh264dec.c;h=85d7163c96d0595c377721d17d6b4902ec14ef02;hb=4888932c4d3ac68c26f44d7adf1e22e0d7ed8808;hp=d1d37ca88aa20c7459aa7f267b435b2cddcc5a94;hpb=c18535399d2597afd8855ccaa4416b0999d0d7f9;p=ffmpeg

diff --git a/libavformat/h264dec.c b/libavformat/h264dec.c
index d1d37ca88aa..85d7163c96d 100644
--- a/libavformat/h264dec.c
+++ b/libavformat/h264dec.c
@@ -19,17 +19,26 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  */
 
+#include "libavcodec/get_bits.h"
+#include "libavcodec/golomb.h"
 #include "avformat.h"
 #include "rawdec.h"
 #include "libavcodec/internal.h"
 
+#define MAX_SPS_COUNT          32
+#define MAX_PPS_COUNT         256
+
 static int h264_probe(AVProbeData *p)
 {
     uint32_t code = -1;
     int sps = 0, pps = 0, idr = 0, res = 0, sli = 0;
-    int i;
+    int i, ret;
+    int pps_ids[MAX_PPS_COUNT+1] = {0};
+    int sps_ids[MAX_SPS_COUNT+1] = {0};
+    unsigned pps_id, sps_id;
+    GetBitContext gb;
 
-    for (i = 0; i < p->buf_size; i++) {
+    for (i = 0; i + 2 < p->buf_size; i++) {
         code = (code << 8) + p->buf[i];
         if ((code & 0xffffff00) == 0x100) {
             int ref_idc = (code >> 5) & 3;
@@ -53,19 +62,48 @@ static int h264_probe(AVProbeData *p)
                     res++;
             }
 
+            ret = init_get_bits8(&gb, p->buf + i + 1, p->buf_size - i - 1);
+            if (ret < 0)
+                return 0;
+
             switch (type) {
             case 1:
-                sli++;
-                break;
             case 5:
-                idr++;
+                get_ue_golomb_long(&gb);
+                if (get_ue_golomb_long(&gb) > 9U)
+                    return 0;
+                pps_id = get_ue_golomb_long(&gb);
+                if (pps_id > MAX_PPS_COUNT)
+                    return 0;
+                if (!pps_ids[pps_id])
+                    break;
+
+                if (type == 1)
+                    sli++;
+                else
+                    idr++;
                 break;
             case 7:
-                if (p->buf[i + 2] & 0x03)
+                skip_bits(&gb, 14);
+                if (get_bits(&gb, 2))
+                    return 0;
+                skip_bits(&gb, 8);
+                sps_id = get_ue_golomb_long(&gb);
+                if (sps_id > MAX_SPS_COUNT)
                     return 0;
+                sps_ids[sps_id] = 1;
                 sps++;
                 break;
             case 8:
+                pps_id = get_ue_golomb_long(&gb);
+                if (pps_id > MAX_PPS_COUNT)
+                    return 0;
+                sps_id = get_ue_golomb_long(&gb);
+                if (sps_id > MAX_SPS_COUNT)
+                    return 0;
+                if (!sps_ids[sps_id])
+                    break;
+                pps_ids[pps_id] = 1;
                 pps++;
                 break;
             }