From 434bbb27b151714fbb268797a5a568cf170e840d Mon Sep 17 00:00:00 2001
From: Kacy Fortner <kacy@kacyfortner.com>
Date: Sat, 14 Feb 2026 17:46:37 -0500
Subject: [PATCH] bound total vector allocation in persistence deserialization

add validate_vector_total(dim, count) check after individual dim/count
validation in snapshot deserialization. without this, a crafted snapshot
with 65536 dims x 10M vectors would attempt ~2.6 TB of allocation.
the new cap limits total f32 elements to 1 billion (~4 GB).
---
 crates/ember-persistence/src/format.rs   | 18 ++++++++++++++++++
 crates/ember-persistence/src/snapshot.rs |  2 ++
 2 files changed, 20 insertions(+)

diff --git a/crates/ember-persistence/src/format.rs b/crates/ember-persistence/src/format.rs
index fa14e7f0..fca09bcc 100644
--- a/crates/ember-persistence/src/format.rs
+++ b/crates/ember-persistence/src/format.rs
@@ -281,6 +281,24 @@ pub const MAX_PERSISTED_VECTOR_DIMS: u32 = 65_536;
 /// Prevents corrupt count fields from causing unbounded loops.
 pub const MAX_PERSISTED_VECTOR_COUNT: u32 = 10_000_000;
 
+/// Maximum total f32 elements (dim * count) for vector deserialization.
+/// Caps total allocation at ~4 GB. Without this, a crafted file with
+/// 65536 dims x 10M vectors would attempt ~2.6 TB.
+pub const MAX_PERSISTED_VECTOR_TOTAL_FLOATS: u64 = 1_000_000_000;
+
+/// Validates that the total vector element budget (dim * count) is within
+/// bounds. Call after validating dim and count individually.
+pub fn validate_vector_total(dim: u32, count: u32) -> Result<(), FormatError> {
+    let total = dim as u64 * count as u64;
+    if total > MAX_PERSISTED_VECTOR_TOTAL_FLOATS {
+        return Err(FormatError::InvalidData(format!(
+            "vector total elements ({dim} dims x {count} vectors = {total}) \
+             exceeds max {MAX_PERSISTED_VECTOR_TOTAL_FLOATS}"
+        )));
+    }
+    Ok(())
+}
+
 /// Verifies that two CRC32 values match.
 pub fn verify_crc32_values(computed: u32, stored: u32) -> Result<(), FormatError> {
     if computed != stored {
diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs
index fca48bbe..2e51d638 100644
--- a/crates/ember-persistence/src/snapshot.rs
+++ b/crates/ember-persistence/src/snapshot.rs
@@ -144,6 +144,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result<SnapValue, FormatError> {
                     format::MAX_PERSISTED_VECTOR_COUNT
                 )));
             }
+            format::validate_vector_total(dim, count)?;
             let mut elements = Vec::with_capacity(format::capped_capacity(count));
             for _ in 0..count {
                 let name = read_snap_string(r, "vector element name")?;
@@ -683,6 +684,7 @@ impl SnapshotReader {
                             format::MAX_PERSISTED_VECTOR_COUNT
                         )));
                     }
+                    format::validate_vector_total(dim, count)?;
                     format::write_u32(&mut buf, count)?;
                     let mut elements = Vec::with_capacity(format::capped_capacity(count));
                     for _ in 0..count {