From 2006ef57602012939dc10d7e8961925b320d3ef6 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 4 Apr 2018 16:26:25 -0700 Subject: [PATCH] Check that n + kBlockTrailerSize does not overflow before reading a block PiperOrigin-RevId: 191666300 --- tensorflow/core/lib/io/format.cc | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorflow/core/lib/io/format.cc b/tensorflow/core/lib/io/format.cc index 64852943ad560e..0c24c660a246ea 100644 --- a/tensorflow/core/lib/io/format.cc +++ b/tensorflow/core/lib/io/format.cc @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include + #include "tensorflow/core/lib/io/format.h" #include "tensorflow/core/lib/core/coding.h" @@ -84,6 +86,11 @@ Status ReadBlock(RandomAccessFile* file, const BlockHandle& handle, // Read the block contents as well as the type/crc footer. // See table_builder.cc for the code that built this structure. size_t n = static_cast(handle.size()); + + if (kBlockTrailerSize > std::numeric_limits::max() - n) { + return errors::DataLoss("handle.size() too big"); + } + char* buf = new char[n + kBlockTrailerSize]; StringPiece contents; Status s = file->Read(handle.offset(), n + kBlockTrailerSize, &contents, buf);