Fix mkv generation by enhancing avc_find_start_code and thus correctly generating...

diff --git a/src/parsers/parser_avc.c b/src/parsers/parser_avc.c
index 54a08917aac7ee3a92e50cb0dea569c2a85e02f3..e5571470c29e61a0739e557e012cf34595d54eec 100644 (file)
--- a/src/parsers/parser_avc.c
+++ b/src/parsers/parser_avc.c
@@ -20,55 +20,77 @@
  */
 
 #include "parser_avc.h"
-
 static const uint8_t *
-avc_find_startcode(const uint8_t *p, const uint8_t *end)
+avc_find_startcode_internal(const uint8_t *p, const uint8_t *end)
 {
-  int i;
-       
-  uint32_t sc=0xFFFFFFFF;
-  size_t len = end - p;
-  for (i=0;i<len;i++)
-    {
-      sc = (sc <<8) | p[i];
-      if((sc & 0xffffff00) == 0x00000100) {
-       return p+i-3;
+  const uint8_t *a = p + 4 - ((intptr_t)p & 3);
+
+  for (end -= 3; p < a && p < end; p++) {
+    if (p[0] == 0 && p[1] == 0 && p[2] == 1)
+      return p;
+  }
+
+  for (end -= 3; p < end; p += 4) {
+    uint32_t x = *(const uint32_t*)p;
+//  if ((x - 0x01000100) & (~x) & 0x80008000) // little endian
+//  if ((x - 0x00010001) & (~x) & 0x00800080) // big endian
+    if ((x - 0x01010101) & (~x) & 0x80808080) { // generic
+      if (p[1] == 0) {
+       if (p[0] == 0 && p[2] == 1)
+         return p;
+       if (p[2] == 0 && p[3] == 1)
+         return p+1;
+      }
+      if (p[3] == 0) {
+       if (p[2] == 0 && p[4] == 1)
+         return p+2;
+       if (p[4] == 0 && p[5] == 1)
+         return p+3;
       }
-               
     }
-  return end;
+  }
+
+  for (end += 3; p < end; p++) {
+    if (p[0] == 0 && p[1] == 0 && p[2] == 1)
+      return p;
+  }
+
+  return end + 3;
 }
-static int
-avc_parse_nal_units(sbuf_t *sb, const uint8_t *buf_in, int size)
+
+static const uint8_t *
+avc_find_startcode(const uint8_t *p, const uint8_t *end)
+{
+    const uint8_t *out= avc_find_startcode_internal(p, end);
+    while(p<out && out<end && !out[-1]) out--;
+    return out;
+}
+
+static int avc_parse_nal_units(sbuf_t *sb, const uint8_t *buf_in, int size)
 {
   const uint8_t *p = buf_in;
   const uint8_t *end = p + size;
-  const uint8_t *nal_start, *nal_end; 
-
-  //printf("CONVERT SIZE %d\n", size);
+  const uint8_t *nal_start, *nal_end;
 
   size = 0;
   nal_start = avc_find_startcode(p, end);
-  while (nal_start < end) {
-    while(!*(nal_start++));
+  for (;;) {
+    while (nal_start < end && !*(nal_start++));
+    if (nal_start == end)
+      break;
+    
     nal_end = avc_find_startcode(nal_start, end);
-    /*printf("%4d bytes  %5d : %d\n", nal_end - nal_start,
-      nal_start - buf_in,
-      nal_end   - buf_in);*/
 
     int l = nal_end - nal_start;
 
-    if (l) {
-      sbuf_put_be32(sb, l);
-      sbuf_append(sb, nal_start, l);
-      size += 4 + l;
-    }
+    sbuf_put_be32(sb, l);
+    sbuf_append(sb, nal_start, l);
+    size += 4 + l;
     nal_start = nal_end;
   }
   return size;
 }
 
-
 static int
 avc_parse_nal_units_buf(const uint8_t *buf_in, uint8_t **buf, int *size)
 {
@@ -95,6 +117,7 @@ RB24(const uint8_t *d)
   return (d[0] << 16) | (d[1] << 8) | d[2];
 }
 
+#define FFMIN(a, b) ((a) > (b) ? (b) : (a))
 
 static int
 isom_write_avcc(sbuf_t *sb, const uint8_t *data, int len)
@@ -116,35 +139,37 @@ isom_write_avcc(sbuf_t *sb, const uint8_t *data, int len)
       end = buf + len;
 
       /* look for sps and pps */
-      while (buf < end) {
+      while (end - buf > 4) {
        unsigned int size;
        uint8_t nal_type;
-       size = RB32(buf);
-       nal_type = buf[4] & 0x1f;
-       if (nal_type == 7) { /* SPS */
+       size = FFMIN(RB32(buf), end - buf - 4);
+       buf += 4;
+       nal_type = buf[0] & 0x1f;
+       
+       if ((nal_type == 7) && (size >= 4) && (size <= UINT16_MAX)) { /* SPS */
          sps_array = realloc(sps_array,sizeof(uint8_t*)*(sps_count+1));
          sps_size_array = realloc(sps_size_array,sizeof(uint32_t)*(sps_count+1));
-         sps_array[sps_count] = buf + 4;
+         sps_array[sps_count] = buf;
          sps_size_array[sps_count] = size;
          sps_count++;
-       } else if (nal_type == 8) { /* PPS */
+       } else if ((nal_type == 8) && (size <= UINT16_MAX)) { /* PPS */
          pps_size_array = realloc(pps_size_array,sizeof(uint32_t)*(pps_count+1));
          pps_array = realloc(pps_array,sizeof (uint8_t*)*(pps_count+1));
-         pps_array[pps_count] = buf + 4;
+         pps_array[pps_count] = buf;
          pps_size_array[pps_count] = size;
          pps_count++;
        }
-       buf += size + 4;
+       buf += size;
       }
       if(!sps_count || !pps_count) {
        free(start);
        if (sps_count) {
          free(sps_array);
-          free(sps_size_array);
+         free(sps_size_array);
         }
        if (pps_count) {
          free(pps_array);
-          free(pps_size_array);
+         free(pps_size_array);
         }
        return -1;
       }
@@ -169,11 +194,11 @@ isom_write_avcc(sbuf_t *sb, const uint8_t *data, int len)
 
       if (sps_count) {
        free(sps_array);
-        free(sps_size_array);
+       free(sps_size_array);
       }
       if (pps_count) {
        free(pps_array);
-        free(pps_size_array);
+       free(pps_size_array);
       }
     } else {
       sbuf_append(sb, data, len);
@@ -182,9 +207,6 @@ isom_write_avcc(sbuf_t *sb, const uint8_t *data, int len)
   return 0;
 }
 
-
-
-
 th_pkt_t *
 avc_convert_pkt(th_pkt_t *src)
 {
@@ -217,4 +239,3 @@ avc_convert_pkt(th_pkt_t *src)
   pkt_ref_dec(src);
   return pkt;
 }
-