From 43200aa0d91c5d638380c1bda158cc6f36fd3592 Mon Sep 17 00:00:00 2001 From: Lee Hannigan Date: Tue, 7 Jul 2026 15:54:41 +0000 Subject: [PATCH] fix: reject cross-segment ExclusiveStartKey in parallel Scan A parallel Scan (TotalSegments/Segment) accepted an ExclusiveStartKey that belongs to a different segment, silently returning a truncated or empty page. Validate the start key against the same segment-assignment function used for the scan and reject a mismatch with a ValidationException naming the correct Segment, so a key returned as a LastEvaluatedKey for one segment is only valid when the same segment is re-scanned. Verified against DynamoDB. Adds a scan integration test. Signed-off-by: Lee Hannigan --- .../storage-postgres/src/data/query_scan.rs | 24 +++++ tests/rust/src/scan.rs | 95 +++++++++++++++++++ 2 files changed, 119 insertions(+) diff --git a/crates/storage-postgres/src/data/query_scan.rs b/crates/storage-postgres/src/data/query_scan.rs index 6d61e9af..1d3899a6 100644 --- a/crates/storage-postgres/src/data/query_scan.rs +++ b/crates/storage-postgres/src/data/query_scan.rs @@ -413,6 +413,30 @@ impl PostgresEngine { "The provided starting key is invalid: The provided key element does not match the schema".to_owned(), )); } + + // Parallel scan: the ExclusiveStartKey must belong to the segment + // being scanned. Validate against the same hashtext-based segment + // function used for assignment above, so a key returned as a + // LastEvaluatedKey for one segment is rejected when replayed against + // another (rather than silently returning a truncated/empty page). + if let (Some(seg), Some(total)) = (segment, total_segments) { + let pk_text = pk_to_text(&start_key[pk_name])?; + let in_segment: bool = + sqlx::query_scalar("SELECT (hashtext($1)::bigint & 2147483647) % $2 = $3") + .bind(pk_text.as_ref()) + .bind(total) + .bind(seg) + .fetch_one(&self.data_pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if !in_segment { + return Err(StorageError::Validation(format!( + "The provided starting key is invalid: Invalid ExclusiveStartKey. \ + Please use ExclusiveStartKey with correct Segment. \ + TotalSegments: {total} Segment: {seg}" + ))); + } + } // Actual PK/SK binding happens in execute_scan_sql. if index_name.is_some() { diff --git a/tests/rust/src/scan.rs b/tests/rust/src/scan.rs index 8d699d01..d7a0fc1b 100755 --- a/tests/rust/src/scan.rs +++ b/tests/rust/src/scan.rs @@ -324,3 +324,98 @@ async fn scan_filter_contains_binary() { .collect(); assert_eq!(keys, vec!["i2"], "NOT contains b\"ell\" should match only i2"); } + +/// Parallel scan validates that an `ExclusiveStartKey` belongs to the segment +/// being scanned: a key is accepted as the start key for its own segment and +/// rejected (ValidationException) for any other segment. Regression for the +/// previously-missing cross-segment check, which silently returned a truncated +/// or empty page instead of erroring. +#[tokio::test] +async fn parallel_scan_rejects_cross_segment_exclusive_start_key() { + use aws_sdk_dynamodb::types::{ + AttributeDefinition, BillingMode, KeySchemaElement, KeyType, ScalarAttributeType, + }; + let c = client(); + let table = format!("ScanEskSeg_{}", ts()); + c.create_table() + .table_name(&table) + .billing_mode(BillingMode::PayPerRequest) + .key_schema( + KeySchemaElement::builder() + .attribute_name(HASH_KEY_S) + .key_type(KeyType::Hash) + .build() + .unwrap(), + ) + .attribute_definitions( + AttributeDefinition::builder() + .attribute_name(HASH_KEY_S) + .attribute_type(ScalarAttributeType::S) + .build() + .unwrap(), + ) + .send() + .await + .unwrap(); + wait_for_active(c, &table).await; + + for i in 0..20 { + let mut item = std::collections::HashMap::new(); + item.insert(HASH_KEY_S.into(), s(&i.to_string())); + c.put_item() + .table_name(&table) + .set_item(Some(item)) + .send() + .await + .unwrap(); + } + + const TOTAL: i32 = 4; + // Find a segment that owns at least one item, and grab one of its keys. + let mut owning_seg = None; + let mut key = None; + for seg in 0..TOTAL { + let r = c + .scan() + .table_name(&table) + .total_segments(TOTAL) + .segment(seg) + .send() + .await + .unwrap(); + if let Some(first) = r.items().first() { + owning_seg = Some(seg); + key = Some(first.clone()); + break; + } + } + let owning_seg = owning_seg.expect("some segment must own an item"); + let key = key.unwrap(); + + // Same-segment ESK is accepted. + c.scan() + .table_name(&table) + .total_segments(TOTAL) + .segment(owning_seg) + .set_exclusive_start_key(Some(key.clone())) + .send() + .await + .expect("ESK for its own segment should be accepted"); + + // Cross-segment ESK is rejected with a ValidationException. + let other_seg = (owning_seg + 1) % TOTAL; + let err = c + .scan() + .table_name(&table) + .total_segments(TOTAL) + .segment(other_seg) + .set_exclusive_start_key(Some(key)) + .send() + .await + .expect_err("ESK from a different segment must be rejected"); + let msg = format!("{err:?}"); + assert!( + msg.contains("Invalid ExclusiveStartKey"), + "expected cross-segment ExclusiveStartKey rejection, got: {msg}" + ); +}