Check the result of serializing a Tensor and return an error if it could not be serialized (most likely because it could not fit into a protobuf). This prevents segfaults when trying to load invalid protobufs.

Change: 111576943

1 files changed, 3 insertions, 1 deletions

diff --git a/tensorflow/core/util/tensor_slice_writer.h b/tensorflow/core/util/tensor_slice_writer.h
index 58eeafc17d..3a7d9979f0 100644
--- a/tensorflow/core/util/tensor_slice_writer.h
+++ b/tensorflow/core/util/tensor_slice_writer.h
@@ -136,7 +136,9 @@ Status TensorSliceWriter::Add(const string& name, const TensorShape& shape,
     // list the tensor slices we want to save and then another pass to actually
     // set the data. Need to figure out if the interface works well.
     std::pair<string, string> key_value(key, "");
-    sts.AppendToString(&key_value.second);
+    if (!sts.AppendToString(&key_value.second)) {
+      return errors::Internal("Error writing Tensor. Possible size overflow.");
+    }
     data_.insert(key_value);
   }
   ++slices_;