fix(auth): enforce dynamodb:LeadingKeys on Query, batch and transaction operations

crates/engine/src/batch_get_item.rs

@@ -16,12 +16,10 @@ use crate::OperationContext;
 use crate::capacity_helpers;
 use crate::create_table::storage_err_to_dynamo;
 use crate::expression_helpers::build_expression_maps;
+use crate::request_limits;
 use crate::serialize_output;
 use crate::{DispatchMetrics, DispatchResult};
 
-/// Maximum number of keys across all tables in a single `BatchGetItem` request.
-const MAX_BATCH_GET_KEYS: usize = 100;
-
 /// Handle a `BatchGetItem` request.
 ///
 /// Reads items from one or more tables by primary key. Each table's keys are
@@ -38,44 +36,15 @@ pub async fn handle_batch_get_item(
     let input: BatchGetItemInput =
         serde_json::from_value(body).map_err(crate::deserialize_error)?;
 
-    // Validate: RequestItems must not be empty. The message is the request
-    // model's length constraint (measured; note "Value at", with no value
-    // echoed). BatchWriteItem keeps its distinct required-parameter sentence,
-    // which is likewise measured.
-    if input.request_items.is_empty() {
-        return Err(DynamoDbError::ValidationException(
-            "1 validation error detected: Value at 'RequestItems' failed to satisfy \
-             constraint: Member must have length greater than or equal to 1"
-                .to_owned(),
-        ));
-    }
-
-    // Validate: per-table keys <= 100
-    for (table_name, ka) in &input.request_items {
-        if ka.keys.len() > MAX_BATCH_GET_KEYS {
-            return Err(DynamoDbError::ValidationException(format!(
-                "1 validation error detected: Value at 'RequestItems.{table_name}.member.Keys' failed to satisfy constraint: \
-                 Member must have length less than or equal to 100"
-            )));
-        }
-    }
-
-    // Validate: total keys across all tables <= 100
-    let total_keys: usize = input.request_items.values().map(|ka| ka.keys.len()).sum();
-    if total_keys > MAX_BATCH_GET_KEYS {
-        return Err(DynamoDbError::ValidationException(
-            "Too many items requested for the BatchGetItem call".to_owned(),
-        ));
-    }
-
-    // Validate: each table must have at least one key
-    for (table_name, ka) in &input.request_items {
-        if ka.keys.is_empty() {
-            return Err(DynamoDbError::ValidationException(format!(
-                "1 validation error detected: Value '[]' at 'requestItems.{table_name}.member.keys' failed to satisfy constraint: Member must have length greater than or equal to 1"
-            )));
-        }
-    }
+    // Validate the table and key counts: RequestItems not empty, at most 100
+    // tables, at most 100 keys per table and in total, at least one key per
+    // table. The server runs the same checks before authorization.
+    let counts: Vec<(&str, Option<usize>)> = input
+        .request_items
+        .iter()
+        .map(|(table_name, ka)| (table_name.as_str(), Some(ka.keys.len())))
+        .collect();
+    request_limits::check_batch_counts("BatchGetItem", &counts)?;
 
     // Reject mixing expression and non-expression parameters anywhere in the
     // request. The check is request-wide: ProjectionExpression on one table and
@@ -189,8 +158,7 @@ pub async fn handle_batch_get_item(
 
     for (table_name, ka) in &input.request_items {
         let key_info = ctx
-            .storage
-            .table_key_info(&ctx.account_id, table_name)
+            .table_key_info_uncached(table_name)
             .await
             .map_err(storage_err_to_dynamo)?;
 

crates/engine/src/batch_write_item.rs

@@ -21,13 +21,11 @@ use extenddb_core::validation::{
 use crate::OperationContext;
 use crate::capacity_helpers;
 use crate::create_table::storage_err_to_dynamo;
+use crate::request_limits;
 use crate::serialize_output;
 use crate::stream_capture;
 use crate::{DispatchMetrics, DispatchResult};
 
-/// Maximum number of write requests across all tables in a single `BatchWriteItem`.
-const MAX_BATCH_WRITE_ITEMS: usize = 25;
-
 /// Handle a `BatchWriteItem` request.
 ///
 /// Writes (put or delete) items across one or more tables. Each operation is
@@ -44,45 +42,15 @@ pub async fn handle_batch_write_item(
     let input: BatchWriteItemInput =
         serde_json::from_value(body).map_err(crate::deserialize_error)?;
 
-    // Validate: RequestItems must not be empty
-    if input.request_items.is_empty() {
-        return Err(DynamoDbError::ValidationException(
-            "The requestItems parameter is required for BatchWriteItem".to_owned(),
-        ));
-    }
-
-    // Validate: per-table operations <= 25
-    for (table_name, reqs) in &input.request_items {
-        if reqs.len() > MAX_BATCH_WRITE_ITEMS {
-            let items_repr = reqs
-                .iter()
-                .map(|_| "WriteRequest")
-                .collect::<Vec<_>>()
-                .join(", ");
-            return Err(DynamoDbError::ValidationException(format!(
-                "1 validation error detected: Value '{{{table_name}=[{items_repr}]}}' at 'requestItems' failed to satisfy constraint: \
-                 Map value must satisfy constraint: [Member must have length less than or equal to 25, \
-                 Member must have length greater than or equal to 1]"
-            )));
-        }
-    }
-
-    // Validate: total operations across all tables <= 25
-    let total_ops: usize = input.request_items.values().map(std::vec::Vec::len).sum();
-    if total_ops > MAX_BATCH_WRITE_ITEMS {
-        return Err(DynamoDbError::ValidationException(
-            "Too many items requested for the BatchWriteItem call".to_owned(),
-        ));
-    }
-
-    // Validate: each table must have at least one request
-    for (table_name, reqs) in &input.request_items {
-        if reqs.is_empty() {
-            return Err(DynamoDbError::ValidationException(format!(
-                "1 validation error detected: Value '[]' at 'requestItems.{table_name}.member' failed to satisfy constraint: Member must have length greater than or equal to 1"
-            )));
-        }
-    }
+    // Validate the table and request counts: RequestItems not empty, at most
+    // 25 tables, at most 25 requests per table and in total, at least one
+    // request per table. The server runs the same checks before authorization.
+    let counts: Vec<(&str, Option<usize>)> = input
+        .request_items
+        .iter()
+        .map(|(table_name, reqs)| (table_name.as_str(), Some(reqs.len())))
+        .collect();
+    request_limits::check_batch_counts("BatchWriteItem", &counts)?;
 
     // Validate: each WriteRequest has exactly one of PutRequest or DeleteRequest
     for reqs in input.request_items.values() {
@@ -113,8 +81,7 @@ pub async fn handle_batch_write_item(
 
     for (table_name, reqs) in &input.request_items {
         let key_info = ctx
-            .storage
-            .table_key_info(&ctx.account_id, table_name)
+            .table_key_info_uncached(table_name)
             .await
             .map_err(storage_err_to_dynamo)?;
 

crates/engine/src/create_table.rs

@@ -110,7 +110,16 @@ pub async fn handle_create_table(
     serialize_output(&output)
 }
 
-pub(crate) fn storage_err_to_dynamo(e: extenddb_storage::error::StorageError) -> DynamoDbError {
+/// Map a storage-layer error to its `DynamoDB` wire error.
+///
+/// Exposed so the server's authorization path classifies a key-info lookup
+/// fault the same way the engine classifies the identical fault on the data
+/// path: a transient or connection error becomes `ServiceUnavailable`, the
+/// catalog and backend-internal errors become `InternalServerError`, and a
+/// not-active table becomes `ResourceInUseException`. Without a single mapper
+/// the same database hiccup returned different wire errors depending on which
+/// layer hit it first.
+pub fn storage_err_to_dynamo(e: extenddb_storage::error::StorageError) -> DynamoDbError {
     use extenddb_storage::error::StorageError;
     match e {
         StorageError::TableNotFound(name) => DynamoDbError::ResourceNotFoundException(

crates/engine/src/lib.rs

@@ -28,6 +28,7 @@ mod list_tables;
 mod put_item;
 mod query;
 mod read_helpers;
+mod request_limits;
 mod scan;
 mod search_vectors;
 pub mod stream_capture;
@@ -44,6 +45,7 @@ mod vector_gate;
 pub use batch_get_item::handle_batch_get_item;
 pub use batch_write_item::handle_batch_write_item;
 pub use create_table::handle_create_table;
+pub use create_table::storage_err_to_dynamo;
 pub use delete_item::handle_delete_item;
 pub use delete_table::handle_delete_table;
 pub use describe_endpoints::handle_describe_endpoints;
@@ -54,6 +56,7 @@ pub use import_export::{handle_export_table, handle_import_table};
 pub use list_tables::handle_list_tables;
 pub use put_item::handle_put_item;
 pub use query::handle_query;
+pub use request_limits::check_request_item_counts;
 pub use scan::handle_scan;
 pub use search_vectors::handle_search_vectors;
 pub use streams::{
@@ -71,6 +74,7 @@ use std::sync::Arc;
 
 use extenddb_core::error::DynamoDbError;
 use extenddb_core::limits::LimitsConfig;
+use extenddb_core::types::TableKeyInfo;
 
 /// Check whether an operation name is recognized by the dispatch table.
 ///
@@ -315,17 +319,24 @@ impl DispatchResult {
     }
 }
 
+/// The `TableKeyInfo` the authorization layer read for a request, by table
+/// name. A `None` value records a table that authorization found missing.
+pub type AuthorizedKeyInfo = std::collections::HashMap<String, Option<TableKeyInfo>>;
+
 /// Context passed to every operation handler.
 /// Fix #9: region and `account_id` use Arc<str> to avoid per-request cloning.
 ///
 /// # Catalog Query Discipline (P118)
 ///
-/// For single-table item operations (`GetItem`, `PutItem`, `DeleteItem`,
-/// `UpdateItem`, `Query`, `Scan`), the auth layer pre-fetches `TableKeyInfo`
-/// and stores it in `pre_fetched_key_info`. Engine handlers MUST use this
-/// pre-fetched value instead of calling `storage.table_key_info()` directly.
-/// New per-request catalog roundtrips require justification in the discussion
-/// file and principal reviewer approval.
+/// The auth layer reads `TableKeyInfo` for every table whose partition key it
+/// evaluates (`dynamodb:LeadingKeys`) and stores what it read in
+/// `authorized_key_info`. Engine handlers MUST resolve those tables from this
+/// snapshot instead of reading the catalog again, through
+/// [`OperationContext::table_key_info`] or
+/// [`OperationContext::table_key_info_uncached`]. Authorization and the engine
+/// then act on the same key schema, and on the same answer to whether the
+/// table exists. New per-request catalog roundtrips require justification in
+/// the discussion file and principal reviewer approval.
 pub struct OperationContext {
     pub storage: Arc<dyn extenddb_storage::StorageEngine>,
     pub limits: Arc<LimitsConfig>,
@@ -337,34 +348,56 @@ pub struct OperationContext {
     /// Allowed directories for export file operations. Empty means exports
     /// are disabled (secure default).
     pub export_paths: Arc<[Arc<PathBuf>]>,
-    /// Pre-fetched `TableKeyInfo` from the auth layer (P118 optimization #2).
-    /// Populated for single-table item-level operations; `None` for table-level
-    /// and batch/transact operations.
-    pub pre_fetched_key_info: Option<extenddb_core::types::TableKeyInfo>,
+    /// The `TableKeyInfo` the auth layer read for this request (P118
+    /// optimization #2). Populated for single-table item operations, `Query`,
+    /// `Scan`, and the batch and transaction operations; empty for table-level
+    /// operations. In developer mode authorization is skipped after the
+    /// single-table read, so only a single-table operation has an entry.
+    pub authorized_key_info: AuthorizedKeyInfo,
     /// Auth/authz cache handles. Used by control-plane and tagging operations
     /// to issue write-through cache invalidations after the underlying state
     /// changes (e.g. `TagResource` invalidates the resource-tags cache).
     pub auth_cache: extenddb_auth::AuthCacheRegistry,
-    /// Optional cached `TableKeyInfo` lookup. When set, batch / transact /
-    /// multi-table engine handlers route through this instead of calling
-    /// `storage.table_key_info` directly. When unset (e.g. unit tests), the
-    /// engine falls back to direct storage lookups.
+    /// Optional cached `TableKeyInfo` lookup, used by
+    /// [`OperationContext::table_key_info`] for a table the auth layer did not
+    /// read. When unset (e.g. unit tests), the engine falls back to direct
+    /// storage lookups.
     pub table_key_info_lookup: Option<Arc<dyn extenddb_storage::TableKeyInfoLookup>>,
 }
 
 impl OperationContext {
-    /// Return pre-fetched `TableKeyInfo` if available and matching the requested
-    /// table, otherwise fetch from storage. This is the single entry point for
-    /// obtaining `TableKeyInfo` in engine handlers (P118 Catalog Query Discipline).
+    /// The auth layer's answer for `table_name`, if it read that table: the key
+    /// info it used, or `TableNotFound` when it found the table missing. A
+    /// missing table stays missing for the rest of the request even if it has
+    /// since become ACTIVE, because authorization left `dynamodb:LeadingKeys`
+    /// absent for it and an absent key satisfies a `ForAllValues` condition.
+    /// The snapshot is authoritative for every table it names: an entry for
+    /// another account is an internal error, never a reason to read the table
+    /// again.
+    fn authorized_table_key_info(
+        &self,
+        table_name: &str,
+    ) -> Option<Result<TableKeyInfo, extenddb_storage::error::StorageError>> {
+        use extenddb_storage::error::StorageError;
+        Some(match self.authorized_key_info.get(table_name)? {
+            Some(ki) if *ki.account_id == *self.account_id => Ok(ki.clone()),
+            Some(_) => Err(StorageError::Internal(format!(
+                "authorized key info for table {table_name} belongs to another account"
+            ))),
+            None => Err(StorageError::TableNotFound(table_name.to_owned())),
+        })
+    }
+
+    /// `TableKeyInfo` for a single-table operation: the auth layer's snapshot
+    /// when it read the table, otherwise the cached lookup, otherwise storage.
+    /// This is the entry point for single-table handlers (P118 Catalog Query
+    /// Discipline).
     pub(crate) async fn table_key_info(
         &self,
         table_name: &str,
-    ) -> Result<extenddb_core::types::TableKeyInfo, extenddb_storage::error::StorageError> {
-        if let Some(ref ki) = self.pre_fetched_key_info
-            && ki.table_name == table_name
-            && *ki.account_id == *self.account_id
-        {
-            return Ok(ki.clone());
+    ) -> Result<TableKeyInfo, extenddb_storage::error::StorageError> {
+        if let Some(result) = self.authorized_table_key_info(table_name) {
+            return result;
         }
         if let Some(ref lookup) = self.table_key_info_lookup {
             return lookup.lookup(&self.account_id, table_name).await;
@@ -373,6 +406,23 @@ impl OperationContext {
             .table_key_info(&self.account_id, table_name)
             .await
     }
+
+    /// `TableKeyInfo` for a batch or transaction handler: the auth layer's
+    /// snapshot when it read the table, otherwise a direct storage read. The
+    /// fallback skips the cache, as these handlers always have, so a request
+    /// that authorization did not evaluate (developer mode) sees the catalog as
+    /// it is now rather than a cached or negative entry.
+    pub(crate) async fn table_key_info_uncached(
+        &self,
+        table_name: &str,
+    ) -> Result<TableKeyInfo, extenddb_storage::error::StorageError> {
+        if let Some(result) = self.authorized_table_key_info(table_name) {
+            return result;
+        }
+        self.storage
+            .table_key_info(&self.account_id, table_name)
+            .await
+    }
 }
 
 /// Dispatch an operation by name.