Fix Neural Solution SQL/CMD injection (#1627)

Signed-off-by: Kaihui-intel <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Kaihui-intel

neural_solution/backend/scheduler.py

@@ -154,7 +154,7 @@ def prepare_task(self, task: Task):
         if not task.optimized:
             # Generate quantization code with Neural Coder API
             neural_coder_cmd = ["python -m neural_coder --enable --approach"]
-            # for users to define approach: "static, ""static_ipex", "dynamic", "auto"
+            # for users to define approach: "static", "static_ipex", "dynamic", "auto"
             approach = task.approach
             neural_coder_cmd.append(approach)
             if is_remote_url(task.script_url):

neural_solution/examples/custom_models_optimized/tf_example1/README.md

@@ -91,7 +91,7 @@ optional arguments:
     "script_url": "tf_example1",
     "optimized": "True",
     "arguments": [
-        "--dataset_location=dataset --model_path=model"
+        "--dataset_location=dataset", "--model_path=model"
     ],
     "approach": "static",
     "requirements": [
@@ -106,7 +106,7 @@ When using distributed quantization, the `workers` needs to be set to greater th
     "script_url": "tf_example1",
     "optimized": "True",
     "arguments": [
-        "--dataset_location=dataset --model_path=model"
+        "--dataset_location=dataset", "--model_path=model"
     ],
     "approach": "static",
     "requirements": [

neural_solution/examples/custom_models_optimized/tf_example1/task_request.json

@@ -2,7 +2,7 @@
     "script_url": "custom_models_optimized/tf_example1",
     "optimized": "True",
     "arguments": [
-        "--dataset_location=dataset --model_path=model"
+        "--dataset_location=dataset", "--model_path=model"
     ],
     "approach": "static",
     "requirements": ["tensorflow"

neural_solution/examples/custom_models_optimized/tf_example1/task_request_distributed.json

@@ -2,7 +2,7 @@
     "script_url": "custom_models_optimized/tf_example1",
     "optimized": "True",
     "arguments": [
-        "--dataset_location=dataset --model_path=model"
+        "--dataset_location=dataset", "--model_path=model"
     ],
     "approach": "static",
     "requirements": ["tensorflow"

neural_solution/examples/hf_models/README.md

@@ -72,7 +72,7 @@ optional arguments:
     "script_url": "https://github.com/huggingface/transformers/blob/v4.21-release/examples/pytorch/text-classification/run_glue.py",
     "optimized": "False",
     "arguments": [
-        "--model_name_or_path bert-base-cased --task_name mrpc --do_eval --output_dir result"
+        "--model_name_or_path=bert-base-cased", "--task_name=mrpc", "--do_eval", "--output_dir=result"
     ],
     "approach": "static",
     "requirements": [],

neural_solution/examples/hf_models/task_request.json

@@ -2,7 +2,7 @@
     "script_url": "https://github.com/huggingface/transformers/blob/v4.21-release/examples/pytorch/text-classification/run_glue.py",
     "optimized": "False",
     "arguments": [
-        "--model_name_or_path bert-base-cased --task_name mrpc --do_eval --output_dir result"
+        "--model_name_or_path=bert-base-cased", "--task_name=mrpc", "--do_eval", "--output_dir=result"
     ],
     "approach": "static",
     "requirements": ["datasets", "transformers=4.21.0", "torch"],

neural_solution/examples/hf_models_grpc/README.md

@@ -68,7 +68,7 @@ optional arguments:
     "script_url": "https://github.com/huggingface/transformers/blob/v4.21-release/examples/pytorch/text-classification/run_glue.py",
     "optimized": "False",
     "arguments": [
-        "--model_name_or_path bert-base-cased --task_name mrpc --do_eval --output_dir result"
+        "--model_name_or_path=bert-base-cased", "--task_name=mrpc", "--do_eval", "--output_dir=result"
     ],
     "approach": "static",
     "requirements": [],

neural_solution/examples/hf_models_grpc/task_request.json

@@ -2,7 +2,7 @@
     "script_url": "https://github.com/huggingface/transformers/blob/v4.21-release/examples/pytorch/text-classification/run_glue.py",
     "optimized": "False",
     "arguments": [
-        "--model_name_or_path bert-base-cased --task_name mrpc --do_eval --output_dir result"
+        "--model_name_or_path=bert-base-cased", "--task_name=mrpc", "--do_eval", "--output_dir=result"
     ],
     "approach": "static",
     "requirements": ["datasets", "transformers=4.21.0", "torch"],

neural_solution/frontend/fastapi/main_server.py

@@ -36,6 +36,7 @@
     get_cluster_info,
     get_cluster_table,
     get_res_during_tuning,
+    is_valid_task,
     list_to_string,
     serialize,
 )
@@ -153,10 +154,14 @@ async def submit_task(task: Task):
     Returns:
         json: status , id of task and messages.
     """
+    if not is_valid_task(task.dict()):
+        raise HTTPException(status_code=422, detail="Invalid task")
+
     msg = "Task submitted successfully"
     status = "successfully"
     # search the current
     db_path = get_db_path(config.workspace)
+
     if os.path.isfile(db_path):
         conn = sqlite3.connect(db_path)
         cursor = conn.cursor()

neural_solution/frontend/utility.py

@@ -295,3 +295,61 @@ def list_to_string(lst: list):
         str: string
     """
     return " ".join(str(i) for i in lst)
+
+
+def is_invalid_str(to_test_str: str):
+    """Verify whether the to_test_str is valid.
+
+    Args:
+        to_test_str (str): string to be tested.
+
+    Returns:
+        bool: valid or invalid
+    """
+    return any(char in to_test_str for char in [" ", '"', "'", "&", "|", ";", "`", ">"])
+
+
+def is_valid_task(task: dict) -> bool:
+    """Verify whether the task is valid.
+
+    Args:
+        task (dict): task request
+
+    Returns:
+        bool: valid or invalid
+    """
+    required_fields = ["script_url", "optimized", "arguments", "approach", "requirements", "workers"]
+
+    for field in required_fields:
+        if field not in task:
+            return False
+
+    if not isinstance(task["script_url"], str) or is_invalid_str(task["script_url"]):
+        return False
+
+    if (isinstance(task["optimized"], str) and task["optimized"] not in ["True", "False"]) or (
+        not isinstance(task["optimized"], str) and not isinstance(task["optimized"], bool)
+    ):
+        return False
+
+    if not isinstance(task["arguments"], list):
+        return False
+    else:
+        for argument in task["arguments"]:
+            if is_invalid_str(argument):
+                return False
+
+    if not isinstance(task["approach"], str) or task["approach"] not in ["static", "static_ipex", "dynamic", "auto"]:
+        return False
+
+    if not isinstance(task["requirements"], list):
+        return False
+    else:
+        for requirement in task["requirements"]:
+            if is_invalid_str(requirement):
+                return False
+
+    if not isinstance(task["workers"], int) or task["workers"] < 1:
+        return False
+
+    return True