@@ -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_31(&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_31(&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_31(&gb);

+                if (sps_id > MAX_SPS_COUNT)

+                    return 0;

+                if (!sps_ids[sps_id])

+                    break;

+                pps_ids[pps_id] = 1;

                 pps++;

                 break;

             }