dnn_backend_native: check operand index

it fixed the issue in https://trac.ffmpeg.org/ticket/8716
zhangxiaoben · Jun 17, 2020 · 0b3bd00 · 0b3bd00
1 parent fc93219
commit 0b3bd00
Show file tree

Hide file tree

Showing 14 changed files with 49 additions and 14 deletions.
diff --git a/libavfilter/dnn/dnn_backend_native.c b/libavfilter/dnn/dnn_backend_native.c
@@ -196,7 +196,7 @@ DNNModel *ff_dnn_load_model_native(const char *model_filename)
         }
 
         network->layers[layer].type = layer_type;
-        parsed_size = layer_funcs[layer_type].pf_load(&network->layers[layer], model_file_context, file_size);
+        parsed_size = layer_funcs[layer_type].pf_load(&network->layers[layer], model_file_context, file_size, network->operands_num);
         if (!parsed_size) {
             goto fail;
         }
@@ -209,6 +209,10 @@ DNNModel *ff_dnn_load_model_native(const char *model_filename)
         int32_t operand_index = (int32_t)avio_rl32(model_file_context);
         dnn_size += 4;
 
+        if (operand_index >= network->operands_num) {
+            goto fail;
+        }
+
         oprd = &network->operands[operand_index];
         name_len = (int32_t)avio_rl32(model_file_context);
         dnn_size += 4;

diff --git a/libavfilter/dnn/dnn_backend_native_layer_conv2d.c b/libavfilter/dnn/dnn_backend_native_layer_conv2d.c
@@ -23,7 +23,7 @@
 
 #define CLAMP_TO_EDGE(x, w) ((x) < 0 ? 0 : ((x) >= (w) ? (w - 1) : (x)))
 
-int dnn_load_layer_conv2d(Layer *layer, AVIOContext *model_file_context, int file_size)
+int dnn_load_layer_conv2d(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num)
 {
     ConvolutionalParams *conv_params;
     int kernel_size;
@@ -80,6 +80,11 @@ int dnn_load_layer_conv2d(Layer *layer, AVIOContext *model_file_context, int fil
     layer->input_operand_indexes[0] = (int32_t)avio_rl32(model_file_context);
     layer->output_operand_index = (int32_t)avio_rl32(model_file_context);
     dnn_size += 8;
+
+    if (layer->input_operand_indexes[0] >= operands_num || layer->output_operand_index >= operands_num) {
+        return 0;
+    }
+
     return dnn_size;
 }
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_conv2d.h b/libavfilter/dnn/dnn_backend_native_layer_conv2d.h
@@ -36,7 +36,7 @@ typedef struct ConvolutionalParams{
     float *biases;
 } ConvolutionalParams;
 
-int dnn_load_layer_conv2d(Layer *layer, AVIOContext *model_file_context, int file_size);
+int dnn_load_layer_conv2d(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num);
 int dnn_execute_layer_conv2d(DnnOperand *operands, const int32_t *input_operand_indexes,
                              int32_t output_operand_index, const void *parameters);
 #endif
diff --git a/libavfilter/dnn/dnn_backend_native_layer_depth2space.c b/libavfilter/dnn/dnn_backend_native_layer_depth2space.c
@@ -27,7 +27,7 @@
 #include "libavutil/avassert.h"
 #include "dnn_backend_native_layer_depth2space.h"
 
-int dnn_load_layer_depth2space(Layer *layer, AVIOContext *model_file_context, int file_size)
+int dnn_load_layer_depth2space(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num)
 {
     DepthToSpaceParams *params;
     int dnn_size = 0;
@@ -42,6 +42,10 @@ int dnn_load_layer_depth2space(Layer *layer, AVIOContext *model_file_context, in
     dnn_size += 8;
     layer->params = params;
 
+    if (layer->input_operand_indexes[0] >= operands_num || layer->output_operand_index >= operands_num) {
+        return 0;
+    }
+
     return dnn_size;
 }
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_depth2space.h b/libavfilter/dnn/dnn_backend_native_layer_depth2space.h
@@ -34,7 +34,7 @@ typedef struct DepthToSpaceParams{
     int block_size;
 } DepthToSpaceParams;
 
-int dnn_load_layer_depth2space(Layer *layer, AVIOContext *model_file_context, int file_size);
+int dnn_load_layer_depth2space(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num);
 int dnn_execute_layer_depth2space(DnnOperand *operands, const int32_t *input_operand_indexes,
                                   int32_t output_operand_index, const void *parameters);
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_mathbinary.c b/libavfilter/dnn/dnn_backend_native_layer_mathbinary.c
@@ -27,7 +27,7 @@
 #include "libavutil/avassert.h"
 #include "dnn_backend_native_layer_mathbinary.h"
 
-int dnn_load_layer_math_binary(Layer *layer, AVIOContext *model_file_context, int file_size)
+int dnn_load_layer_math_binary(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num)
 {
     DnnLayerMathBinaryParams *params;
     int dnn_size = 0;
@@ -45,6 +45,9 @@ int dnn_load_layer_math_binary(Layer *layer, AVIOContext *model_file_context, in
         params->v = av_int2float(avio_rl32(model_file_context));
     } else {
         layer->input_operand_indexes[input_index] = (int32_t)avio_rl32(model_file_context);
+        if (layer->input_operand_indexes[input_index] >= operands_num) {
+            return 0;
+        }
         input_index++;
     }
     dnn_size += 4;
@@ -55,6 +58,9 @@ int dnn_load_layer_math_binary(Layer *layer, AVIOContext *model_file_context, in
         params->v = av_int2float(avio_rl32(model_file_context));
     } else {
         layer->input_operand_indexes[input_index] = (int32_t)avio_rl32(model_file_context);
+        if (layer->input_operand_indexes[input_index] >= operands_num) {
+            return 0;
+        }
         input_index++;
     }
     dnn_size += 4;
@@ -63,6 +69,10 @@ int dnn_load_layer_math_binary(Layer *layer, AVIOContext *model_file_context, in
     dnn_size += 4;
     layer->params = params;
 
+    if (layer->output_operand_index >= operands_num) {
+        return 0;
+    }
+
     return dnn_size;
 }
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_mathbinary.h b/libavfilter/dnn/dnn_backend_native_layer_mathbinary.h
@@ -46,7 +46,7 @@ typedef struct DnnLayerMathBinaryParams{
     float v;
 } DnnLayerMathBinaryParams;
 
-int dnn_load_layer_math_binary(Layer *layer, AVIOContext *model_file_context, int file_size);
+int dnn_load_layer_math_binary(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num);
 int dnn_execute_layer_math_binary(DnnOperand *operands, const int32_t *input_operand_indexes,
                                  int32_t output_operand_index, const void *parameters);
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_mathunary.c b/libavfilter/dnn/dnn_backend_native_layer_mathunary.c
@@ -29,7 +29,7 @@
 #include "libavutil/avassert.h"
 #include "dnn_backend_native_layer_mathunary.h"
 
-int dnn_load_layer_math_unary(Layer *layer, AVIOContext *model_file_context, int file_size)
+int dnn_load_layer_math_unary(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num)
 {
     DnnLayerMathUnaryParams *params;
     int dnn_size = 0;
@@ -44,6 +44,10 @@ int dnn_load_layer_math_unary(Layer *layer, AVIOContext *model_file_context, int
     layer->output_operand_index = (int32_t)avio_rl32(model_file_context);
     dnn_size += 8;
 
+    if (layer->input_operand_indexes[0] >= operands_num || layer->output_operand_index >= operands_num) {
+        return 0;
+    }
+
     return dnn_size;
 
 }

diff --git a/libavfilter/dnn/dnn_backend_native_layer_mathunary.h b/libavfilter/dnn/dnn_backend_native_layer_mathunary.h
@@ -41,7 +41,7 @@ typedef struct DnnLayerMathUnaryParams{
     DNNMathUnaryOperation un_op;
 } DnnLayerMathUnaryParams;
 
-int dnn_load_layer_math_unary(Layer *layer, AVIOContext *model_file_context, int file_size);
+int dnn_load_layer_math_unary(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num);
 int dnn_execute_layer_math_unary(DnnOperand *operands, const int32_t *input_operand_indexes,
                                 int32_t output_operand_index, const void *parameters);
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_maximum.c b/libavfilter/dnn/dnn_backend_native_layer_maximum.c
@@ -27,7 +27,7 @@
 #include "libavutil/avassert.h"
 #include "dnn_backend_native_layer_maximum.h"
 
-int dnn_load_layer_maximum(Layer *layer, AVIOContext *model_file_context, int file_size)
+int dnn_load_layer_maximum(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num)
 {
     DnnLayerMaximumParams *params;
     int dnn_size = 0;
@@ -42,6 +42,10 @@ int dnn_load_layer_maximum(Layer *layer, AVIOContext *model_file_context, int fi
     layer->output_operand_index = (int32_t)avio_rl32(model_file_context);
     dnn_size += 8;
 
+    if (layer->input_operand_indexes[0] >= operands_num || layer->output_operand_index >= operands_num) {
+        return 0;
+    }
+
     return dnn_size;
 }
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_maximum.h b/libavfilter/dnn/dnn_backend_native_layer_maximum.h
@@ -37,7 +37,7 @@ typedef struct DnnLayerMaximumParams{
     }val;
 } DnnLayerMaximumParams;
 
-int dnn_load_layer_maximum(Layer *layer, AVIOContext *model_file_context, int file_size);
+int dnn_load_layer_maximum(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num);
 int dnn_execute_layer_maximum(DnnOperand *operands, const int32_t *input_operand_indexes,
                               int32_t output_operand_index, const void *parameters);
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_pad.c b/libavfilter/dnn/dnn_backend_native_layer_pad.c
@@ -22,7 +22,7 @@
 #include "libavutil/avassert.h"
 #include "dnn_backend_native_layer_pad.h"
 
-int dnn_load_layer_pad(Layer *layer, AVIOContext *model_file_context, int file_size)
+int dnn_load_layer_pad(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num)
 {
     LayerPadParams *params;
     int dnn_size = 0;
@@ -42,6 +42,10 @@ int dnn_load_layer_pad(Layer *layer, AVIOContext *model_file_context, int file_s
     dnn_size += 8;
     layer->params = params;
 
+    if (layer->input_operand_indexes[0] >= operands_num || layer->output_operand_index >= operands_num) {
+        return 0;
+    }
+
     return dnn_size;
 }
 

diff --git a/libavfilter/dnn/dnn_backend_native_layer_pad.h b/libavfilter/dnn/dnn_backend_native_layer_pad.h
@@ -36,7 +36,7 @@ typedef struct LayerPadParams{
     float constant_values;
 } LayerPadParams;
 
-int dnn_load_layer_pad(Layer *layer, AVIOContext *model_file_context, int file_size);
+int dnn_load_layer_pad(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num);
 int dnn_execute_layer_pad(DnnOperand *operands, const int32_t *input_operand_indexes,
                           int32_t output_operand_index, const void *parameters);
 

diff --git a/libavfilter/dnn/dnn_backend_native_layers.h b/libavfilter/dnn/dnn_backend_native_layers.h
@@ -26,7 +26,7 @@
 
 typedef int (*LAYER_EXEC_FUNC)(DnnOperand *operands, const int32_t *input_operand_indexes,
                                int32_t output_operand_index, const void *parameters);
-typedef int (*LAYER_LOAD_FUNC)(Layer *layer, AVIOContext *model_file_context, int file_size);
+typedef int (*LAYER_LOAD_FUNC)(Layer *layer, AVIOContext *model_file_context, int file_size, int operands_num);
 
 typedef struct LayerFunc {
     LAYER_EXEC_FUNC pf_exec;