[vlc-devel] [PATCH] h264_nal: fix parsing of the end of SPS and PPS NAL units

[Thomas Guillem]

SPS and PPS NAL units always start with "00 00 00 01" but others NAL units may
starts with "00 00 01".
---
 modules/packetizer/h264_nal.c | 91 +++++++++++++++++++++++++++----------------
 1 file changed, 57 insertions(+), 34 deletions(-)

diff --git a/modules/packetizer/h264_nal.c b/modules/packetizer/h264_nal.c
index a35f861..a867dc9 100644
--- a/modules/packetizer/h264_nal.c
+++ b/modules/packetizer/h264_nal.c
@@ -20,9 +20,10 @@
 
 #include "h264_nal.h"
 
+#include <assert.h>
 #include <limits.h>
 
-static const uint8_t annexb_startcode[] = { 0x00, 0x00, 0x00, 0x01 };
+static const uint8_t annexb_startcode[] = { 0x00, 0x00, 0x01 };
 
 int convert_sps_pps( decoder_t *p_dec, const uint8_t *p_buf,
                      uint32_t i_buf_size, uint8_t *p_out_buf,
@@ -140,6 +141,15 @@ void convert_h264_to_annexb( uint8_t *p_buf, uint32_t i_len,
     }
 }
 
+static size_t get_spspps_size( uint8_t *p_start, uint8_t *p_end )
+{
+    /* cf B.1.1: don't count trailing_zero_8bits */
+
+    while( *p_end == 0 && p_end > p_start )
+        p_end--;
+    return p_end - p_start + 1;
+}
+
 int h264_get_spspps( uint8_t *p_buf, size_t i_buf,
                      uint8_t **pp_sps, size_t *p_sps_size,
                      uint8_t **pp_pps, size_t *p_pps_size )
@@ -148,57 +158,70 @@ int h264_get_spspps( uint8_t *p_buf, size_t i_buf,
     size_t i_sps_size = 0, i_pps_size = 0;
     int i_nal_type = NAL_UNKNOWN;
 
-    if (unlikely(p_buf == NULL || i_buf == 0))
-        return -1;
+    assert( p_buf );
 
-    while( true )
+    while( i_buf > 0 )
     {
-        int i_inc = 0;
-
-        if( i_buf > 5 && memcmp( p_buf, annexb_startcode, 4 ) == 0 )
+        /* SPS and PPS NAL units always start with "00 00 00 01" but other NAL
+         * units may start with "00 00 01". */
+        if( i_buf > 5 && memcmp( p_buf + 1, annexb_startcode, 3 ) == 0 )
         {
             i_nal_type = p_buf[4] & 0x1F;
 
-            /* size of startcode + nal_type */
-            i_inc = 5;
+            /* The start prefix is always "00 00 00 01" for SPS and PPS */
+            if( ( i_nal_type == NAL_SPS || i_nal_type == NAL_PPS )
+             && p_buf[0] != 0 )
+                return -1;
 
             /* pointer to the beginning of the sps/pps  */
             if( i_nal_type == NAL_SPS && !p_sps )
                 p_sps = p_buf;
             if( i_nal_type == NAL_PPS && !p_pps )
                 p_pps = p_buf;
-        } else {
-            i_inc = 1;
-        }
 
-        /* cf. 7.4.1.2.3 */
-        if( i_nal_type == NAL_UNKNOWN || i_nal_type > 18
-         || (i_nal_type >= 10 && i_nal_type <= 12))
-            return -1;
+            /* cf. 7.4.1.2.3 */
+            if( i_nal_type == NAL_UNKNOWN || i_nal_type > 18
+             || ( i_nal_type >= 10 && i_nal_type <= 12 ) )
+                return -1;
 
-        /* update SPS/PPS size if the new NAL is different than the last one */
-        if( !i_sps_size && p_sps && i_nal_type != NAL_SPS )
-            i_sps_size = p_buf - p_sps;
-        if( !i_pps_size && p_pps && i_nal_type != NAL_PPS )
-            i_pps_size = p_buf - p_pps;
+            /* update SPS/PPS size if the new NAL is different than the last
+             * one */
+            if( !i_sps_size && p_sps && i_nal_type != NAL_SPS )
+                i_sps_size = get_spspps_size( p_sps, p_buf );
+            if( !i_pps_size && p_pps && i_nal_type != NAL_PPS )
+                i_pps_size = get_spspps_size( p_pps, p_buf );
 
-        /* SPS/PPS are before the slices */
-        if (i_nal_type >= NAL_SLICE && i_nal_type <= NAL_SLICE_IDR )
-            break;
+            /* Found a SPS and a PPS */
+            if( i_sps_size && i_pps_size )
+                break;
 
-        i_buf -= i_inc;
-        p_buf += i_inc;
+            /* SPS/PPS are before the slices */
+            if ( i_nal_type >= NAL_SLICE && i_nal_type <= NAL_SLICE_IDR )
+                break;
 
-        if( i_buf == 0 )
-        {
-            /* update SPS/PPS size if we reach the end of the buffer */
-            if( !i_sps_size && p_sps )
-                i_sps_size = p_buf - p_sps;
-            if( !i_pps_size && p_pps )
-                i_pps_size = p_buf - p_pps;
-            break;
+            /* size of startcode + nal_type */
+            i_buf -= 5;
+            p_buf += 5;
+        } else {
+            /* cf B.1: Don't abort if there are leading_zero_8bits (only for
+             * the first NAL unit of the bytestream). */
+            if( i_nal_type == NAL_UNKNOWN && p_buf[0] != 0 )
+                return -1;
+
+            i_buf--;
+            p_buf++;
         }
     }
+
+    if( i_buf == 0 )
+    {
+        /* update SPS/PPS size if we reach the end of the bytestream */
+        if( !i_sps_size && p_sps )
+            i_sps_size = p_buf - p_sps;
+        if( !i_pps_size && p_pps )
+            i_pps_size = p_buf - p_pps;
+    }
+
     if( ( !p_sps || !i_sps_size ) && ( !p_pps || !i_pps_size ) )
         return -1;
     *pp_sps = p_sps;
-- 
2.1.4