Skip to content
Permalink
Browse files Browse the repository at this point in the history
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
  • Loading branch information
k-w-w authored and tensorflower-gardener committed Sep 30, 2021
1 parent 349172c commit 8b202f0
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 10 deletions.
2 changes: 2 additions & 0 deletions RELEASE.md
Expand Up @@ -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
Expand Down
26 changes: 19 additions & 7 deletions tensorflow/python/tools/saved_model_cli.py
Expand Up @@ -20,6 +20,7 @@
"""

import argparse
import ast
import os
import re
import sys
Expand Down Expand Up @@ -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.
Expand All @@ -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


Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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: '
Expand Down
11 changes: 8 additions & 3 deletions tensorflow/python/tools/saved_model_cli_test.py
Expand Up @@ -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'])
Expand Down Expand Up @@ -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')
Expand All @@ -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]])
Expand Down

0 comments on commit 8b202f0

Please sign in to comment.