From 7b7585e8977399e0902a4c7598f48d9a1808cc14 Mon Sep 17 00:00:00 2001 From: Samuel Gbafa Date: Fri, 31 Jul 2026 21:04:55 -0400 Subject: [PATCH] perf: implement TC-410 optimization --- tinycloud-node-server/src/auth_guards.rs | 74 +++++++++-- tinycloud-node-server/src/routes/mod.rs | 151 ++++++++++------------- 2 files changed, 131 insertions(+), 94 deletions(-) diff --git a/tinycloud-node-server/src/auth_guards.rs b/tinycloud-node-server/src/auth_guards.rs index 53134805..49c9e5a2 100644 --- a/tinycloud-node-server/src/auth_guards.rs +++ b/tinycloud-node-server/src/auth_guards.rs @@ -2,7 +2,7 @@ use anyhow::Result; use rocket::{ data::{Capped, FromData}, futures::io::AsyncRead, - http::{ContentType, Header, Status}, + http::{ContentType, Header, HeaderMap, Status}, request::{FromRequest, Outcome, Request}, response::{Responder, Response}, serde::json::Json, @@ -443,9 +443,19 @@ const STORED_OBJECT_HEADERS: &[&str] = &[ "content-disposition", ]; +const X_TINYCLOUD_META_PREFIX: &str = "x-tinycloud-meta-"; + +/// Allocation-free, ASCII-case-insensitive check for whether `name` is +/// permitted to become stored object metadata (TC-410). pub(crate) fn is_storable_object_header(name: &str) -> bool { - let name = name.to_ascii_lowercase(); - STORED_OBJECT_HEADERS.contains(&name.as_str()) || name.starts_with("x-tinycloud-meta-") + STORED_OBJECT_HEADERS + .iter() + .any(|header| name.eq_ignore_ascii_case(header)) + || name + .as_bytes() + .get(..X_TINYCLOUD_META_PREFIX.len()) + .map(|prefix| prefix.eq_ignore_ascii_case(X_TINYCLOUD_META_PREFIX.as_bytes())) + .unwrap_or(false) } pub(crate) fn filter_stored_object_metadata(metadata: Metadata) -> Metadata { @@ -477,16 +487,60 @@ pub(crate) fn is_replayable_object_header(name: &str) -> bool { .any(|header| name.eq_ignore_ascii_case(header)) } +/// Lifetime-bound `/invoke` request guard (TC-410). Borrows the last +/// Rocket-parsed value for authorization-adjacent/control headers with no +/// allocation, and owns only the small allowlist of headers permitted to +/// become stored object metadata. `Authorization`, `Cookie`, and every other +/// unlisted header are never copied here; `Authorization`'s value is read +/// only by [`crate::authorization::AuthHeaderGetter`]. +pub struct InvokeHeaders<'r> { + pub accept: Option<&'r str>, + pub content_type: Option<&'r str>, + pub if_match: Option<&'r str>, + pub if_none_match: Option<&'r str>, + pub expected_version: Option<&'r str>, + pub max_response_bytes: Option<&'r str>, + pub limit: Option<&'r str>, + pub cursor: Option<&'r str>, + pub metadata: Metadata, +} + +/// A header nominated by a comma-separated `Connection` token is hop-by-hop +/// for this request and must never be persisted, even if it would otherwise +/// be storable object metadata. Scans `Connection` values directly instead +/// of allocating a token set. +fn connection_nominates(headers: &HeaderMap<'_>, name: &str) -> bool { + headers + .get("connection") + .flat_map(|value| value.split(',')) + .any(|token| token.trim().eq_ignore_ascii_case(name)) +} + #[async_trait] -impl<'r> FromRequest<'r> for ObjectHeaders { +impl<'r> FromRequest<'r> for InvokeHeaders<'r> { type Error = anyhow::Error; async fn from_request(request: &'r Request<'_>) -> Outcome { - let md: BTreeMap = request - .headers() - .iter() - .map(|h| (h.name.into_string(), h.value.to_string())) - .collect(); - Outcome::Success(ObjectHeaders(Metadata(md))) + let headers = request.headers(); + + let mut metadata = BTreeMap::new(); + for header in headers.iter() { + let name = header.name.as_str(); + if is_storable_object_header(name) && !connection_nominates(headers, name) { + metadata.insert(name.to_string(), header.value.to_string()); + } + } + + Outcome::Success(InvokeHeaders { + accept: headers.get("accept").last(), + content_type: headers.get("content-type").last(), + if_match: headers.get("if-match").last(), + if_none_match: headers.get("if-none-match").last(), + expected_version: headers.get("x-tinycloud-expected-version").last(), + max_response_bytes: headers.get("x-tinycloud-max-response-bytes").last(), + limit: headers.get("x-tinycloud-limit").last(), + cursor: headers.get("x-tinycloud-cursor").last(), + metadata: Metadata(metadata), + }) } } diff --git a/tinycloud-node-server/src/routes/mod.rs b/tinycloud-node-server/src/routes/mod.rs index 86837cb6..9ecbb997 100644 --- a/tinycloud-node-server/src/routes/mod.rs +++ b/tinycloud-node-server/src/routes/mod.rs @@ -19,7 +19,7 @@ use tracing::{info_span, Instrument}; use crate::{ auth_guards::{ filter_stored_object_metadata, is_storable_object_header, DataIn, DataOut, InvOut, - ObjectHeaders, + InvokeHeaders, }, authorization::AuthHeaderGetter, config::Config, @@ -599,7 +599,7 @@ pub struct RevokeResponse { pub async fn invoke( i: AuthHeaderGetter, req_span: TracingSpan, - headers: ObjectHeaders, + headers: InvokeHeaders<'_>, data: DataIn<'_>, staging: &State, tinycloud: &State, @@ -633,7 +633,7 @@ pub async fn invoke( pub async fn invoke( i: AuthHeaderGetter, req_span: TracingSpan, - headers: ObjectHeaders, + headers: InvokeHeaders<'_>, data: DataIn<'_>, staging: &State, tinycloud: &State, @@ -674,23 +674,6 @@ type KvInputMap = HashMap< >; type ExpectedKvBatchInputs = BTreeMap; -fn metadata_header<'a>(metadata: &'a Metadata, name: &str) -> Option<&'a str> { - metadata - .0 - .iter() - .find(|(key, _)| key.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) -} - -fn take_metadata_header(metadata: &mut Metadata, name: &str) -> Option { - let key = metadata - .0 - .keys() - .find(|key| key.eq_ignore_ascii_case(name))? - .clone(); - metadata.0.remove(&key) -} - fn parse_strong_blake3_etag(value: &str) -> Result<[u8; 32], (Status, String)> { let value = value.trim(); let digest = value @@ -724,10 +707,10 @@ fn parse_strong_blake3_etag(value: &str) -> Result<[u8; 32], (Status, String)> { } fn parse_positive_u64_header( - metadata: &mut Metadata, + value: Option<&str>, name: &str, ) -> Result, (Status, String)> { - take_metadata_header(metadata, name) + value .map(|value| { let parsed = value.trim().parse::().map_err(|_| { ( @@ -855,7 +838,7 @@ fn kv_list_target(capabilities: &[Capability]) -> Option<(SpaceId, Path)> { fn kv_invoke_options( invocation: &InvocationInfo, - headers: &mut ObjectHeaders, + headers: &InvokeHeaders<'_>, multipart: bool, cursor_key: &[u8; 32], ) -> Result { @@ -871,7 +854,7 @@ fn kv_invoke_options( #[cfg(test)] fn kv_invoke_options_for_capabilities( capabilities: &[Capability], - headers: &mut ObjectHeaders, + headers: &InvokeHeaders<'_>, multipart: bool, ) -> Result { kv_invoke_options_for_capabilities_with_cursor(capabilities, headers, multipart, None, None) @@ -879,14 +862,14 @@ fn kv_invoke_options_for_capabilities( fn kv_invoke_options_for_capabilities_with_cursor( capabilities: &[Capability], - headers: &mut ObjectHeaders, + headers: &InvokeHeaders<'_>, multipart: bool, cursor_key: Option<&[u8; 32]>, cursor_subject: Option<&str>, ) -> Result { - let if_match = take_metadata_header(&mut headers.0, "if-match"); - let if_none_match = take_metadata_header(&mut headers.0, "if-none-match"); - let expected_version = take_metadata_header(&mut headers.0, "x-tinycloud-expected-version"); + let if_match = headers.if_match; + let if_none_match = headers.if_none_match; + let expected_version = headers.expected_version; if expected_version.is_some() { return Err(( Status::BadRequest, @@ -952,13 +935,13 @@ fn kv_invoke_options_for_capabilities_with_cursor( let (space, path, _) = &mutation_targets[0]; preconditions.insert( (space.clone(), path.clone()), - KvPrecondition::Matches(parse_strong_blake3_etag(&value)?), + KvPrecondition::Matches(parse_strong_blake3_etag(value)?), ); } let max_response_bytes = - parse_positive_u64_header(&mut headers.0, "x-tinycloud-max-response-bytes")?; - let list_limit = parse_positive_u64_header(&mut headers.0, "x-tinycloud-limit")? + parse_positive_u64_header(headers.max_response_bytes, "x-tinycloud-max-response-bytes")?; + let list_limit = parse_positive_u64_header(headers.limit, "x-tinycloud-limit")? .map(|limit| { if limit > 1000 { Err(( @@ -970,7 +953,8 @@ fn kv_invoke_options_for_capabilities_with_cursor( } }) .transpose()?; - let list_cursor = take_metadata_header(&mut headers.0, "x-tinycloud-cursor") + let list_cursor = headers + .cursor .map(|value| { let key = cursor_key .ok_or_else(|| (Status::BadRequest, "Invalid KV list cursor".to_string()))?; @@ -980,7 +964,7 @@ fn kv_invoke_options_for_capabilities_with_cursor( .ok_or_else(|| (Status::BadRequest, "Invalid KV list cursor".to_string()))?; let limit = list_limit .ok_or_else(|| (Status::BadRequest, "Invalid KV list cursor".to_string()))?; - decode_kv_cursor(key, &value, subject, &space, &prefix, limit) + decode_kv_cursor(key, value, subject, &space, &prefix, limit) .map_err(|_| (Status::BadRequest, "Invalid KV list cursor".to_string())) }) .transpose()?; @@ -993,8 +977,9 @@ fn kv_invoke_options_for_capabilities_with_cursor( }) } -fn is_multipart(headers: &ObjectHeaders) -> bool { - metadata_header(&headers.0, "content-type") +fn is_multipart(headers: &InvokeHeaders<'_>) -> bool { + headers + .content_type .map(|value| { value .to_ascii_lowercase() @@ -1182,7 +1167,7 @@ async fn copy_multipart_field_to_stage( async fn build_batch_kv_inputs( data: rocket::Data<'_>, - headers: &ObjectHeaders, + headers: &InvokeHeaders<'_>, expected: &ExpectedKvBatchInputs, staging: &State, tinycloud: &State, @@ -1193,7 +1178,7 @@ async fn build_batch_kv_inputs( return Ok(HashMap::new()); } - let content_type = metadata_header(&headers.0, "content-type").ok_or_else(|| { + let content_type = headers.content_type.ok_or_else(|| { ( Status::BadRequest, "Missing multipart content-type".to_string(), @@ -1300,7 +1285,7 @@ fn classify_invocation_time_rejection( async fn invoke_impl( i: AuthHeaderGetter, req_span: TracingSpan, - mut headers: ObjectHeaders, + headers: InvokeHeaders<'_>, data: DataIn<'_>, staging: &State, tinycloud: &State, @@ -1426,10 +1411,9 @@ async fn invoke_impl( .collect(); if !duckdb_caps.is_empty() { - let arrow_format = headers.0 .0.iter().any(|(k, v)| { - k.eq_ignore_ascii_case("accept") - && v.contains("application/vnd.apache.arrow.stream") - }); + let arrow_format = headers + .accept + .is_some_and(|v| v.contains("application/vnd.apache.arrow.stream")); let result = handle_duckdb_invoke( admitted, data, @@ -1469,7 +1453,7 @@ async fn invoke_impl( let put_caps = kv_put_capabilities(&admitted.invocation().0); let is_multipart_request = is_multipart(&headers); - let kv_options = kv_invoke_options(&admitted.invocation().0, &mut headers, is_multipart_request, &tinycloud.kv_cursor_key())?; + let kv_options = kv_invoke_options(&admitted.invocation().0, &headers, is_multipart_request, &tinycloud.kv_cursor_key())?; let expected_batch_inputs = if is_multipart_request && !put_caps.is_empty() { Some(validate_kv_batch_capabilities(&admitted.invocation().0, &put_caps)?) } else { @@ -1567,7 +1551,7 @@ async fn invoke_impl( let mut inputs = HashMap::new(); inputs.insert( (space.clone(), path.clone()), - (filter_stored_object_metadata(headers.0), stage), + (filter_stored_object_metadata(headers.metadata), stage), ); Ok(inputs) } @@ -3199,32 +3183,30 @@ mod tests { assert_eq!(buffered_hash, original_hash); } + fn test_invoke_headers<'r>() -> InvokeHeaders<'r> { + InvokeHeaders { + accept: None, + content_type: None, + if_match: None, + if_none_match: None, + expected_version: None, + max_response_bytes: None, + limit: None, + cursor: None, + metadata: Metadata(BTreeMap::new()), + } + } + #[tokio::test] - async fn bounded_kv_headers_are_positive_and_removed_from_object_metadata() { - let mut metadata = Metadata(BTreeMap::from([ - ( - "X-TinyCloud-Max-Response-Bytes".to_string(), - "1048576".to_string(), - ), - ("content-type".to_string(), "text/plain".to_string()), - ])); + async fn bounded_kv_headers_are_positive() { assert_eq!( - parse_positive_u64_header(&mut metadata, "x-tinycloud-max-response-bytes").unwrap(), + parse_positive_u64_header(Some("1048576"), "x-tinycloud-max-response-bytes").unwrap(), Some(1_048_576) ); - assert!(metadata_header(&metadata, "x-tinycloud-max-response-bytes").is_none()); - assert_eq!( - metadata_header(&metadata, "content-type"), - Some("text/plain") - ); for invalid in ["0", "-1", "many"] { - let mut metadata = Metadata(BTreeMap::from([( - "x-tinycloud-limit".to_string(), - invalid.to_string(), - )])); assert_eq!( - parse_positive_u64_header(&mut metadata, "x-tinycloud-limit") + parse_positive_u64_header(Some(invalid), "x-tinycloud-limit") .unwrap_err() .0, Status::BadRequest @@ -3237,13 +3219,13 @@ mod tests { let space = test_space_id("conditional-kv-options"); let capability = kv_put_capability(&space, "files/report.txt"); - let mut create_headers = ObjectHeaders(Metadata(BTreeMap::from([( - "If-None-Match".to_string(), - "*".to_string(), - )]))); + let create_headers = InvokeHeaders { + if_none_match: Some("*"), + ..test_invoke_headers() + }; let create = kv_invoke_options_for_capabilities( std::slice::from_ref(&capability), - &mut create_headers, + &create_headers, false, ) .unwrap(); @@ -3254,15 +3236,15 @@ mod tests { )), Some(&KvPrecondition::DoesNotExist) ); - assert!(metadata_header(&create_headers.0, "if-none-match").is_none()); let digest = [7u8; 32]; - let mut replace_headers = ObjectHeaders(Metadata(BTreeMap::from([( - "If-Match".to_string(), - format!("\"blake3-{}\"", hex::encode(digest)), - )]))); + let replace_etag = format!("\"blake3-{}\"", hex::encode(digest)); + let replace_headers = InvokeHeaders { + if_match: Some(&replace_etag), + ..test_invoke_headers() + }; let replace = - kv_invoke_options_for_capabilities(&[capability], &mut replace_headers, false).unwrap(); + kv_invoke_options_for_capabilities(&[capability], &replace_headers, false).unwrap(); assert_eq!( replace .preconditions @@ -3278,23 +3260,24 @@ mod tests { kv_put_capability(&space, "a"), kv_put_capability(&space, "b"), ]; - let mut headers = ObjectHeaders(Metadata(BTreeMap::from([( - "If-Match".to_string(), - format!("\"blake3-{}\"", hex::encode([1u8; 32])), - )]))); + let etag = format!("\"blake3-{}\"", hex::encode([1u8; 32])); + let headers = InvokeHeaders { + if_match: Some(&etag), + ..test_invoke_headers() + }; assert_eq!( - kv_invoke_options_for_capabilities(&capabilities, &mut headers, false) + kv_invoke_options_for_capabilities(&capabilities, &headers, false) .unwrap_err() .0, Status::BadRequest ); - let mut headers = ObjectHeaders(Metadata(BTreeMap::from([( - "If-None-Match".to_string(), - "*".to_string(), - )]))); + let headers = InvokeHeaders { + if_none_match: Some("*"), + ..test_invoke_headers() + }; assert_eq!( - kv_invoke_options_for_capabilities(&capabilities[..1], &mut headers, true) + kv_invoke_options_for_capabilities(&capabilities[..1], &headers, true) .unwrap_err() .0, Status::BadRequest