Experimentng with padding extensions

6 files changed, 326 insertions, 11 deletions

diff --git a/opus_sources.mk b/opus_sources.mk
index 44153b57..91589efa 100644
--- a/opus_sources.mk
+++ b/opus_sources.mk
@@ -2,6 +2,7 @@ OPUS_SOURCES = \
 src/opus.c \
 src/opus_decoder.c \
 src/opus_encoder.c \
+src/extensions.c \
 src/opus_multistream.c \
 src/opus_multistream_encoder.c \
 src/opus_multistream_decoder.c \
diff --git a/src/extensions.c b/src/extensions.c
new file mode 100644
index 00000000..b7bafb99
--- /dev/null
+++ b/src/extensions.c
@@ -0,0 +1,288 @@
+/* Copyright (c) 2022 Amazon */
+/*
+   Redistribution and use in source and binary forms, with or without
+   modification, are permitted provided that the following conditions
+   are met:
+
+   - Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+
+   - Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in the
+   documentation and/or other materials provided with the distribution.
+
+   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+   ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+   A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
+   OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+
+#ifdef HAVE_CONFIG_H
+#include "config.h"
+#endif
+
+
+#include "opus_types.h"
+#include "opus_defines.h"
+#include "arch.h"
+#include "os_support.h"
+#include "opus_private.h"
+
+
+/* Given an extension payload, advance data to the next extension and return the
+   length of the remaining extensions. */
+static opus_int32 skip_extension(const unsigned char **data, opus_int32 len, opus_int32 *header_size)
+{
+   int id, L;
+   if (len==0)
+      return 0;
+   id = **data>>1;
+   L = **data&1;
+   if (id == 0 && L == 1)
+   {
+      *header_size = 1;
+      if (len < 1)
+         return -1;
+      (*data)++;
+      len--;
+      return len;
+   } else if (id > 0 && id < 32)
+   {
+      if (len < 1+L)
+         return -1;
+      *data += 1+L;
+      len -= 1+L;
+      *header_size = 1;
+      return len;
+   } else {
+      if (L==0)
+      {
+         *data += len;
+         *header_size = 1;
+         return 0;
+      } else {
+         opus_int32 bytes=0;
+         *header_size = 1;
+         do {
+            (*data)++;
+            len--;
+            if (len == 0)
+               return -1;
+            bytes += **data;
+            (*header_size)++;
+         } while (**data == 255);
+         (*data)++;
+         len--;
+         if (bytes <= len)
+         {
+            len -= bytes;
+            *data += bytes;
+         } else {
+            return -1;
+         }
+         
+         return len;
+      }
+   }
+}
+
+/* Count the number of extensions, excluding real padding and separators. */
+opus_int32 opus_packet_extensions_count(const unsigned char *data, opus_int32 len)
+{
+   opus_int32 curr_len;
+   opus_int32 count=0;
+   const unsigned char *curr_data = data;
+   curr_len = len;
+   while (curr_len > 0)
+   {
+      int id;
+      opus_int32 header_size;
+      id = *curr_data>>1;
+      if (id > 1)
+         count++;
+      curr_len = skip_extension(&curr_data, curr_len, &header_size);
+   }
+   return count;
+}
+
+/* Extract extensions from Opus padding (excluding real padding and separators) */
+opus_int32 opus_packet_extensions_parse(const unsigned char *data, opus_int32 len, opus_extension_data *extensions, opus_int32 *nb_extensions)
+{
+   const unsigned char *curr_data;
+   opus_int32 curr_len;
+   int curr_frame=0;
+   opus_int32 count=0;
+
+   curr_data = data;
+   curr_len = len;
+   while (curr_len > 0)
+   {
+      int id;
+      opus_int32 header_size;
+      opus_extension_data curr_ext;
+      id = *curr_data>>1;
+      if (id > 1)
+      {
+         curr_ext.id = id;
+         curr_ext.frame = curr_frame;
+         curr_ext.data = curr_data;
+      } else if (id == 1)
+      {
+         int L = *curr_data&1;
+         if (L==0)
+            curr_frame++;
+         else {
+            if (curr_len >= 2)
+               curr_frame += curr_data[1];
+            /* Else we're at the end and it doesn't matter. */
+         }
+         if (curr_frame >= 48)
+         {
+            *nb_extensions = count;
+            return OPUS_INVALID_PACKET;
+         }
+      }
+      curr_len = skip_extension(&curr_data, curr_len, &header_size);
+      /* printf("curr_len = %d, header_size = %d\n", curr_len, header_size); */
+      if (curr_len < 0)
+      {
+         *nb_extensions = count;
+         return OPUS_INVALID_PACKET;
+      }
+      celt_assert(curr_data - data == len - curr_len);
+      if (id > 1)
+      {
+         if (count == *nb_extensions)
+         {
+             return OPUS_BUFFER_TOO_SMALL;
+         }
+         curr_ext.len = curr_data - curr_ext.data - header_size;
+         curr_ext.data += header_size;
+         extensions[count++] = curr_ext;
+      }
+   }
+   celt_assert(curr_len == 0);
+   *nb_extensions = count;
+   return OPUS_OK;
+}
+
+opus_int32 opus_packet_extensions_generate(unsigned char *data, opus_int32 len, const opus_extension_data  *extensions, opus_int32 nb_extensions, int pad)
+{
+   int max_frame=0;
+   opus_int32 i;
+   int frame;
+   int curr_frame = 0;
+   opus_int32 pos = 0;
+   opus_int32 written = 0;
+   for (i=0;i<nb_extensions;i++)
+   {
+      max_frame = IMAX(max_frame, extensions[i].frame);
+      if (extensions[i].id < 2)
+         return OPUS_BAD_ARG;
+   }
+   if (max_frame >= 48) return OPUS_BAD_ARG;
+   for (frame=0;frame<=max_frame;frame++)
+   {
+      for (i=0;i<nb_extensions;i++)
+      {
+         if (extensions[i].frame == frame)
+         {
+            /* Insert separator when needed. */
+            if (frame != curr_frame) {
+               int diff = frame - curr_frame;
+               if (len-pos < 2)
+                  return OPUS_BUFFER_TOO_SMALL;
+               if (diff == 1)
+                  data[pos++] = 0x02;
+               else {
+                  data[pos++] = 0x03;
+                  data[pos++] = diff;
+               }
+               curr_frame = frame;
+            }
+            if (extensions[i].id < 32)
+            {
+               if (extensions[i].len < 0 || extensions[i].len > 1)
+                  return OPUS_BAD_ARG;
+               if (len-pos < extensions[i].len+1)
+                  return OPUS_BUFFER_TOO_SMALL;
+               data[pos++] = (extensions[i].id<<1) + extensions[i].len;
+               if (extensions[i].len > 0)
+                  data[pos++] = extensions[i].data[0];
+            } else {
+               int last;
+               opus_int32 length_bytes;
+               last = (written == nb_extensions - 1);
+               length_bytes = 1 + extensions[i].len/255;
+               if (last)
+                  length_bytes = 0;
+               if (len-pos < 1 + length_bytes + extensions[i].len)
+                  return OPUS_BUFFER_TOO_SMALL;
+               data[pos++] = (extensions[i].id<<1) + !last;
+               if (!last)
+               {
+                  opus_int32 j;
+                  for (j=0;j<extensions[i].len/255;j++)
+                     data[pos++] = 255;
+                  data[pos++] = extensions[i].len % 255;
+               }
+               OPUS_COPY(&data[pos], extensions[i].data, extensions[i].len);
+               pos += extensions[i].len;
+            }
+            written++;
+         }
+      }
+   }
+   /* If we need to pad, just prepend 0x01 bytes. Even better would be to fill the
+      end with zeros, but that requires checking that turning the last extesion into
+      an L=1 case still fits. */
+   if (pad && pos < len)
+   {
+      opus_int32 padding = len - pos;
+      OPUS_MOVE(data+padding, data, pos);
+      for (i=0;i<padding;i++)
+         data[i] = 0x01;
+      pos += padding;
+   }
+   return pos;
+}
+
+#if 0
+#include <stdio.h>
+int main()
+{
+   opus_extension_data ext[] = {{2, 0, (const unsigned char *)"a", 1},
+   {32, 10, (const unsigned char *)"DRED", 4},
+   {33, 1, (const unsigned char *)"NOT DRED", 8},
+   {3, 4, (const unsigned char *)NULL, 0}
+   };
+   opus_extension_data ext2[10];
+   int i, len;
+   int nb_ext = 10;
+   unsigned char packet[10000];
+   len = opus_packet_extensions_generate(packet, 32, ext, 4, 1);
+   for (i=0;i<len;i++)
+   {
+      printf("%#04x ", packet[i]);
+      if (i%16 == 15)
+         printf("\n");
+   }
+   printf("\n");
+   printf("count = %d\n", opus_packet_extensions_count(packet, len));
+   opus_packet_extensions_parse(packet, len, ext2, &nb_ext);
+   for (i=0;i<nb_ext;i++)
+   {
+      int j;
+      printf("%d %d {", ext2[i].id, ext2[i].frame);
+      for (j=0;j<ext2[i].len;j++) printf("%#04x ", ext2[i].data[j]);
+      printf("} %d\n", ext2[i].len);
+   }
+}
+#endif
diff --git a/src/opus_encoder.c b/src/opus_encoder.c
index bbf7f003..3db2295e 100644
--- a/src/opus_encoder.c
+++ b/src/opus_encoder.c
@@ -1003,7 +1003,7 @@ static opus_int32 encode_multiframe_packet(OpusEncoder *st,
       }
    }
 
-   ret = opus_repacketizer_out_range_impl(rp, 0, nb_frames, data, repacketize_len, 0, !st->use_vbr);
+   ret = opus_repacketizer_out_range_impl(rp, 0, nb_frames, data, repacketize_len, 0, !st->use_vbr, NULL, 0);
 
    if (ret<0)
    {
diff --git a/src/opus_multistream_encoder.c b/src/opus_multistream_encoder.c
index 213e3eb2..1725ade7 100644
--- a/src/opus_multistream_encoder.c
+++ b/src/opus_multistream_encoder.c
@@ -1003,7 +1003,7 @@ int opus_multistream_encode_native
          return OPUS_INTERNAL_ERROR;
       }
       len = opus_repacketizer_out_range_impl(&rp, 0, opus_repacketizer_get_nb_frames(&rp),
-            data, max_data_bytes-tot_size, s != st->layout.nb_streams-1, !vbr && s == st->layout.nb_streams-1);
+            data, max_data_bytes-tot_size, s != st->layout.nb_streams-1, !vbr && s == st->layout.nb_streams-1, NULL, 0);
       data += len;
       tot_size += len;
    }
diff --git a/src/opus_private.h b/src/opus_private.h
index 5e2463f5..fe1cc755 100644
--- a/src/opus_private.h
+++ b/src/opus_private.h
@@ -44,6 +44,13 @@ struct OpusRepacketizer {
    int framesize;
 };
 
+typedef struct {
+   int id;
+   int frame;
+   const unsigned char *data;
+   opus_int32 len;
+} opus_extension_data;
+
 typedef struct ChannelLayout {
    int nb_channels;
    int nb_streams;
@@ -168,7 +175,8 @@ int opus_packet_parse_impl(const unsigned char *data, opus_int32 len,
       int *payload_offset, opus_int32 *packet_offset);
 
 opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int end,
-      unsigned char *data, opus_int32 maxlen, int self_delimited, int pad);
+      unsigned char *data, opus_int32 maxlen, int self_delimited, int pad,
+      const opus_extension_data *extensions, int nb_extensions);
 
 int pad_frame(unsigned char *data, opus_int32 len, opus_int32 new_len);
 
@@ -198,4 +206,8 @@ int opus_multistream_decode_native(
   void *user_data
 );
 
+opus_int32 opus_packet_extensions_parse(const unsigned char *data, opus_int32 len, opus_extension_data *extensions, opus_int32 *nb_extensions);
+
+opus_int32 opus_packet_extensions_generate(unsigned char *data, opus_int32 len, const opus_extension_data  *extensions, int nb_extensions, int pad);
+
 #endif /* OPUS_PRIVATE_H */
diff --git a/src/repacketizer.c b/src/repacketizer.c
index bda44a14..0e1e157b 100644
--- a/src/repacketizer.c
+++ b/src/repacketizer.c
@@ -100,7 +100,7 @@ int opus_repacketizer_get_nb_frames(OpusRepacketizer *rp)
 }
 
 opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int end,
-      unsigned char *data, opus_int32 maxlen, int self_delimited, int pad)
+      unsigned char *data, opus_int32 maxlen, int self_delimited, int pad, const opus_extension_data *extensions, int nb_extensions)
 {
    int i, count;
    opus_int32 tot_size;
@@ -148,11 +148,12 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
          ptr += encode_size(len[0], ptr);
       }
    }
-   if (count > 2 || (pad && tot_size < maxlen))
+   if (count > 2 || (pad && tot_size < maxlen) || nb_extensions > 0)
    {
       /* Code 3 */
       int vbr;
       int pad_amount=0;
+      int ext_len=0;
 
       /* Restart the process for the padding case */
       ptr = data;
@@ -188,11 +189,24 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
          *ptr++ = count;
       }
       pad_amount = pad ? (maxlen-tot_size) : 0;
+      if (nb_extensions>0)
+      {
+         ext_len = opus_packet_extensions_generate(&data[tot_size], maxlen-tot_size, extensions, nb_extensions, 0);
+         if (ext_len < 0) return ext_len;
+         if (!pad)
+            pad_amount = ext_len + ext_len/254 + 1;
+      }
       if (pad_amount != 0)
       {
          int nb_255s;
          data[1] |= 0x40;
          nb_255s = (pad_amount-1)/255;
+         if (tot_size + ext_len + nb_255s + 1 > maxlen)
+            return OPUS_BUFFER_TOO_SMALL;
+         OPUS_MOVE(&data[tot_size+pad_amount-ext_len], &data[tot_size], ext_len);
+         /* Prepend 0x01 padding */
+         for (i=tot_size+nb_255s+1;i<tot_size+pad_amount-ext_len;i++)
+            data[i] = 0x01;
          for (i=0;i<nb_255s;i++)
             *ptr++ = 255;
          *ptr++ = pad_amount-255*nb_255s-1;
@@ -218,7 +232,7 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
       OPUS_MOVE(ptr, frames[i], len[i]);
       ptr += len[i];
    }
-   if (pad)
+   if (pad && nb_extensions==0)
    {
       /* Fill padding with zeros. */
       while (ptr<data+maxlen)
@@ -229,12 +243,12 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
 
 opus_int32 opus_repacketizer_out_range(OpusRepacketizer *rp, int begin, int end, unsigned char *data, opus_int32 maxlen)
 {
-   return opus_repacketizer_out_range_impl(rp, begin, end, data, maxlen, 0, 0);
+   return opus_repacketizer_out_range_impl(rp, begin, end, data, maxlen, 0, 0, NULL, 0);
 }
 
 opus_int32 opus_repacketizer_out(OpusRepacketizer *rp, unsigned char *data, opus_int32 maxlen)
 {
-   return opus_repacketizer_out_range_impl(rp, 0, rp->nb_frames, data, maxlen, 0, 0);
+   return opus_repacketizer_out_range_impl(rp, 0, rp->nb_frames, data, maxlen, 0, 0, NULL, 0);
 }
 
 int opus_packet_pad(unsigned char *data, opus_int32 len, opus_int32 new_len)
@@ -253,7 +267,7 @@ int opus_packet_pad(unsigned char *data, opus_int32 len, opus_int32 new_len)
    ret = opus_repacketizer_cat(&rp, data+new_len-len, len);
    if (ret != OPUS_OK)
       return ret;
-   ret = opus_repacketizer_out_range_impl(&rp, 0, rp.nb_frames, data, new_len, 0, 1);
+   ret = opus_repacketizer_out_range_impl(&rp, 0, rp.nb_frames, data, new_len, 0, 1, NULL, 0);
    if (ret > 0)
       return OPUS_OK;
    else
@@ -270,7 +284,7 @@ opus_int32 opus_packet_unpad(unsigned char *data, opus_int32 len)
    ret = opus_repacketizer_cat(&rp, data, len);
    if (ret < 0)
       return ret;
-   ret = opus_repacketizer_out_range_impl(&rp, 0, rp.nb_frames, data, len, 0, 0);
+   ret = opus_repacketizer_out_range_impl(&rp, 0, rp.nb_frames, data, len, 0, 0, NULL, 0);
    celt_assert(ret > 0 && ret <= len);
    return ret;
 }
@@ -335,7 +349,7 @@ opus_int32 opus_multistream_packet_unpad(unsigned char *data, opus_int32 len, in
       ret = opus_repacketizer_cat_impl(&rp, data, packet_offset, self_delimited);
       if (ret < 0)
          return ret;
-      ret = opus_repacketizer_out_range_impl(&rp, 0, rp.nb_frames, dst, len, self_delimited, 0);
+      ret = opus_repacketizer_out_range_impl(&rp, 0, rp.nb_frames, dst, len, self_delimited, 0, NULL, 0);
       if (ret < 0)
          return ret;
       else