From 23d6383eb6c14084a8fc3bdf164043b974818012 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Tue, 3 Aug 2021 11:20:20 -0700 Subject: [PATCH] Use the safer `safe_load` function instead of `unsafe_load` when possible There is no need to open ourselves up to arbitrary code execution,especially since this is not in a performance critical loop, so we can take the slowdown due to safety. PiperOrigin-RevId: 388501098 Change-Id: I3434318a5e07a798490533b554f46752397837e5 --- tensorflow/python/keras/engine/functional.py | 2 +- .../python/keras/engine/functional_test.py | 13 ------- tensorflow/python/keras/engine/training.py | 18 ++++----- .../python/keras/saving/model_config.py | 38 ++++--------------- 4 files changed, 17 insertions(+), 54 deletions(-) diff --git a/tensorflow/python/keras/engine/functional.py b/tensorflow/python/keras/engine/functional.py index fd80e7f8..b0cf778a 100644 --- a/tensorflow/python/keras/engine/functional.py +++ b/tensorflow/python/keras/engine/functional.py @@ -58,7 +58,7 @@ class Functional(training_lib.Model): than with subclassed `Model`s, specifically: - Model cloning (`keras.models.clone`) - - Serialization (`model.get_config()/from_config`, `model.to_json()/to_yaml()` + - Serialization (`model.get_config()/from_config`, `model.to_json()` - Whole-model saving (`model.save()`) A `Functional` model can be instantiated by passing two arguments to diff --git a/tensorflow/python/keras/engine/functional_test.py b/tensorflow/python/keras/engine/functional_test.py index b60373e8..c91026a6 100644 --- a/tensorflow/python/keras/engine/functional_test.py +++ b/tensorflow/python/keras/engine/functional_test.py @@ -52,11 +52,6 @@ from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.platform import test from tensorflow.python.training.tracking.util import Checkpoint -try: - import yaml # pylint:disable=g-import-not-at-top -except ImportError: - yaml = None - class NetworkConstructionTest(keras_parameterized.TestCase): @@ -620,10 +615,6 @@ class NetworkConstructionTest(keras_parameterized.TestCase): json_str = model.to_json() models.model_from_json(json_str) - if yaml is not None: - yaml_str = model.to_yaml() - models.model_from_yaml(yaml_str) - @combinations.generate(combinations.combine(mode=['graph', 'eager'])) def test_invalid_graphs(self): a = layers.Input(shape=(32,), name='input_a') @@ -1261,10 +1252,6 @@ class NetworkConstructionTest(keras_parameterized.TestCase): json_str = model.to_json() models.model_from_json(json_str) - if yaml is not None: - yaml_str = model.to_yaml() - models.model_from_yaml(yaml_str) - def test_subclassed_error_if_init_not_called(self): class MyNetwork(training_lib.Model): diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py index a0ebec4f..e000e62f 100644 --- a/tensorflow/python/keras/engine/training.py +++ b/tensorflow/python/keras/engine/training.py @@ -88,11 +88,6 @@ try: import h5py except ImportError: h5py = None - -try: - import yaml -except ImportError: - yaml = None # pylint: enable=g-import-not-at-top @@ -2258,6 +2253,9 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector): def to_yaml(self, **kwargs): """Returns a yaml string containing the network configuration. + Note: Since TF 2.6, this method is no longer supported and will raise a + RuntimeError. + To load a network from a yaml save file, use `keras.models.model_from_yaml(yaml_string, custom_objects={})`. @@ -2273,12 +2271,12 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector): A YAML string. Raises: - ImportError: if yaml module is not found. + RuntimeError: announces that the method poses a security risk """ - if yaml is None: - raise ImportError( - 'Requires yaml module installed (`pip install pyyaml`).') - return yaml.dump(self._updated_config(), **kwargs) + raise RuntimeError( + 'Method `model.to_yaml()` has been removed due to security risk of ' + 'arbitrary code execution. Please use `model.to_json()` instead.' + ) def reset_states(self): for layer in self.layers: diff --git a/tensorflow/python/keras/saving/model_config.py b/tensorflow/python/keras/saving/model_config.py index 63f82b40..344e543f 100644 --- a/tensorflow/python/keras/saving/model_config.py +++ b/tensorflow/python/keras/saving/model_config.py @@ -23,13 +23,6 @@ import json from tensorflow.python.util.tf_export import keras_export -# pylint: disable=g-import-not-at-top -try: - import yaml -except ImportError: - yaml = None -# pylint: enable=g-import-not-at-top - @keras_export('keras.models.model_from_config') def model_from_config(config, custom_objects=None): @@ -59,17 +52,8 @@ def model_from_config(config, custom_objects=None): def model_from_yaml(yaml_string, custom_objects=None): """Parses a yaml model configuration file and returns a model instance. - Usage: - - >>> model = tf.keras.Sequential([ - ... tf.keras.layers.Dense(5, input_shape=(3,)), - ... tf.keras.layers.Softmax()]) - >>> try: - ... import yaml - ... config = model.to_yaml() - ... loaded_model = tf.keras.models.model_from_yaml(config) - ... except ImportError: - ... pass + Note: Since TF 2.6, this method is no longer supported and will raise a + RuntimeError. Arguments: yaml_string: YAML string or open file encoding a model configuration. @@ -81,19 +65,13 @@ def model_from_yaml(yaml_string, custom_objects=None): A Keras model instance (uncompiled). Raises: - ImportError: if yaml module is not found. + RuntimeError: announces that the method poses a security risk """ - if yaml is None: - raise ImportError('Requires yaml module installed (`pip install pyyaml`).') - # The method unsafe_load only exists in PyYAML 5.x+, so which branch of the - # try block is covered by tests depends on the installed version of PyYAML. - try: - # PyYAML 5.x+ - config = yaml.unsafe_load(yaml_string) - except AttributeError: - config = yaml.load(yaml_string) - from tensorflow.python.keras.layers import deserialize # pylint: disable=g-import-not-at-top - return deserialize(config, custom_objects=custom_objects) + raise RuntimeError( + 'Method `model_from_yaml()` has been removed due to security risk of ' + 'arbitrary code execution. Please use `Model.to_json()` and ' + '`model_from_json()` instead.' + ) @keras_export('keras.models.model_from_json') -- 2.27.0