Ensure mapping matrix size is always valid.

Signed-off-by: Jean-Marc Valin <jmvalin@jmvalin.ca>

5 files changed, 40 insertions, 7 deletions

diff --git a/src/mapping_matrix.c b/src/mapping_matrix.c
index d3d59d61..11a14ce1 100644
--- a/src/mapping_matrix.c
+++ b/src/mapping_matrix.c
@@ -39,9 +39,21 @@
 
 #define MATRIX_INDEX(nb_rows, row, col) (nb_rows * col + row)
 
-int mapping_matrix_get_size(int rows, int cols)
+opus_int32 mapping_matrix_get_size(int rows, int cols)
 {
-  return align(sizeof(MappingMatrix)) + align(rows * cols * sizeof(opus_int16));
+  opus_int32 size;
+
+  /* Mapping Matrix must only support up to 255 channels in or out.
+   * Additionally, the total cell count must be <= 65004 octets in order
+   * for the matrix to be stored in an OGG header.
+   */
+  if (rows > 255 || cols > 255)
+      return 0;
+  size = rows * (opus_int32)cols * sizeof(opus_int16);
+  if (size > 65004)
+    return 0;
+
+  return align(sizeof(MappingMatrix)) + align(size);
 }
 
 opus_int16 *mapping_matrix_get_data(const MappingMatrix *matrix)
diff --git a/src/mapping_matrix.h b/src/mapping_matrix.h
index 381c8e1a..86451688 100644
--- a/src/mapping_matrix.h
+++ b/src/mapping_matrix.h
@@ -50,7 +50,7 @@ typedef struct MappingMatrix
     /* Matrix cell data goes here using col-wise ordering. */
 } MappingMatrix;
 
-int mapping_matrix_get_size(int rows, int cols);
+opus_int32 mapping_matrix_get_size(int rows, int cols);
 
 opus_int16 *mapping_matrix_get_data(const MappingMatrix *matrix);
 
diff --git a/src/opus_projection_decoder.c b/src/opus_projection_decoder.c
index 6cdadbe9..4c910d57 100644
--- a/src/opus_projection_decoder.c
+++ b/src/opus_projection_decoder.c
@@ -110,6 +110,9 @@ opus_int32 opus_projection_decoder_get_size(int channels, int streams,
 
   matrix_size =
     mapping_matrix_get_size(streams + coupled_streams, channels);
+  if (!matrix_size)
+    return 0;
+
   decoder_size = opus_multistream_decoder_get_size(streams, coupled_streams);
   if (!decoder_size)
     return 0;
@@ -147,7 +150,14 @@ int opus_projection_decoder_init(OpusProjectionDecoder *st, opus_int32 Fs,
   }
 
   /* Assign demixing matrix. */
-  st->demixing_matrix_size_in_bytes = mapping_matrix_get_size(channels, nb_input_streams);
+  st->demixing_matrix_size_in_bytes =
+    mapping_matrix_get_size(channels, nb_input_streams);
+  if (!st->demixing_matrix_size_in_bytes)
+  {
+    RESTORE_STACK;
+    return OPUS_BAD_ARG;
+  }
+
   mapping_matrix_init(get_demixing_matrix(st), channels, nb_input_streams, 0,
     buf, demixing_matrix_size);
 
diff --git a/src/opus_projection_encoder.c b/src/opus_projection_encoder.c
index 3cf516ab..7242faaf 100644
--- a/src/opus_projection_encoder.c
+++ b/src/opus_projection_encoder.c
@@ -150,6 +150,9 @@ opus_int32 opus_projection_ambisonics_encoder_get_size(int channels,
 
   matrix_rows = order_plus_one * order_plus_one + 2;
   matrix_size = mapping_matrix_get_size(matrix_rows, matrix_rows);
+  if (!matrix_size)
+    return 0;
+
   encoder_size =
       opus_multistream_encoder_get_size(nb_streams, nb_coupled_streams);
   if (!encoder_size)
@@ -210,6 +213,8 @@ int opus_projection_ambisonics_encoder_init(OpusProjectionEncoder *st, opus_int3
     }
     st->mixing_matrix_size_in_bytes = mapping_matrix_get_size(
       mixing_matrix->rows, mixing_matrix->cols);
+    if (!st->mixing_matrix_size_in_bytes)
+      return OPUS_BAD_ARG;
 
     /* Assign demixing matrix based on available pre-computed matrices. */
     demixing_matrix = get_demixing_matrix(st);
@@ -233,6 +238,8 @@ int opus_projection_ambisonics_encoder_init(OpusProjectionEncoder *st, opus_int3
     }
     st->demixing_matrix_size_in_bytes = mapping_matrix_get_size(
       demixing_matrix->rows, demixing_matrix->cols);
+    if (!st->demixing_matrix_size_in_bytes)
+      return OPUS_BAD_ARG;
   }
   else
     return OPUS_UNIMPLEMENTED;
diff --git a/tests/test_opus_projection.c b/tests/test_opus_projection.c
index 01dc42d6..3068cd38 100644
--- a/tests/test_opus_projection.c
+++ b/tests/test_opus_projection.c
@@ -96,6 +96,7 @@ void test_simple_matrix(void)
     -29491, 3277, 0, -3277};
 
   int i, ret;
+  opus_int32 simple_matrix_size;
   opus_val16 *input_val16;
   opus_val16 *output_val16;
   opus_int16 *output_int16;
@@ -107,9 +108,12 @@ void test_simple_matrix(void)
   output_val16 = (opus_val16 *)opus_alloc(align(sizeof(opus_val16) * SIMPLE_MATRIX_OUTPUT_SIZE));
 
   /* Initialize matrix */
-  simple_matrix = (MappingMatrix *)opus_alloc(
-    mapping_matrix_get_size(simple_matrix_params.rows,
-                            simple_matrix_params.cols));
+  simple_matrix_size = mapping_matrix_get_size(simple_matrix_params.rows,
+    simple_matrix_params.cols);
+  if (!simple_matrix_size)
+    test_failed();
+
+  simple_matrix = (MappingMatrix *)opus_alloc(simple_matrix_size);
   mapping_matrix_init(simple_matrix, simple_matrix_params.rows,
     simple_matrix_params.cols, simple_matrix_params.gain, simple_matrix_data,
     sizeof(simple_matrix_data));