diff --git a/src/shape_tensor.cc b/src/shape_tensor.cc index 49930eb..8cca781 100644 --- a/src/shape_tensor.cc +++ b/src/shape_tensor.cc @@ -1,4 +1,4 @@ -// Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -72,6 +72,11 @@ ShapeTensor::SetDataFromBuffer( *reinterpret_cast(data_.get()) = static_cast(total_batch_size); } + if (size_ < datatype_size) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Unexpected integer underflow while calculating shape tensor size."); + } std::memcpy( data_.get() + datatype_size, data_buffer, size_ - datatype_size); } else {