CVE-2021-41228: Improper Control of Generation of Code ('Code Injection') (CWE-94)
value: TensorFlow is an open source platform for machine learning. In affected versions TensorFlow's `saved_model_cli` tool is vulnerable to a code injection as it calls `eval` on user supplied strings. This can be used by attackers to run arbitrary code on the plaform where the CLI tool runs. However, given that the tool is always run manually, the impact of this is not severe. We have patched this by adding a `safe` flag which defaults to `True` and an explicit warning for users. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
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: 7.8
Published: 2021-11-05T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: RELEASE.md, tensorflow/python/tools/saved_model_cli.py, tensorflow/python/tools/saved_model_cli_test.py
Commit date: 2021-09-30 13:40:31 -0700
Commit message: Remove use of `eval` when evaluating the input example.

Use `ast.eval_literal` instead which safely evaluates the expression.

PiperOrigin-RevId: 400012249
Change-Id: I5ff98608ea2d736d093aa488af723ff4f6707e02

Diff: diff --git a/RELEASE.md b/RELEASE.md
index 23820977eca76..aba46a1fad9fc 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -247,10 +247,12 @@
         `tf.data.experimental.choose_from_datasets` API to
         `tf.data.Dataset.choose_from_datasets` and deprecating the experimental
         endpoint.
 *   TF SavedModel:
     *   Custom gradients are now saved by default. See `tf.saved_model.SaveOptions` to disable this.
+    *   The saved_model_cli's `--input_examples` inputs are now restricted to
+        python literals to avoid code injection.
 *   XLA:
     * Added a new API that allows custom call functions to signal errors. The
       old API will be deprecated in a future release. See
       https://www.tensorflow.org/xla/custom_call for details.
     * XLA:GPU reductions are deterministic by default (reductions within
diff --git a/tensorflow/python/tools/saved_model_cli.py b/tensorflow/python/tools/saved_model_cli.py
index 9c37cc40c50bf..f8d7bbb1eece3 100644
--- a/tensorflow/python/tools/saved_model_cli.py
+++ b/tensorflow/python/tools/saved_model_cli.py
@@ -18,10 +18,11 @@ For detailed usages and examples, please refer to:
 https://www.tensorflow.org/guide/saved_model#cli_to_inspect_and_execute_savedmodel
 
 """
 
 import argparse
+import ast
 import os
 import re
 import sys
 
 from absl import app  # pylint: disable=unused-import
@@ -519,20 +520,22 @@ def preprocess_inputs_arg_string(inputs_str):
             '"<input_key>=<filename>[<variable_name>]"' % input_raw)
 
   return input_dict
 
 
-def preprocess_input_exprs_arg_string(input_exprs_str):
+def preprocess_input_exprs_arg_string(input_exprs_str, safe=True):
   """Parses input arg into dictionary that maps input key to python expression.
 
   Parses input string in the format of 'input_key=<python expression>' into a
   dictionary that maps each input_key to its python expression.
 
   Args:
     input_exprs_str: A string that specifies python expression for input keys.
-    Each input is separated by semicolon. For each input key:
+      Each input is separated by semicolon. For each input key:
         'input_key=<python expression>'
+    safe: Whether to evaluate the python expression as literals or allow
+      arbitrary calls (e.g. numpy usage).
 
   Returns:
     A dictionary that maps input keys to their values.
 
   Raises:
@@ -543,12 +546,19 @@ def preprocess_input_exprs_arg_string(input_exprs_str):
   for input_raw in filter(bool, input_exprs_str.split(';')):
     if '=' not in input_exprs_str:
       raise RuntimeError('--input_exprs "%s" format is incorrect. Please follow'
                          '"<input_key>=<python expression>"' % input_exprs_str)
     input_key, expr = input_raw.split('=', 1)
-    # ast.literal_eval does not work with numpy expressions
-    input_dict[input_key] = eval(expr)  # pylint: disable=eval-used
+    if safe:
+      try:
+        input_dict[input_key] = ast.literal_eval(expr)
+      except:
+        raise RuntimeError(
+            f'Expression "{expr}" is not a valid python literal.')
+    else:
+      # ast.literal_eval does not work with numpy expressions
+      input_dict[input_key] = eval(expr)  # pylint: disable=eval-used
   return input_dict
 
 
 def preprocess_input_examples_arg_string(input_examples_str):
   """Parses input into dict that maps input keys to lists of tf.Example.
@@ -657,11 +667,11 @@ def load_inputs_from_input_arg_string(inputs_str, input_exprs_str,
         more than one numpy ndarrays.
   """
   tensor_key_feed_dict = {}
 
   inputs = preprocess_inputs_arg_string(inputs_str)
-  input_exprs = preprocess_input_exprs_arg_string(input_exprs_str)
+  input_exprs = preprocess_input_exprs_arg_string(input_exprs_str, safe=False)
   input_examples = preprocess_input_examples_arg_string(input_examples_str)
 
   for input_tensor_key, (filename, variable_name) in inputs.items():
     data = np.load(file_io.FileIO(filename, mode='rb'), allow_pickle=True)  # pylint: disable=unexpected-keyword-arg
 
@@ -921,12 +931,14 @@ def add_run_subparser(subparsers):
          ' or \'<input_key>=<filename>[<variable_name>]\', separated by \';\'.'
          ' The file format can only be from .npy, .npz or pickle.')
   parser_run.add_argument('--inputs', type=str, default='', help=msg)
   msg = ('Specifying inputs by python expressions, in the format of'
          ' "<input_key>=\'<python expression>\'", separated by \';\'. '
-         'numpy module is available as \'np\'. '
-         'Will override duplicate input keys from --inputs option.')
+         'numpy module is available as \'np\'. Please note that the expression '
+         'will be evaluated as-is, and is susceptible to code injection. '
+         'When this is set, the value will override duplicate input keys from '
+         '--inputs option.')
   parser_run.add_argument('--input_exprs', type=str, default='', help=msg)
   msg = (
       'Specifying tf.Example inputs as list of dictionaries. For example: '
       '<input_key>=[{feature0:value_list,feature1:value_list}]. Use ";" to '
       'separate input keys. Will override duplicate input keys from --inputs '
diff --git a/tensorflow/python/tools/saved_model_cli_test.py b/tensorflow/python/tools/saved_model_cli_test.py
index 9a4cfcd4a0276..424f156b86b88 100644
--- a/tensorflow/python/tools/saved_model_cli_test.py
+++ b/tensorflow/python/tools/saved_model_cli_test.py
@@ -380,11 +380,11 @@ Defined Functions:
   def testInputPreProcessFormats(self):
     input_str = 'input1=/path/file.txt[ab3];input2=file2'
     input_expr_str = 'input3=np.zeros([2,2]);input4=[4,5]'
     input_dict = saved_model_cli.preprocess_inputs_arg_string(input_str)
     input_expr_dict = saved_model_cli.preprocess_input_exprs_arg_string(
-        input_expr_str)
+        input_expr_str, safe=False)
     self.assertTrue(input_dict['input1'] == ('/path/file.txt', 'ab3'))
     self.assertTrue(input_dict['input2'] == ('file2', None))
     print(input_expr_dict['input3'])
     self.assertAllClose(input_expr_dict['input3'], np.zeros([2, 2]))
     self.assertAllClose(input_expr_dict['input4'], [4, 5])
@@ -416,10 +416,15 @@ Defined Functions:
               }
             }
           }
     """, feature)
 
+  def testInputPreprocessExampleWithCodeInjection(self):
+    input_examples_str = 'inputs=os.system("echo hacked")'
+    with self.assertRaisesRegex(RuntimeError, 'not a valid python literal.'):
+      saved_model_cli.preprocess_input_examples_arg_string(input_examples_str)
+
   def testInputPreProcessFileNames(self):
     input_str = (r'inputx=C:\Program Files\data.npz[v:0];'
                  r'input:0=c:\PROGRA~1\data.npy')
     input_dict = saved_model_cli.preprocess_inputs_arg_string(input_str)
     self.assertTrue(input_dict['inputx'] == (r'C:\Program Files\data.npz',
@@ -432,12 +437,12 @@ Defined Functions:
       saved_model_cli.preprocess_inputs_arg_string(input_str)
     input_str = 'inputx:file'
     with self.assertRaises(RuntimeError):
       saved_model_cli.preprocess_inputs_arg_string(input_str)
     input_str = 'inputx:np.zeros((5))'
-    with self.assertRaises(RuntimeError):
-      saved_model_cli.preprocess_input_exprs_arg_string(input_str)
+    with self.assertRaisesRegex(RuntimeError, 'format is incorrect'):
+      saved_model_cli.preprocess_input_exprs_arg_string(input_str, safe=False)
 
   def testInputParserNPY(self):
     x0 = np.array([[1], [2]])
     x1 = np.array(range(6)).reshape(2, 3)
     input0_path = os.path.join(test.get_temp_dir(), 'input0.npy')

Vulnerable code: def preprocess_input_exprs_arg_string(input_exprs_str):
    Each input is separated by semicolon. For each input key:
    # ast.literal_eval does not work with numpy expressions
    input_dict[input_key] = eval(expr)  # pylint: disable=eval-used
  input_exprs = preprocess_input_exprs_arg_string(input_exprs_str)
         'numpy module is available as \'np\'. '
         'Will override duplicate input keys from --inputs option.')
        input_expr_str)
    with self.assertRaises(RuntimeError):
      saved_model_cli.preprocess_input_exprs_arg_string(input_str)
Fixed code:     *   The saved_model_cli's `--input_examples` inputs are now restricted to
        python literals to avoid code injection.
import ast
def preprocess_input_exprs_arg_string(input_exprs_str, safe=True):
      Each input is separated by semicolon. For each input key:
    safe: Whether to evaluate the python expression as literals or allow
      arbitrary calls (e.g. numpy usage).
    if safe:
      try:
        input_dict[input_key] = ast.literal_eval(expr)
      except:
        raise RuntimeError(
            f'Expression "{expr}" is not a valid python literal.')
    else:
      # ast.literal_eval does not work with numpy expressions
      input_dict[input_key] = eval(expr)  # pylint: disable=eval-used
  input_exprs = preprocess_input_exprs_arg_string(input_exprs_str, safe=False)
         'numpy module is available as \'np\'. Please note that the expression '
         'will be evaluated as-is, and is susceptible to code injection. '
         'When this is set, the value will override duplicate input keys from '
         '--inputs option.')
        input_expr_str, safe=False)
  def testInputPreprocessExampleWithCodeInjection(self):
    input_examples_str = 'inputs=os.system("echo hacked")'
    with self.assertRaisesRegex(RuntimeError, 'not a valid python literal.'):
      saved_model_cli.preprocess_input_examples_arg_string(input_examples_str)

    with self.assertRaisesRegex(RuntimeError, 'format is incorrect'):
      saved_model_cli.preprocess_input_exprs_arg_string(input_str, safe=False)