diff options
| author |  Alexandre Passos <apassos@google.com> | 2018-03-23 14:19:37 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-03-25 03:01:31 -0700 |
| commit | 13ef0af4867477cdda7e0b294e61560c2952df42 (patch) | |
| tree | 83be86239c24bc2aaec64bde42e3c5323f00609a /tensorflow/python/lib | |
| parent | 2219b88a3d5154b9158a1902b061cad6cae2d0a8 (diff) | |
Fix buffer overflow when fetching resources.
PiperOrigin-RevId: 190273682
Diffstat (limited to 'tensorflow/python/lib')
| -rw-r--r-- | tensorflow/python/lib/core/ndarray_tensor.cc | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/tensorflow/python/lib/core/ndarray_tensor.cc b/tensorflow/python/lib/core/ndarray_tensor.cc index 994af69386..a07e305ffb 100644 --- a/tensorflow/python/lib/core/ndarray_tensor.cc +++ b/tensorflow/python/lib/core/ndarray_tensor.cc @@ -267,7 +267,9 @@ gtl::InlinedVector<npy_intp, 4> GetPyArrayDimensionsForTensor( const int ndims = TF_NumDims(tensor); gtl::InlinedVector<npy_intp, 4> dims(ndims); if (TF_TensorType(tensor) == TF_RESOURCE) { - dims[0] = TF_TensorByteSize(tensor); + CHECK_EQ(ndims, 0) + << "Fetching of non-scalar resource tensors is not supported."; + dims.push_back(TF_TensorByteSize(tensor)); *nelems = dims[0]; } else { *nelems = 1; |