Add bounds checking to OgaSequencesGetSequence C API functions

src/ort_genai_c.cpp

@@ -4,6 +4,7 @@
 #include <stdexcept>
 #include <cstdint>
 #include <cstddef>
+#include <limits>
 #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<ONNXTensorElementDataType>(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<const int64_t>{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<uint64_t>(shape_dims[i]) > std::numeric_limits<size_t>::max())
+      throw std::runtime_error("shape dimension exceeds size_t range");
+    const size_t dim = static_cast<size_t>(shape_dims[i]);
+    if (dim != 0 && byte_count > std::numeric_limits<size_t>::max() / dim)
+      throw std::runtime_error("tensor byte count overflow");
+    byte_count *= dim;
+  }
   std::unique_ptr<OrtValue> ort_tensor;
   if (data)
     ort_tensor = OrtValue::CreateTensor(*p_memory_info, data, byte_count, shape, ort_element_type);

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);
 

test/c_api_tests.cpp

@@ -288,6 +288,23 @@ TEST(CAPITests, AppendTokensToSequence) {
 #endif
 }
 
+TEST(CAPITests, SequencesOutOfBoundsAccess) {
+  auto sequences = OgaSequences::Create();
+
+  std::vector<int32_t> 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<int32_t> input_ids_0{1, 2, 3, 5, 8};