CVE-2021-37678: Deserialization of Untrusted Data (CWE-502)
value: TensorFlow is an end-to-end open source platform for machine learning. In affected versions TensorFlow and Keras can be tricked to perform arbitrary code execution when deserializing a Keras model from YAML format. The [implementation](https://github.com/tensorflow/tensorflow/blob/460e000de3a83278fb00b61a16d161b1964f15f4/tensorflow/python/keras/saving/model_config.py#L66-L104) uses `yaml.unsafe_load` which can perform arbitrary code execution on the input. Given that YAML format support requires a significant amount of work, we have removed it for now. We have patched the issue in GitHub commit 23d6383eb6c14084a8fc3bdf164043b974818012. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.
CWE description: The product deserializes untrusted data without sufficiently verifying that the resulting data will be valid.
CVSS v3 base score: 8.8
Published: 2021-08-12T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: RELEASE.md, tensorflow/python/keras/engine/functional.py, tensorflow/python/keras/engine/functional_test.py, tensorflow/python/keras/engine/training.py, tensorflow/python/keras/saving/model_config.py
Commit date: 2021-08-03 11:25:56 -0700
Commit message: 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

Diff: diff --git a/RELEASE.md b/RELEASE.md
index 2e85cdca26b1e..27779563f2922 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -13,10 +13,14 @@
     You can revert back to the previous behavior by adding upranking yourself
     in the `train_step()`/`test_step()`/`predict_step()` methods, e.g.
     `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
     TF SavedModel's Signature concept.
 
diff --git a/tensorflow/python/keras/engine/functional.py b/tensorflow/python/keras/engine/functional.py
index fc1cb8ddcced9..ccaee609dc859 100644
--- a/tensorflow/python/keras/engine/functional.py
+++ b/tensorflow/python/keras/engine/functional.py
@@ -51,11 +51,11 @@ class Functional(training_lib.Model):
   and `Sequential` (a special case of `Functional`).
   In general, more Keras features are supported with `Functional`
   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
   `__init__`. The first argument is the `keras.Input` Tensors that represent
   the inputs to the model. The second argument specifies the output
diff --git a/tensorflow/python/keras/engine/functional_test.py b/tensorflow/python/keras/engine/functional_test.py
index 8421070805d42..83f6707516603 100644
--- a/tensorflow/python/keras/engine/functional_test.py
+++ b/tensorflow/python/keras/engine/functional_test.py
@@ -45,15 +45,10 @@ from tensorflow.python.ops import state_ops
 from tensorflow.python.ops import string_ops
 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):
 
   def test_default_model_name(self):
     inputs = input_layer_lib.Input(shape=(1,))
@@ -625,14 +620,10 @@ class NetworkConstructionTest(keras_parameterized.TestCase):
 
       model.summary()
       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')
     b = layers.Input(shape=(32,), name='input_b')
 
@@ -1359,14 +1350,10 @@ class NetworkConstructionTest(keras_parameterized.TestCase):
         run_eagerly=testing_utils.should_run_eagerly())
 
     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):
 
       def __init__(self):
diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index 32cc2a301d223..14278b25fa321 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -85,15 +85,10 @@ from tensorflow.tools.docs import doc_controls
 # pylint: disable=g-import-not-at-top
 try:
   import h5py
 except ImportError:
   h5py = None
-
-try:
-  import yaml
-except ImportError:
-  yaml = None
 # pylint: enable=g-import-not-at-top
 
 
 def disable_multi_worker(method):
   """Decorator that disallows multi-worker use of `method`."""
@@ -2414,10 +2409,13 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
         model_config, default=json_utils.get_json_type, **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={})`.
 
     `custom_objects` should be a dictionary mapping
     the names of custom losses / layers / etc to the corresponding
@@ -2429,16 +2427,16 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
 
     Returns:
         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:
       if hasattr(layer, 'reset_states') and getattr(layer, 'stateful', False):
         layer.reset_states()
diff --git a/tensorflow/python/keras/saving/model_config.py b/tensorflow/python/keras/saving/model_config.py
index efe2a94e4f63c..1f4309e0c461e 100644
--- a/tensorflow/python/keras/saving/model_config.py
+++ b/tensorflow/python/keras/saving/model_config.py
@@ -16,22 +16,15 @@
 """Functions that save the model's config into different formats."""
 
 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
   tf.keras.Model().from_config(model.get_config())
 
@@ -61,21 +54,12 @@ def model_from_config(config, custom_objects=None):
 
 @keras_export('keras.models.model_from_yaml')
 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.
       custom_objects: Optional dictionary mapping names
           (strings) to custom classes or functions to be
@@ -83,23 +67,17 @@ def model_from_yaml(yaml_string, custom_objects=None):
 
   Returns:
       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')
 def model_from_json(json_string, custom_objects=None):
   """Parses a JSON model configuration string and returns a model instance.

Vulnerable code:   - Serialization (`model.get_config()/from_config`, `model.to_json()/to_yaml()`
try:
  import yaml  # pylint:disable=g-import-not-at-top
except ImportError:
  yaml = None

      if yaml is not None:
        yaml_str = model.to_yaml()
        models.model_from_yaml(yaml_str)

    if yaml is not None:
      yaml_str = model.to_yaml()
      models.model_from_yaml(yaml_str)


try:
  import yaml
except ImportError:
  yaml = None
        ImportError: if yaml module is not found.
    if yaml is None:
      raise ImportError(
          'Requires yaml module installed (`pip install pyyaml`).')
    return yaml.dump(self._updated_config(), **kwargs)
# pylint: disable=g-import-not-at-top
try:
  import yaml
except ImportError:
  yaml = None
# pylint: enable=g-import-not-at-top

 
  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
      ImportError: if yaml module is not found.
  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)
Fixed code:   * 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.
  - Serialization (`model.get_config()/from_config`, `model.to_json()`
    Note: Since TF 2.6, this method is no longer supported and will raise a
    RuntimeError.

        RuntimeError: announces that the method poses a security risk
    raise RuntimeError(
        'Method `model.to_yaml()` has been removed due to security risk of '
        'arbitrary code execution. Please use `model.to_json()` instead.'
    )

  Note: Since TF 2.6, this method is no longer supported and will raise a
  RuntimeError.
      RuntimeError: announces that the method poses a security risk
  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.'
  )