handle extensions in opus_repacketizer_out_range_impl

6 files changed, 167 insertions, 61 deletions

diff --git a/Makefile.am b/Makefile.am
index 1b772446..d09c1771 100644
--- a/Makefile.am
+++ b/Makefile.am
@@ -202,9 +202,6 @@ tests_test_opus_api_LDADD = libopus.la $(NE10_LIBS) $(LIBM)
 tests_test_opus_encode_SOURCES = tests/test_opus_encode.c tests/opus_encode_regressions.c tests/test_opus_common.h
 tests_test_opus_encode_LDADD = libopus.la $(NE10_LIBS) $(LIBM)
 
-tests_test_opus_extensions_SOURCES = tests/test_opus_extensions.c tests/test_opus_common.h
-tests_test_opus_extensions_LDADD = $(NE10_LIBS) $(LIBM)
-
 tests_test_opus_decode_SOURCES = tests/test_opus_decode.c tests/test_opus_common.h
 tests_test_opus_decode_LDADD = libopus.la $(NE10_LIBS) $(LIBM)
 
@@ -219,6 +216,12 @@ SILK_OBJ = $(SILK_SOURCES:.c=.lo)
 LPCNET_OBJ = $(LPCNET_SOURCES:.c=.lo)
 OPUS_OBJ = $(OPUS_SOURCES:.c=.lo)
 
+tests_test_opus_extensions_SOURCES = tests/test_opus_extensions.c tests/test_opus_common.h
+tests_test_opus_extensions_LDADD = $(OPUS_OBJ) $(SILK_OBJ) $(LPCNET_OBJ) $(CELT_OBJ) $(NE10_LIBS) $(LIBM)
+if OPUS_ARM_EXTERNAL_ASM
+tests_test_opus_extensions_LDADD += libarmasm.la
+endif
+
 tests_test_opus_projection_SOURCES = tests/test_opus_projection.c tests/test_opus_common.h
 tests_test_opus_projection_LDADD = $(OPUS_OBJ) $(SILK_OBJ) $(LPCNET_OBJ) $(CELT_OBJ) $(NE10_LIBS) $(LIBM)
 if OPUS_ARM_EXTERNAL_ASM
diff --git a/src/opus_encoder.c b/src/opus_encoder.c
index d2f251f5..53c899a0 100644
--- a/src/opus_encoder.c
+++ b/src/opus_encoder.c
@@ -974,7 +974,7 @@ static opus_int32 encode_multiframe_packet(OpusEncoder *st,
    opus_int32 bytes_per_frame;
    opus_int32 cbr_bytes;
    opus_int32 repacketize_len;
-   int tmp_len, first_len;
+   int tmp_len;
    ALLOC_STACK;
 
    /* Worst cases:
@@ -1035,49 +1035,9 @@ static opus_int32 encode_multiframe_packet(OpusEncoder *st,
          RESTORE_STACK;
          return OPUS_INTERNAL_ERROR;
       }
-      if (first_frame) first_len = tmp_len;
    }
 
-   {
-      /* Handle extensions here. We only take the extensions from the first frame */
-      /* FIXME: handle extensions in a more general manner in OpusRepacketizer */
-      opus_int16 size[48];
-      const unsigned char *padding;
-      opus_int32 padding_len;
-      opus_int32 extensions_count;
-      VARDECL(opus_extension_data, extensions);
-      int extensions_alloc;
-
-      ret = opus_packet_parse_impl(tmp_data, first_len, 0, NULL, NULL, size,
-         NULL, NULL, &padding, &padding_len);
-      if (ret<0)
-      {
-         RESTORE_STACK;
-         return OPUS_INTERNAL_ERROR;
-      }
-      extensions_count = opus_packet_extensions_count(padding, padding_len);
-      if (extensions_count > 0)
-      {
-         extensions_alloc = extensions_count;
-      }
-      else
-      {
-         extensions_alloc = ALLOC_NONE;
-         extensions_count = 0;
-      }
-
-      ALLOC(extensions, extensions_alloc, opus_extension_data);
-      if (extensions_count > 0)
-      {
-         ret = opus_packet_extensions_parse(padding, padding_len, extensions, &extensions_count);
-         if (ret<0)
-         {
-            RESTORE_STACK;
-            return OPUS_INTERNAL_ERROR;
-         }
-      }
-      ret = opus_repacketizer_out_range_impl(rp, 0, nb_frames, data, repacketize_len, 0, !st->use_vbr, extensions, extensions_count);
-   }
+   ret = opus_repacketizer_out_range_impl(rp, 0, nb_frames, data, repacketize_len, 0, !st->use_vbr, NULL, 0);
    if (ret<0)
    {
       RESTORE_STACK;
diff --git a/src/opus_private.h b/src/opus_private.h
index 2f8f4f43..ecadbc77 100644
--- a/src/opus_private.h
+++ b/src/opus_private.h
@@ -42,6 +42,8 @@ struct OpusRepacketizer {
    const unsigned char *frames[48];
    opus_int16 len[48];
    int framesize;
+   const unsigned char *paddings[48];
+   opus_int32 padding_len[48];
 };
 
 typedef struct {
diff --git a/src/repacketizer.c b/src/repacketizer.c
index cea88e3d..c56f9267 100644
--- a/src/repacketizer.c
+++ b/src/repacketizer.c
@@ -32,6 +32,7 @@
 #include "opus.h"
 #include "opus_private.h"
 #include "os_support.h"
+#include "stack_alloc.h"
 
 
 int opus_repacketizer_get_size(void)
@@ -82,10 +83,18 @@ static int opus_repacketizer_cat_impl(OpusRepacketizer *rp, const unsigned char
       return OPUS_INVALID_PACKET;
    }
 
-   ret=opus_packet_parse_impl(data, len, self_delimited, &tmp_toc, &rp->frames[rp->nb_frames], &rp->len[rp->nb_frames], NULL, NULL, NULL, NULL);
+   ret=opus_packet_parse_impl(data, len, self_delimited, &tmp_toc, &rp->frames[rp->nb_frames], &rp->len[rp->nb_frames],
+       NULL, NULL, &rp->paddings[rp->nb_frames], &rp->padding_len[rp->nb_frames]);
    if(ret<1)return ret;
 
-   rp->nb_frames += curr_nb_frames;
+   /* set padding length to zero for all but the first frame */
+   while (curr_nb_frames > 1)
+   {
+      rp->nb_frames++;
+      rp->padding_len[rp->nb_frames] = 0;
+      curr_nb_frames--;
+   }
+   rp->nb_frames++;
    return OPUS_OK;
 }
 
@@ -109,10 +118,14 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
    unsigned char * ptr;
    int ones_begin=0, ones_end=0;
    int ext_begin=0, ext_len=0;
+   int ext_count, total_ext_count;
+   VARDECL(opus_extension_data, all_extensions);
+   ALLOC_STACK;
 
    if (begin<0 || begin>=end || end>rp->nb_frames)
    {
       /*fprintf(stderr, "%d %d %d\n", begin, end, rp->nb_frames);*/
+      RESTORE_STACK;
       return OPUS_BAD_ARG;
    }
    count = end-begin;
@@ -124,13 +137,50 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
    else
       tot_size = 0;
 
+   /* figure out total number of extensions */
+   total_ext_count = nb_extensions;
+   for (i=begin;i<end;i++)
+   {
+      int n = opus_packet_extensions_count(rp->paddings[i], rp->padding_len[i]);
+      if (n > 0) total_ext_count += n;
+   }
+   ALLOC(all_extensions, total_ext_count ? total_ext_count : ALLOC_NONE, opus_extension_data);
+   /* copy over any extensions that were passed in */
+   for (ext_count=0;ext_count<nb_extensions;ext_count++)
+   {
+      all_extensions[ext_count] = extensions[ext_count];
+   }
+
+   /* incorporate any extensions from the repacketizer padding */
+   for (i=begin;i<end;i++)
+   {
+      int frame_ext_count, j;
+      frame_ext_count = total_ext_count - ext_count;
+      int ret = opus_packet_extensions_parse(rp->paddings[i], rp->padding_len[i],
+         &all_extensions[ext_count], &frame_ext_count);
+      if (ret<0)
+      {
+         RESTORE_STACK;
+         return OPUS_INTERNAL_ERROR;
+      }
+      /* renumber the extension frame numbers */
+      for (j=0;j<frame_ext_count;j++)
+      {
+         all_extensions[ext_count+j].frame += i-begin;
+      }
+      ext_count += frame_ext_count;
+   }
+
    ptr = data;
    if (count==1)
    {
       /* Code 0 */
       tot_size += len[0]+1;
       if (tot_size > maxlen)
+      {
+         RESTORE_STACK;
          return OPUS_BUFFER_TOO_SMALL;
+      }
       *ptr++ = rp->toc&0xFC;
    } else if (count==2)
    {
@@ -139,18 +189,24 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
          /* Code 1 */
          tot_size += 2*len[0]+1;
          if (tot_size > maxlen)
+         {
+            RESTORE_STACK;
             return OPUS_BUFFER_TOO_SMALL;
+         }
          *ptr++ = (rp->toc&0xFC) | 0x1;
       } else {
          /* Code 2 */
          tot_size += len[0]+len[1]+2+(len[0]>=252);
          if (tot_size > maxlen)
+         {
+            RESTORE_STACK;
             return OPUS_BUFFER_TOO_SMALL;
+         }
          *ptr++ = (rp->toc&0xFC) | 0x2;
          ptr += encode_size(len[0], ptr);
       }
    }
-   if (count > 2 || (pad && tot_size < maxlen) || nb_extensions > 0)
+   if (count > 2 || (pad && tot_size < maxlen) || ext_count > 0)
    {
       /* Code 3 */
       int vbr;
@@ -179,20 +235,27 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
          tot_size += len[count-1];
 
          if (tot_size > maxlen)
+         {
+            RESTORE_STACK;
             return OPUS_BUFFER_TOO_SMALL;
+         }
          *ptr++ = (rp->toc&0xFC) | 0x3;
          *ptr++ = count | 0x80;
       } else {
          tot_size += count*len[0]+2;
          if (tot_size > maxlen)
+         {
+            RESTORE_STACK;
             return OPUS_BUFFER_TOO_SMALL;
+         }
          *ptr++ = (rp->toc&0xFC) | 0x3;
          *ptr++ = count;
       }
       pad_amount = pad ? (maxlen-tot_size) : 0;
-      if (nb_extensions>0)
+      if (ext_count>0)
       {
-         ext_len = opus_packet_extensions_generate(NULL, maxlen-tot_size, extensions, nb_extensions, 0);
+         /* figure out how much space we need for the extensions */
+         ext_len = opus_packet_extensions_generate(NULL, maxlen-tot_size, all_extensions, ext_count, 0);
          if (ext_len < 0) return ext_len;
          if (!pad)
             pad_amount = ext_len + ext_len/254 + 1;
@@ -203,7 +266,10 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
          data[1] |= 0x40;
          nb_255s = (pad_amount-1)/255;
          if (tot_size + ext_len + nb_255s + 1 > maxlen)
+         {
+            RESTORE_STACK;
             return OPUS_BUFFER_TOO_SMALL;
+         }
          ext_begin = tot_size+pad_amount-ext_len;
          /* Prepend 0x01 padding */
          ones_begin = tot_size+nb_255s+1;
@@ -234,12 +300,12 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
       ptr += len[i];
    }
    if (ext_len > 0) {
-      int ret = opus_packet_extensions_generate(&data[ext_begin], ext_len, extensions, nb_extensions, 0);
+      int ret = opus_packet_extensions_generate(&data[ext_begin], ext_len, all_extensions, ext_count, 0);
       celt_assert(ret == ext_len);
    }
    for (i=ones_begin;i<ones_end;i++)
       data[i] = 0x01;
-   if (pad && nb_extensions==0)
+   if (pad && ext_count==0)
    {
       /* Fill padding with zeros. */
       while (ptr<data+maxlen)
diff --git a/tests/meson.build b/tests/meson.build
index 5f78985f..18b6fafc 100644
--- a/tests/meson.build
+++ b/tests/meson.build
@@ -19,7 +19,7 @@ foreach t : opus_tests
 
   exe_kwargs = {}
   # This test uses private symbols
-  if test_name == 'test_opus_projection'
+  if test_name == 'test_opus_projection' or test_name == 'test_opus_extensions'
     exe_kwargs = {
       'link_with': [celt_lib, silk_lib, dnn_lib],
       'objects': opus_lib.extract_all_objects(),
diff --git a/tests/test_opus_extensions.c b/tests/test_opus_extensions.c
index 430f6714..c5112a09 100644
--- a/tests/test_opus_extensions.c
+++ b/tests/test_opus_extensions.c
@@ -40,16 +40,9 @@
 #define getpid _getpid
 #endif
 
-/* including sources directly to test internal APIs */
-#define CELT_C /* to make celt_assert work */
-#include "../src/extensions.c"
+#include "../src/opus_private.h"
 #include "test_opus_common.h"
 
-const char *opus_get_version_string(void)
-{
-   return "unknown";
-}
-
 void test_extensions_generate_success(void)
 {
    static const opus_extension_data ext[] = {
@@ -343,6 +336,87 @@ void test_random_extensions_parse(void)
    }
 }
 
+void test_opus_repacketizer_out_range_impl(void)
+{
+   OpusRepacketizer rp;
+   unsigned char packet[1024];
+   unsigned char packet_out[1024];
+   opus_int16 size[48];
+   const unsigned char *padding;
+   opus_int32 padding_len;
+   opus_extension_data ext_out[10];
+   int nb_ext;
+   int res, len;
+   int first_count = 0, second_count = 0;
+   static const opus_extension_data ext[] = {
+      {33, 0, (const unsigned char *)"abcdefg", 7},
+      {100, 0, (const unsigned char *)"uvwxyz", 6},
+   };
+
+   opus_repacketizer_init(&rp);
+
+   memset(packet, 0, sizeof(packet));
+   /* Hybrid Packet with 20 msec frames, Code 3 */
+   packet[0] = (15 << 3) | 3;
+   /* Code 3, padding bit set, 1 frame */
+   packet[1] = 1 << 6 | 1;
+   packet[2] = 0;
+   packet[3] = 0;
+
+   /* generate 2 extensions, id 33 and 100 */
+   len = opus_packet_extensions_generate(&packet[4], sizeof(packet)-4, ext, 2, 0);
+   /* update the padding length */
+   packet[2] = len;
+
+   /* concatenate 3 frames */
+   res = opus_repacketizer_cat(&rp, packet, 4+len);
+   /* for the middle frame, no padding, no extensions */
+   packet[1] = 1;
+   res = opus_repacketizer_cat(&rp, packet, 4);
+   /* switch back to extensions for the last frame extensions */
+   packet[1] = 1 << 6 | 1;
+   res = opus_repacketizer_cat(&rp, packet, 4+len);
+
+   expect_true(rp.nb_frames == 3, "Expected 3 frames");
+   res = opus_repacketizer_out_range_impl(&rp,
+      0, 3, /* begin, end */
+      packet_out, /* unsigned char *data */
+      sizeof(packet_out), /* opus_int32 maxlen */
+      0, /*int self_delimited */
+      0, /* int pad */
+      NULL, /* const opus_extension_data *extensions */
+      0 /* int nb_extensions */);
+   expect_true(res > 0, "expected valid packet length");
+
+   /* now verify that we have the expected extensions */
+   res = opus_packet_parse_impl(packet_out, res, 0, NULL, NULL, size,
+      NULL, NULL, &padding, &padding_len);
+   nb_ext = 10;
+   res = opus_packet_extensions_parse(padding, padding_len, ext_out, &nb_ext);
+   expect_true(nb_ext == 4, "Expected 4 extensions");
+   for (int i = 0 ; i < nb_ext; i++)
+   {
+      if (ext_out[i].id == 33)
+      {
+         opus_test_assert(ext_out[i].len == ext[0].len);
+         opus_test_assert(0 == memcmp(ext_out[i].data, ext[0].data, ext[0].len));
+         first_count++;
+      }
+      else if (ext_out[i].id == 100)
+      {
+         opus_test_assert(ext_out[i].len == ext[1].len);
+         opus_test_assert(0 == memcmp(ext_out[i].data, ext[1].data, ext[1].len));
+         second_count++;
+      }
+      if (i < 2)
+         opus_test_assert(ext_out[i].frame == 0)
+      else
+         opus_test_assert(ext_out[i].frame == 2)
+   }
+   opus_test_assert(first_count == 2);
+   opus_test_assert(second_count == 2);
+}
+
 int main(int argc, char **argv)
 {
    int env_used;
@@ -369,6 +443,7 @@ int main(int argc, char **argv)
    test_extensions_parse_zero();
    test_extensions_parse_fail();
    test_random_extensions_parse();
+   test_opus_repacketizer_out_range_impl();
    fprintf(stderr,"Tests completed successfully.\n");
    return 0;
 }