CVE-2022-45907: Improper Control of Generation of Code ('Code Injection') (CWE-94)
value: In PyTorch before trunk/89695, torch.jit.annotations.parse_type_line can cause arbitrary code execution because eval is used unsafely.
CWE description: The product constructs all or part of a code segment using externally-influenced input from an upstream component, but it does not neutralize or incorrectly neutralizes special elements that could modify the syntax or behavior of the intended code segment.
CVSS v3 base score: 9.8
Published: 2022-11-26T02:15Z
Repository: https://github.com/pytorch/pytorch
Language: Python
Files changed: test/test_jit.py, torch/csrc/jit/frontend/script_type_parser.cpp, torch/jit/annotations.py
Commit date: 2022-11-17 22:05:30 +0000
Commit message: [JIT][Security] Do not blindly eval input string (#89189)

Introduce `_eval_no_call` method, that evaluates statement only if it
does not contain any calls(done by examining the bytecode), thus preventing command injection exploit

Added simple unit test to check for that
`torch.jit.annotations.get_signature` would not result in calling random
code.

Although, this code path exists for Python-2 compatibility, and perhaps
should be simply removed.

Fixes https://github.com/pytorch/pytorch/issues/88868

Pull Request resolved: https://github.com/pytorch/pytorch/pull/89189
Approved by: https://github.com/suo

Diff: diff --git a/test/test_jit.py b/test/test_jit.py
index 13c27b0efa..6cbc091d50 100644
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -3949,10 +3949,18 @@ def foo(x):
             def invalid4(a):
                 # type: (Int, Int) -> Int
                 return a + 2
             torch.jit.script(invalid4)
 
+    def test_calls_in_type_annotations(self):
+        with self.assertRaisesRegex(RuntimeError, "Type annotation should not contain calls"):
+            def spooky(a):
+                # type: print("Hello") -> Tensor # noqa: F723
+                return a + 2
+            print(torch.__file__)
+            torch.jit.annotations.get_signature(spooky, None, 1, True)
+
     def test_is_optional(self):
         ann = Union[List[int], List[float]]
         torch._jit_internal.is_optional(ann)
 
     def test_interpreter_fuzz(self):
diff --git a/torch/csrc/jit/frontend/script_type_parser.cpp b/torch/csrc/jit/frontend/script_type_parser.cpp
index f5d6f640d4..d05ec95fb9 100644
--- a/torch/csrc/jit/frontend/script_type_parser.cpp
+++ b/torch/csrc/jit/frontend/script_type_parser.cpp
@@ -314,11 +314,11 @@ std::vector<IValue> ScriptTypeParser::evaluateDefaults(
   // To evaluate the default expressions, we create a graph with no inputs,
   // and whose returns are the default values we need.
   // We then run constant prop on this graph and check the results are
   // constant. This approach avoids having to have separate handling of
   // default arguments from standard expressions by piecing together existing
-  // machinery for graph generation, constant propgation, and constant
+  // machinery for graph generation, constant propagation, and constant
   // extraction.
   auto tuple_type = Subscript::create(
       r,
       Var::create(r, Ident::create(r, "Tuple")),
       List<Expr>::create(r, default_types));
diff --git a/torch/jit/annotations.py b/torch/jit/annotations.py
index a4a36ce36a..a6ff2d04d2 100644
--- a/torch/jit/annotations.py
+++ b/torch/jit/annotations.py
@@ -1,6 +1,7 @@
 import ast
+import dis
 import enum
 import inspect
 import re
 import builtins
 import torch
@@ -142,29 +143,38 @@ def check_fn(fn, loc):
             loc, f"Cannot instantiate class '{py_ast.body[0].name}' in a script function")
     if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef):
         raise torch.jit.frontend.FrontendError(loc, "Expected a single top-level function")
 
 
+def _eval_no_call(stmt, glob, loc):
+    """Evaluate statement as long as it does not contain any method/function calls"""
+    bytecode = compile(stmt, "", mode="eval")
+    for insn in dis.get_instructions(bytecode):
+        if "CALL" in insn.opname:
+            raise RuntimeError(f"Type annotation should not contain calls, but '{stmt}' does")
+    return eval(bytecode, glob, loc)  # type: ignore[arg-type] # noqa: P204
+
+
 def parse_type_line(type_line, rcb, loc):
     """Parses a type annotation specified as a comment.
 
     Example inputs:
         # type: (Tensor, torch.Tensor) -> Tuple[Tensor]
         # type: (Tensor, Tuple[Tensor, Tensor]) -> Tensor
     """
     arg_ann_str, ret_ann_str = split_type_line(type_line)
 
     try:
-        arg_ann = eval(arg_ann_str, {}, EvalEnv(rcb))  # type: ignore[arg-type] # noqa: P204
+        arg_ann = _eval_no_call(arg_ann_str, {}, EvalEnv(rcb))
     except (NameError, SyntaxError) as e:
         raise RuntimeError("Failed to parse the argument list of a type annotation") from e
 
     if not isinstance(arg_ann, tuple):
         arg_ann = (arg_ann,)
 
     try:
-        ret_ann = eval(ret_ann_str, {}, EvalEnv(rcb))  # type: ignore[arg-type] # noqa: P204
+        ret_ann = _eval_no_call(ret_ann_str, {}, EvalEnv(rcb))
     except (NameError, SyntaxError) as e:
         raise RuntimeError("Failed to parse the return type of a type annotation") from e
 
     arg_types = [ann_to_type(ann, loc) for ann in arg_ann]
     return arg_types, ann_to_type(ret_ann, loc)

Vulnerable code:   // machinery for graph generation, constant propgation, and constant
        arg_ann = eval(arg_ann_str, {}, EvalEnv(rcb))  # type: ignore[arg-type] # noqa: P204
        ret_ann = eval(ret_ann_str, {}, EvalEnv(rcb))  # type: ignore[arg-type] # noqa: P204
Fixed code:     def test_calls_in_type_annotations(self):
        with self.assertRaisesRegex(RuntimeError, "Type annotation should not contain calls"):
            def spooky(a):
                # type: print("Hello") -> Tensor # noqa: F723
                return a + 2
            print(torch.__file__)
            torch.jit.annotations.get_signature(spooky, None, 1, True)

  // machinery for graph generation, constant propagation, and constant
import dis
def _eval_no_call(stmt, glob, loc):
    """Evaluate statement as long as it does not contain any method/function calls"""
    bytecode = compile(stmt, "", mode="eval")
    for insn in dis.get_instructions(bytecode):
        if "CALL" in insn.opname:
            raise RuntimeError(f"Type annotation should not contain calls, but '{stmt}' does")
    return eval(bytecode, glob, loc)  # type: ignore[arg-type] # noqa: P204


        arg_ann = _eval_no_call(arg_ann_str, {}, EvalEnv(rcb))
        ret_ann = _eval_no_call(ret_ann_str, {}, EvalEnv(rcb))