Code injection in `saved_model_cli` in TensorFlow

@@ -684,7 +684,7 @@ def load_inputs_from_input_arg_string(inputs_str, input_exprs_str,
   tensor_key_feed_dict = {}
 
   inputs = preprocess_inputs_arg_string(inputs_str)
-  input_exprs = preprocess_input_exprs_arg_string(input_exprs_str, safe=False)
+  input_exprs = preprocess_input_exprs_arg_string(input_exprs_str)
   input_examples = preprocess_input_examples_arg_string(input_examples_str)
 
   for input_tensor_key, (filename, variable_name) in inputs.items():

@@ -486,43 +486,6 @@ def testInputParserPickle(self):
     self.assertTrue(np.all(feed_dict['y'] == pkl1))
     self.assertTrue(np.all(feed_dict['z'] == pkl2))
 
-  def testInputParserPythonExpression(self):
-    x1 = np.ones([2, 10])
-    x2 = np.array([[1], [2], [3]])
-    x3 = np.mgrid[0:5, 0:5]
-    x4 = [[3], [4]]
-    input_expr_str = ('x1=np.ones([2,10]);x2=np.array([[1],[2],[3]]);'
-                      'x3=np.mgrid[0:5,0:5];x4=[[3],[4]]')
-    feed_dict = saved_model_cli.load_inputs_from_input_arg_string(
-        '', input_expr_str, '')
-    self.assertTrue(np.all(feed_dict['x1'] == x1))
-    self.assertTrue(np.all(feed_dict['x2'] == x2))
-    self.assertTrue(np.all(feed_dict['x3'] == x3))
-    self.assertTrue(np.all(feed_dict['x4'] == x4))
-
-  def testInputParserBoth(self):
-    x0 = np.array([[1], [2]])
-    input_path = os.path.join(test.get_temp_dir(), 'input.npz')
-    np.savez(input_path, a=x0)
-    x1 = np.ones([2, 10])
-    input_str = 'x0=' + input_path + '[a]'
-    input_expr_str = 'x1=np.ones([2,10])'
-    feed_dict = saved_model_cli.load_inputs_from_input_arg_string(
-        input_str, input_expr_str, '')
-    self.assertTrue(np.all(feed_dict['x0'] == x0))
-    self.assertTrue(np.all(feed_dict['x1'] == x1))
-
-  def testInputParserBothDuplicate(self):
-    x0 = np.array([[1], [2]])
-    input_path = os.path.join(test.get_temp_dir(), 'input.npz')
-    np.savez(input_path, a=x0)
-    x1 = np.ones([2, 10])
-    input_str = 'x0=' + input_path + '[a]'
-    input_expr_str = 'x0=np.ones([2,10])'
-    feed_dict = saved_model_cli.load_inputs_from_input_arg_string(
-        input_str, input_expr_str, '')
-    self.assertTrue(np.all(feed_dict['x0'] == x1))
-
   def testInputParserErrorNoName(self):
     x0 = np.array([[1], [2]])
     x1 = np.array(range(5))
@@ -629,7 +592,7 @@ def testRunCommandInvalidInputKeyError(self, use_tfrt):
     base_path = test.test_src_dir_path(SAVED_MODEL_PATH)
     args = self.parser.parse_args([
         'run', '--dir', base_path, '--tag_set', 'serve', '--signature_def',
-        'regress_x2_to_y3', '--input_exprs', 'x2=np.ones((3,1))'
+        'regress_x2_to_y3', '--input_exprs', 'x2=[1,2,3]'
     ] + (['--use_tfrt'] if use_tfrt else []))
     with self.assertRaises(ValueError):
       saved_model_cli.run(args)
@@ -640,7 +603,7 @@ def testRunCommandInvalidSignature(self, use_tfrt):
     base_path = test.test_src_dir_path(SAVED_MODEL_PATH)
     args = self.parser.parse_args([
         'run', '--dir', base_path, '--tag_set', 'serve', '--signature_def',
-        'INVALID_SIGNATURE', '--input_exprs', 'x2=np.ones((3,1))'
+        'INVALID_SIGNATURE', '--input_exprs', 'x2=[1,2,3]'
     ] + (['--use_tfrt'] if use_tfrt else []))
     with self.assertRaisesRegex(ValueError,
                                 'Could not find signature "INVALID_SIGNATURE"'):

@@ -249,6 +249,8 @@
         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

@@ -20,6 +20,7 @@
 """
 
 import argparse
+import ast
 import os
 import re
 import sys
@@ -521,16 +522,18 @@ def preprocess_inputs_arg_string(inputs_str):
   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.
@@ -545,8 +548,15 @@ def preprocess_input_exprs_arg_string(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
 
 
@@ -659,7 +669,7 @@ def load_inputs_from_input_arg_string(inputs_str, input_exprs_str,
   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():
@@ -923,8 +933,10 @@ def add_run_subparser(subparsers):
   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: '

@@ -382,7 +382,7 @@ def testInputPreProcessFormats(self):
     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'])
@@ -418,6 +418,11 @@ def testInputPreProcessExamplesWithStrAndBytes(self):
           }
     """, 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')
@@ -434,8 +439,8 @@ def testInputPreProcessErrorBadFormat(self):
     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]])