Skip to content
Permalink
Browse files Browse the repository at this point in the history
Use the safer safe_load function instead of unsafe_load when poss…
…ible

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
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Aug 3, 2021
1 parent e1ec21f commit 23d6383
Show file tree
Hide file tree
Showing 5 changed files with 22 additions and 55 deletions.
4 changes: 4 additions & 0 deletions RELEASE.md
Expand Up @@ -15,6 +15,10 @@
`if x.shape.rank == 1: x = tf.expand_dims(x, axis=-1)`.
Functional models as well as Sequential models built with an explicit
input shape are not affected.
* The methods `Model.to_yaml()` and `keras.models.model_from_yaml` have been
replaced to raise a `RuntimeError` as they can be abused to cause arbitrary
code execution. It is recommended to use JSON serialization instead of YAML,
or, a better alternative, serialize to H5.

* `tf.lite`:
* Rename fields `SignatureDef` table in schema to maximize the parity with
Expand Down
2 changes: 1 addition & 1 deletion tensorflow/python/keras/engine/functional.py
Expand Up @@ -53,7 +53,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
Expand Down
13 changes: 0 additions & 13 deletions tensorflow/python/keras/engine/functional_test.py
Expand Up @@ -47,11 +47,6 @@
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):

Expand Down Expand Up @@ -627,10 +622,6 @@ def test_multi_input_multi_output_recursion(self):
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')
Expand Down Expand Up @@ -1361,10 +1352,6 @@ def test_constant_initializer_with_numpy(self):
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):
Expand Down
18 changes: 8 additions & 10 deletions tensorflow/python/keras/engine/training.py
Expand Up @@ -87,11 +87,6 @@
import h5py
except ImportError:
h5py = None

try:
import yaml
except ImportError:
yaml = None
# pylint: enable=g-import-not-at-top


Expand Down Expand Up @@ -2416,6 +2411,9 @@ def to_json(self, **kwargs):
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={})`.
Expand All @@ -2431,12 +2429,12 @@ def to_yaml(self, **kwargs):
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:
Expand Down
40 changes: 9 additions & 31 deletions tensorflow/python/keras/saving/model_config.py
Expand Up @@ -18,18 +18,11 @@
from tensorflow.python.keras.saving.saved_model import json_utils
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):
"""Instantiates a Keras model from its config.
Usage:
```
# for a Functional API model
Expand Down Expand Up @@ -63,17 +56,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.
Args:
yaml_string: YAML string or open file encoding a model configuration.
Expand All @@ -85,19 +69,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')
Expand Down

0 comments on commit 23d6383

Please sign in to comment.