diff --git a/src/ort_genai_c.cpp b/src/ort_genai_c.cpp index 6bf1443a56..69a123a0fc 100644 --- a/src/ort_genai_c.cpp +++ b/src/ort_genai_c.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "span.h" #include "ort_genai_c.h" #include "generators.h" @@ -161,10 +162,16 @@ size_t OGA_API_CALL OgaSequencesCount(const OgaSequences* p) { } size_t OGA_API_CALL OgaSequencesGetSequenceCount(const OgaSequences* p, size_t sequence) { + if (sequence >= p->size()) { + return 0; + } return (*p)[sequence].size(); } const int32_t* OGA_API_CALL OgaSequencesGetSequenceData(const OgaSequences* p, size_t sequence) { + if (sequence >= p->size()) { + return nullptr; + } return (*p)[sequence].data(); } @@ -749,9 +756,19 @@ OgaResult* OGA_API_CALL OgaCreateTensorFromBuffer(void* data, const int64_t* sha auto p_memory_info = OrtMemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); auto ort_element_type = static_cast(element_type); size_t byte_count = Ort::SizeOf(ort_element_type); + if (shape_dims_count > 0 && shape_dims == nullptr) + throw std::runtime_error("shape_dims must not be null when shape_dims_count is non-zero"); auto shape = std::span{shape_dims, shape_dims_count}; - for (size_t i = 0; i < shape_dims_count; i++) - byte_count *= shape_dims[i]; + for (size_t i = 0; i < shape_dims_count; i++) { + if (shape_dims[i] < 0) + throw std::runtime_error("shape dimension must be non-negative"); + if (static_cast(shape_dims[i]) > std::numeric_limits::max()) + throw std::runtime_error("shape dimension exceeds size_t range"); + const size_t dim = static_cast(shape_dims[i]); + if (dim != 0 && byte_count > std::numeric_limits::max() / dim) + throw std::runtime_error("tensor byte count overflow"); + byte_count *= dim; + } std::unique_ptr ort_tensor; if (data) ort_tensor = OrtValue::CreateTensor(*p_memory_info, data, byte_count, shape, ort_element_type); diff --git a/src/ort_genai_c.h b/src/ort_genai_c.h index c553caa17c..2f0cc86aab 100644 --- a/src/ort_genai_c.h +++ b/src/ort_genai_c.h @@ -174,7 +174,8 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaAppendTokenToSequence(int32_t token, OgaSe * \brief Returns the number of tokens in the sequence at the given index. * \param[in] sequences OgaSequences to use. * \param[in] sequence_index index of the sequence to use. - * \return The number of tokens in the sequence at the given index + * \return The number of tokens in the sequence at the given index. Returns 0 if + * sequence_index is out of bounds (i.e. >= OgaSequencesCount(sequences)). */ OGA_EXPORT size_t OGA_API_CALL OgaSequencesGetSequenceCount(const OgaSequences* sequences, size_t sequence_index); @@ -184,6 +185,7 @@ OGA_EXPORT size_t OGA_API_CALL OgaSequencesGetSequenceCount(const OgaSequences* * \param[in] sequences OgaSequences to use. * \param[in] sequence_index index of the sequence to use. * \return The pointer to the sequence data at the given index. The pointer is valid until the OgaSequences is destroyed. + * Returns nullptr if sequence_index is out of bounds (i.e. >= OgaSequencesCount(sequences)). */ OGA_EXPORT const int32_t* OGA_API_CALL OgaSequencesGetSequenceData(const OgaSequences* sequences, size_t sequence_index); diff --git a/test/c_api_tests.cpp b/test/c_api_tests.cpp index f15fc170e3..a388429909 100644 --- a/test/c_api_tests.cpp +++ b/test/c_api_tests.cpp @@ -288,6 +288,23 @@ TEST(CAPITests, AppendTokensToSequence) { #endif } +TEST(CAPITests, SequencesOutOfBoundsAccess) { + auto sequences = OgaSequences::Create(); + + std::vector tokens{100, 200, 300}; + sequences->Append(tokens.data(), tokens.size()); + + ASSERT_EQ(sequences->Count(), 1u); + EXPECT_EQ(sequences->SequenceCount(0), tokens.size()); + EXPECT_NE(sequences->SequenceData(0), nullptr); + + // Out-of-bounds indices must not read past the underlying storage. + EXPECT_EQ(sequences->SequenceCount(1), 0u); + EXPECT_EQ(sequences->SequenceData(1), nullptr); + EXPECT_EQ(sequences->SequenceCount(1000), 0u); + EXPECT_EQ(sequences->SequenceData(1000), nullptr); +} + TEST(CAPITests, MaxLength) { // Batch size 1 case std::vector input_ids_0{1, 2, 3, 5, 8};