feat: harden paged KV handoff deserialization boundary

src/distributed/kv_cache_serde/deserialize.rs

@@ -19,19 +19,44 @@
 //!
 //! Used by: decode node in disaggregated serving pipeline
 
+use std::collections::{HashMap, HashSet};
+
 use anyhow::{Context, Result, bail};
 
 use super::serialize::CACHE_FORMAT_VERSION;
 use super::types::{
-    CACHE_FORMAT_VERSION_V1, CacheType, RawTensorData, SerializableCacheState,
-    SerializableSequenceBackend, validate_raw_tensor,
+    CACHE_FORMAT_VERSION_V1, CacheIngestLimits, CacheType, ExpectedBlockGeometry, RawTensorData,
+    SerializableCacheState, SerializableSequenceBackend, validate_raw_tensor,
 };
 
-/// Deserialize a `SerializableCacheState` from the binary wire format.
+/// Deserialize a `SerializableCacheState` from the binary wire format, applying
+/// the default [`CacheIngestLimits`] frame cap.
 ///
 /// Returns the deserialized state ready for reconstruction into live
 /// cache objects.
 pub fn deserialize_cache_state(buf: &[u8]) -> Result<SerializableCacheState> {
+    deserialize_cache_state_with_limits(buf, &CacheIngestLimits::default())
+}
+
+/// Deserialize a `SerializableCacheState`, rejecting any frame larger than
+/// `limits.max_frame_bytes` before parsing.
+///
+/// The cap is the first line of defense at the deserialization boundary: an
+/// untrusted peer cannot force a multi-gigabyte allocation by sending an
+/// oversized buffer or a forged metadata length. The disaggregated decode path
+/// passes a config-derived cap; [`deserialize_cache_state`] uses the generous
+/// default backstop.
+pub fn deserialize_cache_state_with_limits(
+    buf: &[u8],
+    limits: &CacheIngestLimits,
+) -> Result<SerializableCacheState> {
+    if buf.len() > limits.max_frame_bytes {
+        bail!(
+            "cache state frame is {} bytes, exceeds limit {}",
+            buf.len(),
+            limits.max_frame_bytes
+        );
+    }
     if buf.len() < 10 {
         bail!(
             "cache state buffer too short: need at least 10 bytes, got {}",
@@ -51,6 +76,14 @@ pub fn deserialize_cache_state(buf: &[u8]) -> Result<SerializableCacheState> {
     let num_layers = u32::from_le_bytes(buf[2..6].try_into()?) as usize;
     let metadata_len = u32::from_le_bytes(buf[6..10].try_into()?) as usize;
 
+    if metadata_len > limits.max_frame_bytes {
+        bail!(
+            "cache metadata length {} exceeds frame limit {}",
+            metadata_len,
+            limits.max_frame_bytes
+        );
+    }
+
     let metadata_end = 10 + metadata_len;
     if buf.len() < metadata_end {
         bail!(
@@ -188,15 +221,186 @@ pub fn restore_into_sequence_cache_set(
     Ok(())
 }
 
-/// Restore serialized cache state directly into an active `CachePool` slot.
+/// Validate a deserialized paged handoff payload against ingest limits and,
+/// when supplied, the decode model's real KV geometry, BEFORE any pool block is
+/// acquired. This is the structural gate for a cross-node paged restore.
 ///
-/// Used by: decode-side distributed cache ingestion where paged allocator
-/// bookkeeping must be rebuilt alongside the sequence state.
+/// Always enforced (no live model needed):
+/// - the transferred block count and each slab's byte size stay within `limits`;
+/// - every block the restored table references has matching transferred
+///   contents on the SAME layer, and every transferred block is referenced by
+///   the table on its declared layer (no orphan contents, which would leak a
+///   pool block; no cross-layer aliasing; no duplicate id disagreeing on layer);
+/// - a paged sequence carries a coherent, non-empty content path (an empty
+///   block table, or a table that references blocks with no contents, is
+///   rejected; a metadata-only paged payload is not a valid cross-node handoff
+///   because a fresh decode pool holds no pre-existing block contents).
+///
+/// Enforced when `expected` is provided:
+/// - every slab's shape is exactly `[block_size, n_kv_heads, head_dim]` and its
+///   dtype matches the decode model, anchoring pool geometry before the first
+///   write captures it from the (untrusted) slab.
+///
+/// Dense payloads (no paged state, no blocks) pass unchanged.
+pub fn validate_paged_payload(
+    state: &SerializableCacheState,
+    limits: &CacheIngestLimits,
+    expected: Option<&ExpectedBlockGeometry>,
+) -> Result<()> {
+    let is_paged = state.sequence_backend == SerializableSequenceBackend::PagedKvCache
+        || state.paged_state.is_some()
+        || !state.paged_blocks.is_empty();
+    if !is_paged {
+        return Ok(());
+    }
+
+    if state.paged_blocks.len() > limits.max_paged_blocks {
+        bail!(
+            "paged handoff carries {} blocks, exceeds limit {}",
+            state.paged_blocks.len(),
+            limits.max_paged_blocks
+        );
+    }
+
+    let num_layers = expected
+        .map(|g| g.num_layers)
+        .or_else(|| state.paged_state.as_ref().map(|s| s.layers.len()));
+
+    // Per-block: byte caps, layer range, optional geometry anchor. Index each
+    // supplied block id to its declared layer for the table cross-check below.
+    let mut block_layer: HashMap<u64, usize> = HashMap::new();
+    for (i, block) in state.paged_blocks.iter().enumerate() {
+        for (name, slab) in [("keys", &block.keys), ("values", &block.values)] {
+            if slab.data.len() > limits.max_block_slab_bytes {
+                bail!(
+                    "paged block {i} {name} slab is {} bytes, exceeds limit {}",
+                    slab.data.len(),
+                    limits.max_block_slab_bytes
+                );
+            }
+        }
+        if let Some(n) = num_layers
+            && block.layer_idx >= n
+        {
+            bail!(
+                "paged block {i} layer_idx {} out of range for {n} layers",
+                block.layer_idx
+            );
+        }
+        if let Some(g) = expected {
+            let want = g.slab_shape();
+            for (name, slab) in [("keys", &block.keys), ("values", &block.values)] {
+                if slab.shape != want {
+                    bail!(
+                        "paged block {i} {name} shape {:?} does not match decode model geometry {:?}",
+                        slab.shape,
+                        want
+                    );
+                }
+                if slab.dtype != g.dtype {
+                    bail!(
+                        "paged block {i} {name} dtype {} does not match decode model dtype {}",
+                        slab.dtype,
+                        g.dtype
+                    );
+                }
+            }
+        }
+        if let Some(prev) = block_layer.insert(block.block_id, block.layer_idx)
+            && prev != block.layer_idx
+        {
+            bail!(
+                "paged block id {} supplied for both layer {prev} and layer {} (block ids are layer-unique)",
+                block.block_id,
+                block.layer_idx
+            );
+        }
+    }
+
+    let Some(paged_state) = state.paged_state.as_ref() else {
+        bail!("paged sequence payload is missing paged_state metadata");
+    };
+
+    let has_blocks = !state.paged_blocks.is_empty();
+    let mut referenced: HashSet<u64> = HashSet::new();
+    let mut referenced_count = 0usize;
+    for (layer_idx, layer) in paged_state.layers.iter().enumerate() {
+        for &block_id in &layer.block_ids {
+            referenced_count += 1;
+            referenced.insert(block_id);
+            if has_blocks {
+                match block_layer.get(&block_id) {
+                    Some(&l) if l == layer_idx => {}
+                    Some(&l) => bail!(
+                        "block {block_id} is referenced by layer {layer_idx} but its transferred contents declare layer {l}"
+                    ),
+                    None => bail!(
+                        "block {block_id} referenced by layer {layer_idx} has no transferred contents"
+                    ),
+                }
+            }
+        }
+    }
+
+    if referenced_count == 0 {
+        bail!(
+            "paged handoff carries an empty block table; a cross-node decode handoff must carry a prefilled sequence"
+        );
+    }
+    if !has_blocks {
+        bail!(
+            "paged handoff references {referenced_count} blocks but carries no block contents; metadata-only paged restore is invalid on a fresh decode pool"
+        );
+    }
+    for block in &state.paged_blocks {
+        if !referenced.contains(&block.block_id) {
+            bail!(
+                "paged block id {} has transferred contents but is not referenced by any layer table",
+                block.block_id
+            );
+        }
+    }
+
+    Ok(())
+}
+
+/// Restore serialized cache state directly into an active `CachePool` slot,
+/// applying the default [`CacheIngestLimits`] and no model geometry anchor.
+///
+/// Used by: in-process round-trips and the decode-side ingestion that does not
+/// (yet) supply the decode model's geometry. Prefer
+/// [`restore_into_cache_pool_sequence_anchored`] on the live disaggregated path
+/// so transferred slabs are pinned to the decode model.
 pub fn restore_into_cache_pool_sequence(
     state: &SerializableCacheState,
     cache_pool: &mut mlxcel_core::cache::CachePool,
     seq_id: mlxcel_core::cache::SequenceId,
 ) -> Result<()> {
+    restore_into_cache_pool_sequence_anchored(
+        state,
+        cache_pool,
+        seq_id,
+        &CacheIngestLimits::default(),
+        None,
+    )
+}
+
+/// Restore serialized cache state into an active `CachePool` slot, validating
+/// the payload against `limits` and (when supplied) the decode model's real KV
+/// geometry before any pool block is acquired.
+///
+/// Used by: decode-side distributed cache ingestion where paged allocator
+/// bookkeeping must be rebuilt alongside the sequence state, and the untrusted
+/// wire payload must be anchored to the local model.
+pub fn restore_into_cache_pool_sequence_anchored(
+    state: &SerializableCacheState,
+    cache_pool: &mut mlxcel_core::cache::CachePool,
+    seq_id: mlxcel_core::cache::SequenceId,
+    limits: &CacheIngestLimits,
+    expected: Option<&ExpectedBlockGeometry>,
+) -> Result<()> {
+    validate_paged_payload(state, limits, expected)?;
+
     // A true pool-backed cross-node handoff carries the pool block CONTENTS in
     // `paged_blocks` (#125). That path materializes fresh pool blocks rather
     // than restoring dense per-layer buffers (empty for pool-backed sequences)

src/distributed/kv_cache_serde/mod.rs

@@ -60,17 +60,18 @@ pub mod types;
 mod tests;
 
 pub use deserialize::{
-    deserialize_cache_state, reconstruct_mlx_array, restore_into_cache_pool_sequence,
-    restore_into_kv_caches, restore_into_sequence_cache_set,
+    deserialize_cache_state, deserialize_cache_state_with_limits, reconstruct_mlx_array,
+    restore_into_cache_pool_sequence, restore_into_cache_pool_sequence_anchored,
+    restore_into_kv_caches, restore_into_sequence_cache_set, validate_paged_payload,
 };
 pub use serialize::{
     CACHE_FORMAT_VERSION, extract_chunked_cache_entry, extract_kv_cache_entry,
     extract_rotating_cache_entry, serialize_cache_pool_sequence, serialize_cache_state,
     serialize_sequence_cache_set,
 };
 pub use types::{
-    CACHE_FORMAT_VERSION_V1, CACHE_FORMAT_VERSION_V2, CacheMetadata, CacheType, RawTensorData,
-    SerializableCacheEntry, SerializableCacheState, SerializablePagedBlock,
-    SerializablePagedLayerState, SerializablePagedSequenceState, SerializableSamplingState,
-    SerializableSequenceBackend, validate_raw_tensor,
+    CACHE_FORMAT_VERSION_V1, CACHE_FORMAT_VERSION_V2, CacheIngestLimits, CacheMetadata, CacheType,
+    ExpectedBlockGeometry, RawTensorData, SerializableCacheEntry, SerializableCacheState,
+    SerializablePagedBlock, SerializablePagedLayerState, SerializablePagedSequenceState,
+    SerializableSamplingState, SerializableSequenceBackend, validate_raw_tensor,
 };