diff --git a/extension/data_loader/buffer_data_loader.h b/extension/data_loader/buffer_data_loader.h index 64e26d3dcb8..7f4211546e0 100644 --- a/extension/data_loader/buffer_data_loader.h +++ b/extension/data_loader/buffer_data_loader.h @@ -40,7 +40,7 @@ class BufferDataLoader final : public executorch::runtime::DataLoader { ET_CHECK_OR_RETURN_ERROR( !overflow && total_size <= size_, InvalidArgument, - "offset %zu + size %zu > size_ %zu", + "offset %zu + size %zu > size_ %zu, or overflow detected", offset, size, size_); diff --git a/extension/data_loader/file_data_loader.cpp b/extension/data_loader/file_data_loader.cpp index 997585aa9d1..bc5c17ef33f 100644 --- a/extension/data_loader/file_data_loader.cpp +++ b/extension/data_loader/file_data_loader.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -143,10 +144,12 @@ Result FileDataLoader::load( fd_ >= 0, InvalidState, "Uninitialized"); + size_t total_size; + bool overflow = c10::add_overflows(offset, size, &total_size); ET_CHECK_OR_RETURN_ERROR( - offset + size <= file_size_, + !overflow && total_size <= file_size_, InvalidArgument, - "File %s: offset %zu + size %zu > file_size_ %zu", + "File %s: offset %zu + size %zu > file_size_ %zu, or overflow detected", file_name_, offset, size, @@ -204,10 +207,12 @@ ET_NODISCARD Error FileDataLoader::load_into( fd_ >= 0, InvalidState, "Uninitialized"); + size_t total_size; + bool overflow = c10::add_overflows(offset, size, &total_size); ET_CHECK_OR_RETURN_ERROR( - offset + size <= file_size_, + !overflow && total_size <= file_size_, InvalidArgument, - "File %s: offset %zu + size %zu > file_size_ %zu", + "File %s: offset %zu + size %zu > file_size_ %zu, or overflow detected", file_name_, offset, size, diff --git a/extension/data_loader/file_descriptor_data_loader.cpp b/extension/data_loader/file_descriptor_data_loader.cpp index a57f2ce3640..3b6423d21cb 100644 --- a/extension/data_loader/file_descriptor_data_loader.cpp +++ b/extension/data_loader/file_descriptor_data_loader.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -157,10 +158,12 @@ Result FileDescriptorDataLoader::load( fd_ >= 0, InvalidState, "Uninitialized"); + size_t total_size; + bool overflow = c10::add_overflows(offset, size, &total_size); ET_CHECK_OR_RETURN_ERROR( - offset + size <= file_size_, + !overflow && total_size <= file_size_, InvalidArgument, - "File %s: offset %zu + size %zu > file_size_ %zu", + "File %s: offset %zu + size %zu > file_size_ %zu, or overflow detected", file_descriptor_uri_, offset, size, @@ -218,10 +221,12 @@ ET_NODISCARD Error FileDescriptorDataLoader::load_into( fd_ >= 0, InvalidState, "Uninitialized"); + size_t total_size; + bool overflow = c10::add_overflows(offset, size, &total_size); ET_CHECK_OR_RETURN_ERROR( - offset + size <= file_size_, + !overflow && total_size <= file_size_, InvalidArgument, - "File %s: offset %zu + size %zu > file_size_ %zu", + "File %s: offset %zu + size %zu > file_size_ %zu, or overflow detected", file_descriptor_uri_, offset, size, diff --git a/extension/data_loader/mmap_data_loader.cpp b/extension/data_loader/mmap_data_loader.cpp index 2271d5a3690..5d77b67cc59 100644 --- a/extension/data_loader/mmap_data_loader.cpp +++ b/extension/data_loader/mmap_data_loader.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -159,10 +160,12 @@ Error MmapDataLoader::validate_input(size_t offset, size_t size) const { fd_ >= 0, InvalidState, "Uninitialized"); + size_t total_size; + bool overflow = c10::add_overflows(offset, size, &total_size); ET_CHECK_OR_RETURN_ERROR( - offset + size <= file_size_, + !overflow && total_size <= file_size_, InvalidArgument, - "File %s: offset %zu + size %zu > file_size_ %zu", + "File %s: offset %zu + size %zu > file_size_ %zu, or overflow detected", file_name_, offset, size, diff --git a/extension/data_loader/shared_ptr_data_loader.h b/extension/data_loader/shared_ptr_data_loader.h index 551ab4d498c..78fdc2b7157 100644 --- a/extension/data_loader/shared_ptr_data_loader.h +++ b/extension/data_loader/shared_ptr_data_loader.h @@ -8,6 +8,7 @@ #pragma once +#include #include #include #include @@ -33,10 +34,12 @@ class SharedPtrDataLoader final : public executorch::runtime::DataLoader { size_t offset, size_t size, ET_UNUSED const DataLoader::SegmentInfo& segment_info) const override { + size_t total_size; + bool overflow = c10::add_overflows(offset, size, &total_size); ET_CHECK_OR_RETURN_ERROR( - offset + size <= size_, + !overflow && total_size <= size_, InvalidArgument, - "offset %zu + size %zu > size_ %zu", + "offset %zu + size %zu > size_ %zu, or overflow detected", offset, size, size_);