From 61c51875814c65b6d1d877cc80c5920bbeae196a Mon Sep 17 00:00:00 2001
From: Jean-Marc Valin <jmvalin@amazon.com>
Date: Mon, 8 May 2023 17:41:56 -0400
Subject: Fixes corruption when using extensions

Now generating the extension in place once all the data is already
in the right place.
---
 src/extensions.c   | 44 ++++++++++++++++++++++++++++----------------
 src/repacketizer.c | 10 +++++++---
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/src/extensions.c b/src/extensions.c
index e899ad2b..1c286fcd 100644
--- a/src/extensions.c
+++ b/src/extensions.c
@@ -199,11 +199,14 @@ opus_int32 opus_packet_extensions_generate(unsigned char *data, opus_int32 len,
                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;
+               if (diff == 1) {
+                  if (data) data[pos] = 0x02;
+                  pos++;
+               } else {
+                  if (data) data[pos] = 0x03;
+                  pos++;
+                  if (data) data[pos] = diff;
+                  pos++;
                }
                curr_frame = frame;
             }
@@ -213,9 +216,12 @@ opus_int32 opus_packet_extensions_generate(unsigned char *data, opus_int32 len,
                   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];
+               if (data) data[pos] = (extensions[i].id<<1) + extensions[i].len;
+               pos++;
+               if (extensions[i].len > 0) {
+                  data[pos] = extensions[i].data[0];
+                  pos++;
+               }
             } else {
                int last;
                opus_int32 length_bytes;
@@ -225,15 +231,19 @@ opus_int32 opus_packet_extensions_generate(unsigned char *data, opus_int32 len,
                   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 (data) data[pos] = (extensions[i].id<<1) + !last;
+               pos++;
                if (!last)
                {
                   opus_int32 j;
-                  for (j=0;j<extensions[i].len/255;j++)
-                     data[pos++] = 255;
-                  data[pos++] = extensions[i].len % 255;
+                  for (j=0;j<extensions[i].len/255;j++) {
+                     if (data) data[pos] = 255;
+                     pos++;
+                  }
+                  if (data) data[pos] = extensions[i].len % 255;
+                  pos++;
                }
-               OPUS_COPY(&data[pos], extensions[i].data, extensions[i].len);
+               if (data) OPUS_COPY(&data[pos], extensions[i].data, extensions[i].len);
                pos += extensions[i].len;
             }
             written++;
@@ -246,9 +256,11 @@ opus_int32 opus_packet_extensions_generate(unsigned char *data, opus_int32 len,
    if (pad && pos < len)
    {
       opus_int32 padding = len - pos;
-      OPUS_MOVE(data+padding, data, pos);
-      for (i=0;i<padding;i++)
-         data[i] = 0x01;
+      if (data) {
+         OPUS_MOVE(data+padding, data, pos);
+         for (i=0;i<padding;i++)
+            data[i] = 0x01;
+      }
       pos += padding;
    }
    return pos;
diff --git a/src/repacketizer.c b/src/repacketizer.c
index 7e468eb8..cea88e3d 100644
--- a/src/repacketizer.c
+++ b/src/repacketizer.c
@@ -108,6 +108,7 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
    const unsigned char **frames;
    unsigned char * ptr;
    int ones_begin=0, ones_end=0;
+   int ext_begin=0, ext_len=0;
 
    if (begin<0 || begin>=end || end>rp->nb_frames)
    {
@@ -154,7 +155,6 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
       /* Code 3 */
       int vbr;
       int pad_amount=0;
-      int ext_len=0;
 
       /* Restart the process for the padding case */
       ptr = data;
@@ -192,7 +192,7 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
       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);
+         ext_len = opus_packet_extensions_generate(NULL, maxlen-tot_size, extensions, nb_extensions, 0);
          if (ext_len < 0) return ext_len;
          if (!pad)
             pad_amount = ext_len + ext_len/254 + 1;
@@ -204,7 +204,7 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
          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);
+         ext_begin = tot_size+pad_amount-ext_len;
          /* Prepend 0x01 padding */
          ones_begin = tot_size+nb_255s+1;
          ones_end = tot_size+pad_amount-ext_len;
@@ -233,6 +233,10 @@ opus_int32 opus_repacketizer_out_range_impl(OpusRepacketizer *rp, int begin, int
       OPUS_MOVE(ptr, frames[i], len[i]);
       ptr += len[i];
    }
+   if (ext_len > 0) {
+      int ret = opus_packet_extensions_generate(&data[ext_begin], ext_len, extensions, nb_extensions, 0);
+      celt_assert(ret == ext_len);
+   }
    for (i=ones_begin;i<ones_end;i++)
       data[i] = 0x01;
    if (pad && nb_extensions==0)
-- 
cgit v1.2.1