(mahout) branch main updated: fix(qdp): guard batch size multiplication against overflow (#1324)

[hcr]

This is an automated email from the ASF dual-hosted git repository.

ryankert01 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/mahout.git


The following commit(s) were added to refs/heads/main by this push:
     new ac30a8c95 fix(qdp): guard batch size multiplication against overflow 
(#1324)
ac30a8c95 is described below

commit ac30a8c954f692f0e57e53cef185497b68b61217
Author: KUAN-HAO HUANG <101171023+rich7...@users.noreply.github.com>
AuthorDate: Mon May 18 19:49:33 2026 +0800

    fix(qdp): guard batch size multiplication against overflow (#1324)
    
    `Preprocessor::validate_batch` and `AmplitudeEncoder::encode_batch_f32`
    computed `num_samples * sample_size` without overflow checks, while the
    analogous angle encoder paths already use `checked_mul`. On 64-bit hosts
    the product can wrap and let an undersized buffer pass validation,
    leading to out-of-bounds reads in the encoder.
    
    Mirror the angle.rs pattern: use `checked_mul` and surface an
    `InvalidInput` error on overflow. Add a unit test covering the
    overflow path through the shared preprocessor.
---
 qdp/qdp-core/src/gpu/encodings/amplitude.rs | 10 ++++++++--
 qdp/qdp-core/src/preprocessing.rs           | 30 ++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/qdp/qdp-core/src/gpu/encodings/amplitude.rs 
b/qdp/qdp-core/src/gpu/encodings/amplitude.rs
index e3081b65d..67d9f06ac 100644
--- a/qdp/qdp-core/src/gpu/encodings/amplitude.rs
+++ b/qdp/qdp-core/src/gpu/encodings/amplitude.rs
@@ -482,12 +482,18 @@ impl QuantumEncoder for AmplitudeEncoder {
                 sample_size, state_len, num_qubits
             )));
         }
-        if batch_data.len() != num_samples * sample_size {
+        let expected_len = num_samples.checked_mul(sample_size).ok_or_else(|| {
+            MahoutError::InvalidInput(format!(
+                "Batch size overflow: num_samples {} * sample_size {}",
+                num_samples, sample_size
+            ))
+        })?;
+        if batch_data.len() != expected_len {
             return Err(MahoutError::InvalidInput(format!(
                 "batch_data length mismatch (expected {} * {} = {}, got {})",
                 num_samples,
                 sample_size,
-                num_samples * sample_size,
+                expected_len,
                 batch_data.len()
             )));
         }
diff --git a/qdp/qdp-core/src/preprocessing.rs 
b/qdp/qdp-core/src/preprocessing.rs
index c8469aa39..625c9fa59 100644
--- a/qdp/qdp-core/src/preprocessing.rs
+++ b/qdp/qdp-core/src/preprocessing.rs
@@ -98,7 +98,13 @@ impl Preprocessor {
             ));
         }
 
-        if batch_data.len() != num_samples * sample_size {
+        let expected_len = num_samples.checked_mul(sample_size).ok_or_else(|| {
+            MahoutError::InvalidInput(format!(
+                "Batch size overflow: num_samples {} * sample_size {}",
+                num_samples, sample_size
+            ))
+        })?;
+        if batch_data.len() != expected_len {
             return Err(MahoutError::InvalidInput(format!(
                 "Batch data length {} doesn't match num_samples {} * 
sample_size {}",
                 batch_data.len(),
@@ -160,3 +166,25 @@ impl Preprocessor {
             .collect()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn validate_batch_rejects_size_overflow() {
+        let data = [0.0_f64; 4];
+        let err = Preprocessor::validate_batch(&data, usize::MAX, 2, 1)
+            .expect_err("expected overflow error");
+        match err {
+            MahoutError::InvalidInput(msg) => {
+                assert!(
+                    msg.contains("overflow"),
+                    "unexpected error message: {}",
+                    msg
+                );
+            }
+            other => panic!("expected InvalidInput, got {:?}", other),
+        }
+    }
+}