From 113357235f6ab3caa127450fedf56ee906b8ec94 Mon Sep 17 00:00:00 2001
From: wdxwj <wdxwj@vip.qq.com>
Date: Tue, 2 Mar 2021 09:16:27 +0800
Subject: [PATCH] To verify the input value length of open interface in case of
 overflow attacking

---
 mindinsight/mindconverter/cli.py                  | 15 ++++++++++++---
 .../graph_based_converter/constant.py             |  6 +++++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/mindinsight/mindconverter/cli.py b/mindinsight/mindconverter/cli.py
index c3461c3f..4a6f0518 100644
--- a/mindinsight/mindconverter/cli.py
+++ b/mindinsight/mindconverter/cli.py
@@ -21,7 +21,7 @@ import mindinsight
 from mindinsight.mindconverter.converter import main
 from mindinsight.mindconverter.graph_based_converter.common.utils import get_framework_type
 from mindinsight.mindconverter.graph_based_converter.constant import ARGUMENT_LENGTH_LIMIT, \
-    FrameworkType
+ARGUMENT_NUM_LIMIT, ARGUMENT_LEN_LIMIT, FrameworkType
 from mindinsight.mindconverter.graph_based_converter.framework import main_graph_base_converter
 
 from mindinsight.mindconverter.common.log import logger as log, logger_console as log_console
@@ -250,7 +250,15 @@ class ShapeAction(argparse.Action):
             return [int(num_shape) for num_shape in shape_list.split(',')]
 
         try:
-            in_shape = [_convert_to_int(shape) for shape in values]
+            if len(values) > ARGUMENT_NUM_LIMIT:
+                parser_in.error(f"The length of {option_string} {values} should be no more than {ARGUMENT_NUM_LIMIT}.")
+            in_shape = []
+            for v in values:
+                shape = _convert_to_int(v)
+                if len(shape) > ARGUMENT_LEN_LIMIT:
+                    parser_in.error(
+                        f"The length of {option_string} {shape} should be no more than {ARGUMENT_LEN_LIMIT}.")
+                in_shape.append(shape)
             setattr(namespace, self.dest, in_shape)
         except ValueError:
             parser_in.error(
@@ -273,7 +281,8 @@ class NodeAction(argparse.Action):
         """
 
         ArgsCheck.check_repeated(namespace, self.dest, self.default, option_string, parser_in)
-
+        if len(values) > ARGUMENT_NUM_LIMIT:
+            parser_in.error(f"The length of {option_string} {values} should be no more than {ARGUMENT_NUM_LIMIT}.")
         for v in values:
             if len(v) > ARGUMENT_LENGTH_LIMIT:
                 parser_in.error(
diff --git a/mindinsight/mindconverter/graph_based_converter/constant.py b/mindinsight/mindconverter/graph_based_converter/constant.py
index 0e89745d..7b14f776 100644
--- a/mindinsight/mindconverter/graph_based_converter/constant.py
+++ b/mindinsight/mindconverter/graph_based_converter/constant.py
@@ -93,7 +93,11 @@ BINARY_HEADER_PYTORCH_FILE = \
 TENSORFLOW_MODEL_SUFFIX = "pb"
 BINARY_HEADER_PYTORCH_BITS = 32
 
-ARGUMENT_LENGTH_LIMIT = 512
+ARGUMENT_LENGTH_LIMIT = 128
+
+ARGUMENT_NUM_LIMIT = 32
+
+ARGUMENT_LEN_LIMIT = 64
 
 EXPECTED_NUMBER = 1