Buffer size mismatch in tensorflow/lite/kernels/stablehlo_pad.cc

[PaDarochek]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

2.16

Custom code

No

OS platform and distribution

No response

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

Pointers this->edge_pad_low_, this->edge_pad_high_, this->interior_pad_ reference memory locations of size 48 bytes as they point to arrays of int64_t of kMaxDims elements, where kMaxDims == 6:

tensorflow/tensorflow/lite/kernels/stablehlo_pad.cc

Line 36 in bd1c3bf

 static constexpr int kMaxDims = 6; 

tensorflow/tensorflow/lite/kernels/stablehlo_pad.cc

Lines 221 to 224 in bd1c3bf

 private: 
  int64_t edge_pad_low_[kMaxDims]; 
  int64_t edge_pad_high_[kMaxDims]; 
  int64_t interior_pad_[kMaxDims]; 

These pointers are passed as parameters to function 'memcpy' with a size parameter TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT * 8 that is always equal to 64 bytes:

tensorflow/tensorflow/lite/kernels/stablehlo_pad.cc

Lines 98 to 108 in bd1c3bf

 explicit PadData(const TfLiteStablehloPadParams& params) { 
   std::memcpy( 
       edge_pad_low_, params.edge_padding_low, 
       TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT * sizeof(int64_t)); 
   std::memcpy( 
       edge_pad_high_, params.edge_padding_high, 
       TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT * sizeof(int64_t)); 
   std::memcpy( 
       interior_pad_, params.interior_padding, 
       TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT * sizeof(int64_t)); 
 } 

tensorflow/tensorflow/lite/core/c/builtin_op_data.h

Line 38 in b19a54c

 #define TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT 8 

This can lead to a buffer overflow.

It's worth noting that in tensorflow/lite/core/api/flatbuffer_conversions_test.cc
var kMaxDims is explicitly assigned the value of the constant TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT:

tensorflow/tensorflow/lite/core/api/flatbuffer_conversions_test.cc

Lines 737 to 738 in b19a54c

 static constexpr int kMaxDims = 
     TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT; 

Standalone code to reproduce the issue

Bug was found by Svace static analysis tool.

Relevant log output

No response