From 213bcc2675dafbf615b3ea35c3facb0e5a7c21ca Mon Sep 17 00:00:00 2001 From: Phoebe Goldman Date: Tue, 28 Jul 2026 17:21:43 -0400 Subject: [PATCH 1/7] Middleware and metric to track per-database HTTP response egress This commit adds a new Prometheus metric, `spacetime_http_response_size_bytes_total`, which tracks bytes sent as responses to HTTP requests related to the database. This includes the `sql` and `call` routes, guest-defined HTTP handlers, logs, plus some misc. management routes. Notably, the `subscribe` route, which initiates a long-lived WebSocket connection, is not counted by the new metric, as we already track its egress separately. The new metric is tracked by an Axum middleware. --- crates/client-api/src/routes/database.rs | 373 ++++++++++++++++++++++- crates/core/src/host/host_controller.rs | 7 +- crates/datastore/src/db_metrics/mod.rs | 7 + 3 files changed, 375 insertions(+), 12 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 8c13f3975c2..d008be214b7 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -42,6 +42,7 @@ use spacetimedb_client_api_messages::name::{ self, DatabaseName, DomainName, MigrationPolicy, PrePublishAutoMigrateResult, PrePublishManualMigrateResult, PrePublishResult, PrettyPrintStyle, PublishOp, PublishResult, }; +use spacetimedb_datastore::db_metrics::DB_METRICS; use spacetimedb_lib::db::raw_def::v10::RawModuleDefV10; use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9; use spacetimedb_lib::{http as st_http, ConnectionId}; @@ -1621,7 +1622,10 @@ where S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, { pub fn into_router(self, ctx: S) -> axum::Router { - let db_router = axum::Router::::new() + let egress_metrics_middleware = + axum::middleware::from_fn_with_state(ctx.clone(), http_response_egress_metrics_middleware::); + + let counted_db_router = axum::Router::::new() .route("/", self.db_put) .route("/", self.db_get) .route("/", self.db_delete) @@ -1629,7 +1633,6 @@ where .route("/names", self.names_post) .route("/names", self.names_put) .route("/identity", self.identity_get) - .route("/subscribe", self.subscribe_get) .route("/call/:reducer", self.call_reducer_procedure_post) .route("/schema", self.schema_get) .route("/logs", self.logs_get) @@ -1639,7 +1642,13 @@ where .route("/pre_publish", self.pre_publish) .route("/reset", self.db_reset) .route("/lock", self.lock_post) - .route("/unlock", self.unlock_post); + .route("/unlock", self.unlock_post) + .route_layer(egress_metrics_middleware.clone()); + + // Add the subscribe route after `egress_metrics_middleware` + // so that its egress bytes don't get counted into `http_response_size_bytes`; + // we have different metrics tracking WebSocket message size. + let db_router = counted_db_router.route("/subscribe", self.subscribe_get); let authed_root_router = axum::Router::new().route( "/", @@ -1661,7 +1670,8 @@ where let http_route_router = axum::Router::::new() .route("/:name_or_identity/route", self.http_route_root) .route("/:name_or_identity/route/", self.http_route_root_slash) - .route("/:name_or_identity/route/*path", self.http_route); + .route("/:name_or_identity/route/*path", self.http_route) + .route_layer(egress_metrics_middleware); axum::Router::new() .merge(authed_root_router) @@ -1670,6 +1680,68 @@ where } } +/// Middleware which counts response bytes in the metric `spacetime_http_response_size_bytes_total`. +/// +/// This middleware is intended to be supplied to all HTTP routes which apply to a particular database, +/// *except* the WebSocket connection route (confusingly named `subscribe`), whose egress is measured separately. +async fn http_response_egress_metrics_middleware( + State(worker_ctx): State, + Path(DatabaseParam { name_or_identity }): Path, + request: Request, + next: axum::middleware::Next, +) -> axum::response::Response +where + S: ControlStateDelegate + Clone + Send + Sync + 'static, +{ + let Ok(Ok(database_identity)) = name_or_identity.try_resolve(&worker_ctx).await else { + // The provided name doesn't map to an `Identity`. + // Run the route unchanged (as opposed to returning an error from the middleware) + // to preserve the error-handling behavior of the route. + return next.run(request).await; + }; + + if !matches!( + worker_ctx_find_database(&worker_ctx, &database_identity).await, + Ok(Some(_)) + ) { + // We have what appears to be an `Identity`, but it doesn't name any database. + // Don't create a metrics label for it, + // and run the route unchanged (as opposed to returning an error from the middleware) + // to preserve the error-handling behavior of the route. + return next.run(request).await; + } + + let response = next.run(request).await; + let (parts, body) = response.into_parts(); + + // Count the number of bytes used by the headers. + // For guest-defined routes bound to HTTP handlers, these may be arbitrarily large and are worth billing for; + // for built-in routes they will be small and it doesn't really matter one way or another whether we do or don't bill. + // N.b. headers installed by other middleware may or may not be counted here, + // depending on the order in which the middleware applies. + let header_bytes: usize = parts + .headers + .iter() + .map(|(name, value)| name.as_str().len() + value.as_bytes().len()) + .sum(); + + let counter = DB_METRICS + .http_response_size_bytes + .with_label_values(&database_identity); + counter.inc_by(header_bytes as u64); + + // `/logs?follow=true` can stream indefinitely. + // Counting frames as they are emitted preserves streaming behavior and avoids buffering the response. + let body = body.map_frame(move |frame| { + if let Some(data) = frame.data_ref() { + counter.inc_by(data.len() as u64); + } + frame + }); + + axum::response::Response::from_parts(parts, Body::new(body)) +} + #[cfg(test)] mod tests { use super::*; @@ -1690,10 +1762,15 @@ mod tests { use spacetimedb_client_api_messages::name::{ DomainName, InsertDomainResult, RegisterTldResult, SetDomainsResult, Tld, }; + use spacetimedb_lib::Hash; use spacetimedb_paths::server::ModuleLogsDir; use spacetimedb_paths::FromPathUnchecked; use spacetimedb_schema::auto_migrate::{MigrationPolicy, PrettyPrintStyle}; + use std::collections::HashMap; use tower::util::ServiceExt; + + static HTTP_EGRESS_METRIC_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + #[derive(Clone, Default)] struct DummyValidator; @@ -1710,8 +1787,14 @@ mod tests { } impl TokenSigner for DummyJwtProvider { - fn sign(&self, _claims: &T) -> Result { - Err(JwtError::from(JwtErrorKind::InvalidSignature)) + fn sign(&self, claims: &T) -> Result { + use base64::{engine::general_purpose, Engine}; + + let payload = serde_json::to_vec(claims).map_err(|_| JwtError::from(JwtErrorKind::InvalidSignature))?; + Ok(format!( + "test.{}.signature", + general_purpose::URL_SAFE_NO_PAD.encode(payload) + )) } } @@ -1736,6 +1819,8 @@ mod tests { jwt: DummyJwtProvider, client_actor_index: std::sync::Arc, module_logs_dir: ModuleLogsDir, + databases: std::sync::Arc>, + dns: std::sync::Arc>, } impl DummyState { @@ -1746,8 +1831,83 @@ mod tests { }, client_actor_index: std::sync::Arc::new(ClientActorIndex::new()), module_logs_dir: ModuleLogsDir::from_path_unchecked(std::env::temp_dir()), + databases: std::sync::Arc::new(HashMap::new()), + dns: std::sync::Arc::new(HashMap::new()), } } + + fn with_database(mut self, database_identity: Identity) -> Self { + let mut databases = HashMap::new(); + databases.insert(database_identity, test_database(database_identity)); + self.databases = std::sync::Arc::new(databases); + self + } + + fn with_dns(mut self, name: &str, database_identity: Identity) -> Self { + let mut dns = HashMap::new(); + dns.insert(name.to_owned(), database_identity); + self.dns = std::sync::Arc::new(dns); + self + } + } + + fn test_identity(byte: u8) -> Identity { + Identity::from_byte_array([byte; 32]) + } + + fn test_database(database_identity: Identity) -> Database { + Database { + id: u64::from(database_identity.to_byte_array()[0]), + database_identity, + owner_identity: test_identity(254), + host_type: HostType::Wasm, + initial_program: Hash::from_byte_array([0; 32]), + bootstrap_generation: 0, + } + } + + fn http_response_size_metric(database_identity: Identity) -> u64 { + DB_METRICS + .http_response_size_bytes + .with_label_values(&database_identity) + .get() + } + + fn collected_http_response_size_metric(database_identity: Identity) -> Option { + let db_label = database_identity.to_hex(); + for metric_family in prometheus::core::Collector::collect(&DB_METRICS.http_response_size_bytes) { + if metric_family.name() != "spacetime_http_response_size_bytes_total" { + continue; + } + + for metric in metric_family.get_metric() { + let has_db_label = metric + .get_label() + .iter() + .any(|label| label.name() == "db" && label.value() == db_label.as_str()); + if has_db_label { + return Some(metric.get_counter().value() as u64); + } + } + } + + None + } + + fn remove_http_response_size_metric(database_identity: Identity) { + let _ = DB_METRICS + .http_response_size_bytes + .remove_label_values(&database_identity); + } + + fn text_plain_header_bytes() -> u64 { + "content-type".len() as u64 + "text/plain; charset=utf-8".len() as u64 + } + + fn http_egress_metric_test_guard() -> std::sync::MutexGuard<'static, ()> { + HTTP_EGRESS_METRIC_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) } impl HasWebSocketOptions for DummyState { @@ -1817,8 +1977,8 @@ mod tests { async fn get_database_by_id(&self, _id: u64) -> anyhow::Result> { Ok(None) } - async fn get_database_by_identity(&self, _database_identity: &Identity) -> anyhow::Result> { - Ok(None) + async fn get_database_by_identity(&self, database_identity: &Identity) -> anyhow::Result> { + Ok(self.databases.get(database_identity).cloned()) } async fn get_databases(&self) -> anyhow::Result> { Ok(Vec::new()) @@ -1835,8 +1995,8 @@ mod tests { async fn get_energy_balance(&self, _identity: &Identity) -> anyhow::Result> { Ok(None) } - async fn lookup_database_identity(&self, _domain: &str) -> anyhow::Result> { - Ok(None) + async fn lookup_database_identity(&self, domain: &str) -> anyhow::Result> { + Ok(self.dns.get(domain).copied()) } async fn reverse_lookup(&self, _database_identity: &Identity) -> anyhow::Result> { Ok(Vec::new()) @@ -1966,4 +2126,197 @@ mod tests { // - `name_or_identity.resolve(worker_ctx)` -> `NameOrIdentity::resolve` assert_eq!(body, "`not-a-database` not found"); } + + #[tokio::test] + async fn http_response_egress_metric_counts_database_routes() { + let _guard = http_egress_metric_test_guard(); + let database_identity = test_identity(11); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { ([(http::header::CONTENT_TYPE, "text/plain")], "hello") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "hello"); + assert_eq!( + http_response_size_metric(database_identity), + "content-type".len() as u64 + "text/plain".len() as u64 + "hello".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_counts_user_http_route_headers_and_body() { + let _guard = http_egress_metric_test_guard(); + let database_identity = test_identity(12); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new() + .with_database(database_identity) + .with_dns("metric-test", database_identity); + let app = DatabaseRoutes:: { + http_route: axum::routing::any(|| async { ([("x-large-test-header", "abcdef")], "route-body") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri("/metric-test/route/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "route-body"); + assert_eq!( + http_response_size_metric(database_identity), + "x-large-test-header".len() as u64 + + "abcdef".len() as u64 + + text_plain_header_bytes() + + "route-body".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_skips_unresolved_and_nonexistent_databases() { + let _guard = http_egress_metric_test_guard(); + let resolved_but_missing_identity = test_identity(13); + let arbitrary_identity = test_identity(14); + remove_http_response_size_metric(resolved_but_missing_identity); + remove_http_response_size_metric(arbitrary_identity); + + let state = DummyState::new().with_dns("missing-db", resolved_but_missing_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { "not counted" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let _ = app + .clone() + .oneshot(Request::builder().uri("/unresolved-name").body(Body::empty()).unwrap()) + .await + .unwrap() + .into_body() + .collect() + .await + .unwrap(); + let _ = app + .clone() + .oneshot(Request::builder().uri("/missing-db").body(Body::empty()).unwrap()) + .await + .unwrap() + .into_body() + .collect() + .await + .unwrap(); + let _ = app + .oneshot( + Request::builder() + .uri(format!("/{arbitrary_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + .into_body() + .collect() + .await + .unwrap(); + + assert_eq!(collected_http_response_size_metric(resolved_but_missing_identity), None); + assert_eq!(collected_http_response_size_metric(arbitrary_identity), None); + + remove_http_response_size_metric(resolved_but_missing_identity); + remove_http_response_size_metric(arbitrary_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_does_not_count_subscribe() { + let _guard = http_egress_metric_test_guard(); + let database_identity = test_identity(15); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + subscribe_get: axum::routing::get(|| async { "subscribe response" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}/subscribe")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "subscribe response"); + assert_eq!(http_response_size_metric(database_identity), 0); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_counts_error_responses_for_existing_database() { + let _guard = http_egress_metric_test_guard(); + let database_identity = test_identity(16); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { (StatusCode::BAD_REQUEST, "bad request body") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "bad request body"); + assert_eq!( + http_response_size_metric(database_identity), + text_plain_header_bytes() + "bad request body".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } } diff --git a/crates/core/src/host/host_controller.rs b/crates/core/src/host/host_controller.rs index 60ffb1ba042..fbd717b51f9 100644 --- a/crates/core/src/host/host_controller.rs +++ b/crates/core/src/host/host_controller.rs @@ -1578,8 +1578,10 @@ pub async fn extract_schema(program_bytes: Box<[u8]>, host_type: HostType) -> an .await } -// Remove all gauges associated with a database. -// This is useful if a database is being deleted. +/// Removes metrics associated with a database when the database is deleted. +/// +/// Despite the historical function name, this cleans up per-database metric +/// series even when they are not literally `Gauge`s or `IntGauge`s. pub fn remove_database_gauges<'a, I>(db: &Identity, table_names: I) where I: IntoIterator, @@ -1609,4 +1611,5 @@ where V8HeapMetrics::remove_all_metric_label_values_for_database(db); let _ = WORKER_METRICS.v8_request_queue_length.remove_label_values(db); + let _ = DB_METRICS.http_response_size_bytes.remove_label_values(db); } diff --git a/crates/datastore/src/db_metrics/mod.rs b/crates/datastore/src/db_metrics/mod.rs index b8366034bac..2c7375362b0 100644 --- a/crates/datastore/src/db_metrics/mod.rs +++ b/crates/datastore/src/db_metrics/mod.rs @@ -225,6 +225,13 @@ An individual HTTP response's size in bytes is the sum of the sizes of the heade #[labels(db: Identity)] pub procedure_http_response_size_bytes: IntCounterVec, + #[name = spacetime_http_response_size_bytes_total] + #[help = "Total logical bytes sent in HTTP responses for routes scoped to a specific database. + +An individual HTTP response's size in bytes is the sum of the sizes of the header names, header values and body."] + #[labels(db: Identity)] + pub http_response_size_bytes: IntCounterVec, + #[name = spacetime_procedure_num_http_requests] #[help = "Number of HTTP requests performed by procedures running in databases. From 1d3b6508916a5ff54b0116ce1ae65ac848d6e40a Mon Sep 17 00:00:00 2001 From: Phoebe Goldman Date: Wed, 29 Jul 2026 14:49:24 -0400 Subject: [PATCH 2/7] Remove unnecessary mutex from tests --- crates/client-api/src/routes/database.rs | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index d008be214b7..0a52272153e 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -1769,8 +1769,6 @@ mod tests { use std::collections::HashMap; use tower::util::ServiceExt; - static HTTP_EGRESS_METRIC_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); - #[derive(Clone, Default)] struct DummyValidator; @@ -1904,12 +1902,6 @@ mod tests { "content-type".len() as u64 + "text/plain; charset=utf-8".len() as u64 } - fn http_egress_metric_test_guard() -> std::sync::MutexGuard<'static, ()> { - HTTP_EGRESS_METRIC_TEST_LOCK - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - } - impl HasWebSocketOptions for DummyState { fn websocket_options(&self) -> WebSocketOptions { WebSocketOptions::default() @@ -2129,7 +2121,6 @@ mod tests { #[tokio::test] async fn http_response_egress_metric_counts_database_routes() { - let _guard = http_egress_metric_test_guard(); let database_identity = test_identity(11); remove_http_response_size_metric(database_identity); @@ -2163,7 +2154,6 @@ mod tests { #[tokio::test] async fn http_response_egress_metric_counts_user_http_route_headers_and_body() { - let _guard = http_egress_metric_test_guard(); let database_identity = test_identity(12); remove_http_response_size_metric(database_identity); @@ -2202,7 +2192,6 @@ mod tests { #[tokio::test] async fn http_response_egress_metric_skips_unresolved_and_nonexistent_databases() { - let _guard = http_egress_metric_test_guard(); let resolved_but_missing_identity = test_identity(13); let arbitrary_identity = test_identity(14); remove_http_response_size_metric(resolved_but_missing_identity); @@ -2257,7 +2246,6 @@ mod tests { #[tokio::test] async fn http_response_egress_metric_does_not_count_subscribe() { - let _guard = http_egress_metric_test_guard(); let database_identity = test_identity(15); remove_http_response_size_metric(database_identity); @@ -2288,7 +2276,6 @@ mod tests { #[tokio::test] async fn http_response_egress_metric_counts_error_responses_for_existing_database() { - let _guard = http_egress_metric_test_guard(); let database_identity = test_identity(16); remove_http_response_size_metric(database_identity); From 6860f4c6cac5beb906785eb2c4c7fbdd3640dbeb Mon Sep 17 00:00:00 2001 From: Phoebe Goldman Date: Wed, 5 Aug 2026 13:48:35 -0400 Subject: [PATCH 3/7] Revise comment per Joshua's review --- crates/core/src/host/host_controller.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/core/src/host/host_controller.rs b/crates/core/src/host/host_controller.rs index fbd717b51f9..9fc3e012fa0 100644 --- a/crates/core/src/host/host_controller.rs +++ b/crates/core/src/host/host_controller.rs @@ -1578,7 +1578,10 @@ pub async fn extract_schema(program_bytes: Box<[u8]>, host_type: HostType) -> an .await } -/// Removes metrics associated with a database when the database is deleted. +/// Removes metrics associated with a database. +/// +/// This is called when a database's [`ModuleHost`] exits, +/// including (but not limited to) when a database is deleted. /// /// Despite the historical function name, this cleans up per-database metric /// series even when they are not literally `Gauge`s or `IntGauge`s. From 500d72f9f2bdfdcab81bc00b103882ce14f38924 Mon Sep 17 00:00:00 2001 From: Phoebe Goldman Date: Thu, 6 Aug 2026 16:53:54 -0400 Subject: [PATCH 4/7] Expand responsibility of middleware to also resolve databases (#5681) Based on and targeting #5611 . I thought this was a significant enough change to be worth reviewing separately. # Description of Changes The middleware responsible for the metric `spacetime_http_response_size_bytes_total` has to resolve a request's `:name_or_identity` from name to identity in order to compute the correct metrics label, and from identity to referent database in order to avoid allocating metrics labels for non-existent databases in response to ill-formed requests. Prior to this commit, the middleware did said resolution, then discarded the results, and the specific route handlers did the same resolution again. With this commit, the middleware is expanded to be fully responsible for database resolution. It attaches the resolved `Database` record to the request as an extension, and the route handlers read that extension rather than reading and resolving the `:name_or_identity` themselves. A small number of routes have behavior on non-existent databases other than returning 404. These 404s may occur either when a name is not bound to an identity, or when an identity is not bound to a database. (N.b. a name may be bound to an identity without that identity being bound to a database, which is a somewhat silly situation.) Rather than increasing the complexity of the middleware to cope with these behaviors or changing the behavior of these routes by applying the middleware to them, we opt to simply exclude these routes from the middleware, applying it only to routes which want the 404 response behavior. The specific routes exluded from the middleware are: - `PUT /database/:name_or_identity`, which creates a new database if the name is not already in use. - `DELETE /database/:name_or_identity`, which is a no-op on a non-existent database. - `/database/:name_or_identity/identity`, which only resolves a name to an identity, and does not check whether that identity refers to a database. - `GET /database/:name_or_identity/subscribe`, already excluded prior to this commit, whose response egress is tracked by a different metric than the one used by this middleware. All of the newly excluded routes are infrequent operations with small responses, and said responses are not user-controlled, so it's probably fine that they don't get counted by the egress metric. # API and ABI breaking changes N/a # Expected complexity level and risk # Testing - [x] New automated tests that affected routes don't re-resolve names unnecessarily. - [x] Other behavior should be covered by existing tests. --- crates/client-api/src/routes/database.rs | 422 ++++++++++++----------- crates/client-api/src/routes/mcp.rs | 38 +- crates/pg/src/pg_server.rs | 21 +- 3 files changed, 247 insertions(+), 234 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 0a52272153e..a690458d0b0 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -90,7 +90,6 @@ fn allow_creation(auth: &SpacetimeAuth) -> Result<(), ErrorResponse> { } #[derive(Deserialize)] pub struct CallParams { - name_or_identity: NameOrIdentity, reducer: String, } @@ -141,10 +140,8 @@ fn map_procedure_error(e: ProcedureCallError, procedure: &str) -> (StatusCode, S pub async fn call( State(worker_ctx): State, Extension(auth): Extension, - Path(CallParams { - name_or_identity, - reducer, - }): Path, + Extension(ResolvedDatabase(database)): Extension, + Path(CallParams { reducer }): Path, TypedHeader(content_type): TypedHeader, ByteStringBody(body): ByteStringBody, ) -> axum::response::Result { @@ -156,7 +153,8 @@ pub async fn call( let caller_auth: ConnectionAuthCtx = auth.into(); - let (module, Database { owner_identity, .. }) = find_module_and_database(&worker_ctx, name_or_identity).await?; + let owner_identity = database.owner_identity; + let module = find_database_module(&worker_ctx, &database).await?; let fut = async move |module: ModuleHost, caller_identity: Identity, connection_id: ConnectionId| { let result = match module @@ -215,42 +213,37 @@ pub async fn call( with_connection(module, caller_auth, caller_identity, fut).await } -#[derive(Deserialize)] -pub struct HttpRouteRootParams { - name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct HttpRouteParams { - name_or_identity: NameOrIdentity, path: String, } pub async fn handle_http_route_root( State(worker_ctx): State, - Path(HttpRouteRootParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, OriginalUri(original_uri): OriginalUri, request: Request, ) -> axum::response::Result { - handle_http_route_impl(worker_ctx, name_or_identity, "".to_string(), original_uri, request).await + handle_http_route_impl(worker_ctx, database, "".to_string(), original_uri, request).await } pub async fn handle_http_route_root_slash( State(worker_ctx): State, - Path(HttpRouteRootParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, OriginalUri(original_uri): OriginalUri, request: Request, ) -> axum::response::Result { - handle_http_route_impl(worker_ctx, name_or_identity, "/".to_string(), original_uri, request).await + handle_http_route_impl(worker_ctx, database, "/".to_string(), original_uri, request).await } pub async fn handle_http_route( State(worker_ctx): State, - Path(HttpRouteParams { name_or_identity, path }): Path, + Extension(ResolvedDatabase(database)): Extension, + Path(HttpRouteParams { path }): Path, OriginalUri(original_uri): OriginalUri, request: Request, ) -> axum::response::Result { - handle_http_route_impl(worker_ctx, name_or_identity, format!("/{path}"), original_uri, request).await + handle_http_route_impl(worker_ctx, database, format!("/{path}"), original_uri, request).await } /// Error response body for unknown user-defined HTTP route. @@ -258,7 +251,7 @@ const NO_SUCH_ROUTE: &str = "Database has not registered a handler for this rout async fn handle_http_route_impl( worker_ctx: S, - name_or_identity: NameOrIdentity, + database: Database, handler_path: String, original_uri: http::Uri, request: Request, @@ -266,7 +259,7 @@ async fn handle_http_route_impl( let (parts, body) = request.into_parts(); let st_method = http_method_to_st(&parts.method); - let (module, _database) = find_module_and_database(&worker_ctx, name_or_identity).await?; + let module = find_database_module(&worker_ctx, &database).await?; let module_def = &module.info().module_def; let Some((handler_id, _handler_def, _route_def)) = module_def.match_http_route(&st_method, &handler_path) else { @@ -461,31 +454,23 @@ pub(crate) fn client_disconnected_error_to_response(err: ReducerCallError) -> Er (StatusCode::INTERNAL_SERVER_ERROR, format!("{:#}", anyhow::anyhow!(err))).into() } -pub(crate) async fn find_leader_and_database( +pub(crate) async fn find_database_leader( worker_ctx: &S, - name_or_identity: NameOrIdentity, -) -> axum::response::Result<(Host, Database)> { - let db_identity = name_or_identity.resolve(worker_ctx).await?; - let database = worker_ctx_find_database(worker_ctx, &db_identity) - .await? - .ok_or_else(|| { - log::error!("Could not find database: {}", db_identity.to_hex()); - NO_SUCH_DATABASE - })?; - + database: &Database, +) -> axum::response::Result { let leader = worker_ctx.leader(database.id).await.map_err(Into::into)?; - Ok((leader, database)) + Ok(leader) } -pub(crate) async fn find_module_and_database( +pub(crate) async fn find_database_module( worker_ctx: &S, - name_or_identity: NameOrIdentity, -) -> axum::response::Result<(ModuleHost, Database)> { - let (leader, database) = find_leader_and_database(worker_ctx, name_or_identity).await?; + database: &Database, +) -> axum::response::Result { + let leader = find_database_leader(worker_ctx, database).await?; let module = leader.module().await.map_err(log_and_500)?; - Ok((module, database)) + Ok(module) } #[derive(Debug, derive_more::From)] @@ -502,10 +487,6 @@ fn procedure_outcome_response(return_val: AlgebraicValue) -> (StatusCode, axum:: ) } -#[derive(Deserialize)] -pub struct SchemaParams { - name_or_identity: NameOrIdentity, -} #[derive(Deserialize)] pub struct SchemaQueryParams { version: SchemaVersion, @@ -521,14 +502,14 @@ enum SchemaVersion { pub async fn schema( State(worker_ctx): State, - Path(SchemaParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Query(SchemaQueryParams { version }): Query, Extension(auth): Extension, ) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate, { - let (leader, _) = find_leader_and_database(&worker_ctx, name_or_identity).await?; + let leader = find_database_leader(&worker_ctx, &database).await?; // Wait for the module to finish loading rather than returning an immediate // 500 error. The database may still be initializing (replaying the log, // running init reducers, etc.). @@ -561,6 +542,9 @@ pub struct DatabaseParam { name_or_identity: NameOrIdentity, } +#[derive(Clone)] +pub struct ResolvedDatabase(pub Database); + #[derive(sats::Serialize)] struct DatabaseResponse { database_identity: Identity, @@ -581,26 +565,12 @@ impl From for DatabaseResponse { } pub async fn db_info( - State(worker_ctx): State, - Path(DatabaseParam { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, ) -> axum::response::Result { - log::trace!("Trying to resolve database identity: {name_or_identity:?}"); - let database_identity = name_or_identity.resolve(&worker_ctx).await?; - log::trace!("Resolved identity to: {database_identity:?}"); - let database = worker_ctx_find_database(&worker_ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; - log::trace!("Fetched database from the worker db for database identity: {database_identity:?}"); - let response = DatabaseResponse::from(database); Ok(axum::Json(sats::serde::SerdeWrapper(response))) } -#[derive(Deserialize)] -pub struct LogsParams { - name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct LogsQuery { num_lines: Option, @@ -610,7 +580,7 @@ pub struct LogsQuery { pub async fn logs( State(worker_ctx): State, - Path(LogsParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Query(LogsQuery { num_lines, follow }): Query, Extension(auth): Extension, ) -> axum::response::Result @@ -620,10 +590,7 @@ where // You should not be able to read the logs from a database that you do not own // so, unless you are the owner, this will fail. - let database_identity: Identity = name_or_identity.resolve(&worker_ctx).await?; - let database = worker_ctx_find_database(&worker_ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; + let database_identity = database.database_identity; worker_ctx .authorize_action(auth.claims.identity, database.database_identity, Action::ViewModuleLogs) @@ -687,11 +654,6 @@ pub(crate) async fn worker_ctx_find_database( .map_err(log_and_500) } -#[derive(Deserialize)] -pub struct SqlParams { - pub name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct SqlQueryParams { /// If `true`, return the query result only after its transaction offset @@ -742,7 +704,7 @@ where pub async fn sql_direct( worker_ctx: S, - SqlParams { name_or_identity }: SqlParams, + database: Database, SqlQueryParams { confirmed }: SqlQueryParams, caller_identity: Identity, caller_auth: ConnectionAuthCtx, @@ -751,7 +713,7 @@ pub async fn sql_direct( where S: NodeDelegate + ControlStateDelegate + Authorization + 'static, { - let (host, database) = find_leader_and_database(&worker_ctx, name_or_identity).await?; + let host = find_database_leader(&worker_ctx, &database).await?; let module = host.module().await.map_err(log_and_500)?; let fut = async move |_module: ModuleHost, caller_identity: Identity, _connection_id: ConnectionId| { @@ -773,7 +735,7 @@ where pub async fn sql( State(worker_ctx): State, - Path(name_or_identity): Path, + Extension(ResolvedDatabase(database)): Extension, Query(params): Query, Extension(auth): Extension, body: String, @@ -783,7 +745,7 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let json = sql_direct(worker_ctx, name_or_identity, params, caller_identity, caller_auth, body).await?; + let json = sql_direct(worker_ctx, database, params, caller_identity, caller_auth, body).await?; let total_duration = json.iter().fold(0, |acc, x| acc + x.total_duration_micros); @@ -798,11 +760,6 @@ pub struct DNSParams { name_or_identity: NameOrIdentity, } -#[derive(Deserialize)] -pub struct ReverseDNSParams { - name_or_identity: NameOrIdentity, -} - #[derive(Deserialize)] pub struct DNSQueryParams {} @@ -817,9 +774,9 @@ pub async fn get_identity( pub async fn get_names( State(ctx): State, - Path(ReverseDNSParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, ) -> axum::response::Result { - let database_identity = name_or_identity.resolve(&ctx).await?; + let database_identity = database.database_identity; let names = ctx .reverse_lookup(&database_identity) @@ -847,6 +804,7 @@ pub struct ResetDatabaseQueryParams { pub async fn reset( State(ctx): State, + Extension(ResolvedDatabase(database)): Extension, Path(ResetDatabaseParams { name_or_identity }): Path, Query(ResetDatabaseQueryParams { num_replicas, @@ -855,10 +813,7 @@ pub async fn reset( Extension(auth): Extension, program_bytes: Option, ) -> axum::response::Result> { - let database_identity = name_or_identity.resolve(&ctx).await?; - let database = worker_ctx_find_database(&ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database.database_identity, Action::ResetDatabase) .await?; @@ -965,12 +920,8 @@ pub async fn publish( .ok_or_else(|| bad_request("Clear database requires database name or identity".into()))?; let database_identity = name_or_identity.try_resolve(&ctx).await.map_err(log_and_500)?; if let Ok(identity) = database_identity { - let exists = ctx - .get_database_by_identity(&identity) - .await - .map_err(log_and_500)? - .is_some(); - if exists { + let database = ctx.get_database_by_identity(&identity).await.map_err(log_and_500)?; + if let Some(database) = database { if parent.is_some() { return Err(bad_request( "Setting the parent of an existing database is not supported".into(), @@ -979,6 +930,7 @@ pub async fn publish( return self::reset( State(ctx), + Extension(ResolvedDatabase(database)), Path(ResetDatabaseParams { name_or_identity: name_or_identity.clone(), }), @@ -1228,11 +1180,6 @@ fn bad_request(message: Cow<'static, str>) -> ErrorResponse { (StatusCode::BAD_REQUEST, message).into() } -#[derive(serde::Deserialize)] -pub struct PrePublishParams { - name_or_identity: NameOrIdentity, -} - #[derive(serde::Deserialize)] pub struct PrePublishQueryParams { #[serde(default)] @@ -1243,13 +1190,13 @@ pub struct PrePublishQueryParams { pub async fn pre_publish( State(ctx): State, - Path(PrePublishParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Query(PrePublishQueryParams { style, host_type }): Query, Extension(auth): Extension, program_bytes: Bytes, ) -> axum::response::Result> { // User should not be able to print migration plans for a database that they do not own - let database_identity = resolve_and_authenticate(&ctx, &name_or_identity, &auth).await?; + let database_identity = resolve_and_authenticate(&ctx, database, &auth).await?; let style = match style { PrettyPrintStyle::NoColor => AutoMigratePrettyPrintStyle::NoColor, PrettyPrintStyle::AnsiColor => AutoMigratePrettyPrintStyle::AnsiColor, @@ -1311,17 +1258,13 @@ pub async fn pre_publish .map(axum::Json) } -/// Resolves the [`NameOrIdentity`] to a database identity and checks if the -/// `auth` identity owns the database. +/// Checks if the `auth` identity owns the middleware-resolved database. async fn resolve_and_authenticate( ctx: &S, - name_or_identity: &NameOrIdentity, + database: Database, auth: &SpacetimeAuth, ) -> axum::response::Result { - let database_identity = name_or_identity.resolve(ctx).await?; - let database = worker_ctx_find_database(ctx, &database_identity) - .await? - .ok_or(NO_SUCH_DATABASE)?; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database.database_identity, Action::UpdateDatabase) .await?; @@ -1364,13 +1307,10 @@ pub async fn delete_database( pub async fn lock_database( State(ctx): State, - Path(DeleteDatabaseParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, ) -> axum::response::Result { - let database_identity = name_or_identity.resolve(&ctx).await?; - let Some(_database) = worker_ctx_find_database(&ctx, &database_identity).await? else { - return Err(StatusCode::NOT_FOUND.into()); - }; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database_identity, Action::DeleteDatabase) .await?; @@ -1384,13 +1324,10 @@ pub async fn lock_database( pub async fn unlock_database( State(ctx): State, - Path(DeleteDatabaseParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, ) -> axum::response::Result { - let database_identity = name_or_identity.resolve(&ctx).await?; - let Some(_database) = worker_ctx_find_database(&ctx, &database_identity).await? else { - return Err(StatusCode::NOT_FOUND.into()); - }; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database_identity, Action::DeleteDatabase) .await?; @@ -1402,19 +1339,14 @@ pub async fn unlock_database( Ok(()) } -#[derive(Deserialize)] -pub struct AddNameParams { - name_or_identity: NameOrIdentity, -} - pub async fn add_name( State(ctx): State, - Path(AddNameParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, name: String, ) -> axum::response::Result { let name = DatabaseName::try_from(name).map_err(|err| (StatusCode::BAD_REQUEST, err.to_string()))?; - let database_identity = name_or_identity.resolve(&ctx).await?; + let database_identity = database.database_identity; let response = ctx .create_dns_record(&auth.claims.identity, &name.into(), &database_identity) @@ -1432,14 +1364,9 @@ pub async fn add_name( Ok((code, axum::Json(response))) } -#[derive(Deserialize)] -pub struct SetNamesParams { - name_or_identity: NameOrIdentity, -} - pub async fn set_names( State(ctx): State, - Path(SetNamesParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, names: axum::Json>, ) -> axum::response::Result { @@ -1450,18 +1377,7 @@ pub async fn set_names( .collect::, _>>() .map_err(|(input, e)| (StatusCode::BAD_REQUEST, format!("Error parsing `{input}`: {e}")))?; - let database_identity = name_or_identity.resolve(&ctx).await?; - - let database = ctx - .get_database_by_identity(&database_identity) - .await - .map_err(log_and_500)?; - let Some(database) = database else { - return Ok(( - StatusCode::NOT_FOUND, - axum::Json(name::SetDomainsResult::DatabaseNotFound), - )); - }; + let database_identity = database.database_identity; ctx.authorize_action(auth.claims.identity, database.database_identity, Action::RenameDatabase) .await @@ -1509,30 +1425,15 @@ pub async fn set_names( Ok((status, axum::Json(response))) } -#[derive(serde::Deserialize)] -pub struct TimestampParams { - name_or_identity: NameOrIdentity, -} - /// Returns the database's view of the current time, /// as a SATS-JSON encoded [`Timestamp`]. /// -/// Takes a particular database's [`NameOrIdentity`] as an argument +/// Takes a particular database as an argument /// because in a clusterized SpacetimeDB-cloud deployment, /// this request will be routed to the node running the requested database. -async fn get_timestamp( - State(worker_ctx): State, - Path(TimestampParams { name_or_identity }): Path, +async fn get_timestamp( + Extension(ResolvedDatabase(_database)): Extension, ) -> axum::response::Result { - let db_identity = name_or_identity.resolve(&worker_ctx).await?; - - let _database = worker_ctx_find_database(&worker_ctx, &db_identity) - .await? - .ok_or_else(|| { - log::error!("Could not find database: {}", db_identity.to_hex()); - NO_SUCH_DATABASE - })?; - Ok(axum::Json(sats::serde::SerdeWrapper(Timestamp::now())).into_response()) } @@ -1607,7 +1508,7 @@ where mcp_post: post(crate::routes::mcp::mcp::), pre_publish: post(pre_publish::), db_reset: put(reset::), - timestamp_get: get(get_timestamp::), + timestamp_get: get(get_timestamp), lock_post: post(lock_database::), unlock_post: post(unlock_database::), http_route_root: any(handle_http_route_root::), @@ -1622,17 +1523,16 @@ where S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, { pub fn into_router(self, ctx: S) -> axum::Router { - let egress_metrics_middleware = - axum::middleware::from_fn_with_state(ctx.clone(), http_response_egress_metrics_middleware::); + let resolving_egress_metrics_middleware = axum::middleware::from_fn_with_state( + ctx.clone(), + resolve_database_name_and_count_response_egress_middleware::, + ); let counted_db_router = axum::Router::::new() - .route("/", self.db_put) .route("/", self.db_get) - .route("/", self.db_delete) .route("/names", self.names_get) .route("/names", self.names_post) .route("/names", self.names_put) - .route("/identity", self.identity_get) .route("/call/:reducer", self.call_reducer_procedure_post) .route("/schema", self.schema_get) .route("/logs", self.logs_get) @@ -1643,12 +1543,24 @@ where .route("/reset", self.db_reset) .route("/lock", self.lock_post) .route("/unlock", self.unlock_post) - .route_layer(egress_metrics_middleware.clone()); + .route_layer(resolving_egress_metrics_middleware.clone()); + + // Publishing can create a database for a new name, so it bypasses existing-database resolution. + // Publish operations are infrequent and the responses are small, so we don't mind that we don't measure them. + let db_router = counted_db_router.route("/", self.db_put); + + // These routes have different behavior on a non-existent database than a 404, + // and so resolve the database themselves rather than having the middleware do so. + // Like publish, these operations are infrequent and the responses are small, + // so we don't mind that we don't measure them. + let db_router = db_router + .route("/", self.db_delete) + .route("/identity", self.identity_get); - // Add the subscribe route after `egress_metrics_middleware` + // Add the subscribe route after `resolving_egress_metrics_middleware` // so that its egress bytes don't get counted into `http_response_size_bytes`; // we have different metrics tracking WebSocket message size. - let db_router = counted_db_router.route("/subscribe", self.subscribe_get); + let db_router = db_router.route("/subscribe", self.subscribe_get); let authed_root_router = axum::Router::new().route( "/", @@ -1671,7 +1583,7 @@ where .route("/:name_or_identity/route", self.http_route_root) .route("/:name_or_identity/route/", self.http_route_root_slash) .route("/:name_or_identity/route/*path", self.http_route) - .route_layer(egress_metrics_middleware); + .route_layer(resolving_egress_metrics_middleware); axum::Router::new() .merge(authed_root_router) @@ -1680,36 +1592,31 @@ where } } -/// Middleware which counts response bytes in the metric `spacetime_http_response_size_bytes_total`. +/// Resolves an existing database, attaches it as [`ResolvedDatabase`], +/// and counts response bytes in the metric `spacetime_http_response_size_bytes_total`. /// -/// This middleware is intended to be supplied to all HTTP routes which apply to a particular database, -/// *except* the WebSocket connection route (confusingly named `subscribe`), whose egress is measured separately. -async fn http_response_egress_metrics_middleware( +/// This middleware returns established name and database `404`s before the handler runs. +/// It is intended for HTTP routes that require an existing database, +/// except WebSocket `subscribe`, whose egress is measured separately. +async fn resolve_database_name_and_count_response_egress_middleware( State(worker_ctx): State, Path(DatabaseParam { name_or_identity }): Path, - request: Request, + mut request: Request, next: axum::middleware::Next, ) -> axum::response::Response where S: ControlStateDelegate + Clone + Send + Sync + 'static, { - let Ok(Ok(database_identity)) = name_or_identity.try_resolve(&worker_ctx).await else { - // The provided name doesn't map to an `Identity`. - // Run the route unchanged (as opposed to returning an error from the middleware) - // to preserve the error-handling behavior of the route. - return next.run(request).await; + let database_identity = match name_or_identity.resolve(&worker_ctx).await { + Ok(database_identity) => database_identity, + Err(response) => return Err::<(), ErrorResponse>(response).into_response(), }; - - if !matches!( - worker_ctx_find_database(&worker_ctx, &database_identity).await, - Ok(Some(_)) - ) { - // We have what appears to be an `Identity`, but it doesn't name any database. - // Don't create a metrics label for it, - // and run the route unchanged (as opposed to returning an error from the middleware) - // to preserve the error-handling behavior of the route. - return next.run(request).await; - } + let database = match worker_ctx_find_database(&worker_ctx, &database_identity).await { + Ok(Some(database)) => database, + Ok(None) => return NO_SUCH_DATABASE.into_response(), + Err(response) => return Err::<(), ErrorResponse>(response).into_response(), + }; + request.extensions_mut().insert(ResolvedDatabase(database.clone())); let response = next.run(request).await; let (parts, body) = response.into_parts(); @@ -1727,7 +1634,7 @@ where let counter = DB_METRICS .http_response_size_bytes - .with_label_values(&database_identity); + .with_label_values(&database.database_identity); counter.inc_by(header_bytes as u64); // `/logs?follow=true` can stream indefinitely. @@ -1767,6 +1674,7 @@ mod tests { use spacetimedb_paths::FromPathUnchecked; use spacetimedb_schema::auto_migrate::{MigrationPolicy, PrettyPrintStyle}; use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; use tower::util::ServiceExt; #[derive(Clone, Default)] @@ -1819,6 +1727,7 @@ mod tests { module_logs_dir: ModuleLogsDir, databases: std::sync::Arc>, dns: std::sync::Arc>, + dns_lookups: std::sync::Arc, } impl DummyState { @@ -1831,6 +1740,7 @@ mod tests { module_logs_dir: ModuleLogsDir::from_path_unchecked(std::env::temp_dir()), databases: std::sync::Arc::new(HashMap::new()), dns: std::sync::Arc::new(HashMap::new()), + dns_lookups: std::sync::Arc::new(AtomicUsize::new(0)), } } @@ -1847,6 +1757,10 @@ mod tests { self.dns = std::sync::Arc::new(dns); self } + + fn dns_lookups(&self) -> usize { + self.dns_lookups.load(Ordering::Relaxed) + } } fn test_identity(byte: u8) -> Identity { @@ -1988,6 +1902,7 @@ mod tests { Ok(None) } async fn lookup_database_identity(&self, domain: &str) -> anyhow::Result> { + self.dns_lookups.fetch_add(1, Ordering::Relaxed); Ok(self.dns.get(domain).copied()) } async fn reverse_lookup(&self, _database_identity: &Identity) -> anyhow::Result> { @@ -2113,9 +2028,8 @@ mod tests { assert_eq!(response.status(), StatusCode::NOT_FOUND); let body = response.into_body().collect().await.unwrap().to_bytes(); // We'll get this error message out of the stack: - // - `find_module_and_database` - // - `find_leader_and_database` - // - `name_or_identity.resolve(worker_ctx)` -> `NameOrIdentity::resolve` + // The database-resolution middleware returns this established error + // before the unauthenticated HTTP route handler is invoked. assert_eq!(body, "`not-a-database` not found"); } @@ -2191,7 +2105,7 @@ mod tests { } #[tokio::test] - async fn http_response_egress_metric_skips_unresolved_and_nonexistent_databases() { + async fn resolving_middleware_returns_not_found_without_creating_metrics_labels() { let resolved_but_missing_identity = test_identity(13); let arbitrary_identity = test_identity(14); remove_http_response_size_metric(resolved_but_missing_identity); @@ -2205,25 +2119,29 @@ mod tests { .into_router(state.clone()) .with_state(state); - let _ = app + let response = app .clone() .oneshot(Request::builder().uri("/unresolved-name").body(Body::empty()).unwrap()) .await - .unwrap() - .into_body() - .collect() - .await .unwrap(); - let _ = app + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "`unresolved-name` not found" + ); + + let response = app .clone() .oneshot(Request::builder().uri("/missing-db").body(Body::empty()).unwrap()) .await - .unwrap() - .into_body() - .collect() - .await .unwrap(); - let _ = app + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "No such database." + ); + + let response = app .oneshot( Request::builder() .uri(format!("/{arbitrary_identity}")) @@ -2231,11 +2149,12 @@ mod tests { .unwrap(), ) .await - .unwrap() - .into_body() - .collect() - .await .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "No such database." + ); assert_eq!(collected_http_response_size_metric(resolved_but_missing_identity), None); assert_eq!(collected_http_response_size_metric(arbitrary_identity), None); @@ -2244,6 +2163,95 @@ mod tests { remove_http_response_size_metric(arbitrary_identity); } + #[tokio::test] + async fn resolving_middleware_attaches_database_and_resolves_a_name_once() { + let database_identity = test_identity(17); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new() + .with_database(database_identity) + .with_dns("named-database", database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get( + |Extension(ResolvedDatabase(database)): Extension| async move { + database.database_identity.to_string() + }, + ), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state.clone()); + + let response = app + .oneshot(Request::builder().uri("/named-database").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + database_identity.to_string() + ); + assert_eq!(state.dns_lookups(), 1); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn db_info_resolves_a_database_name_once() { + let database_identity = test_identity(19); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new() + .with_database(database_identity) + .with_dns("db-info", database_identity); + let app = DatabaseRoutes::::default() + .into_router(state.clone()) + .with_state(state.clone()); + + let response = app + .oneshot(Request::builder().uri("/db-info").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(state.dns_lookups(), 1); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn publish_delete_and_identity_bypass_resolving_middleware() { + let missing_identity = test_identity(18); + remove_http_response_size_metric(missing_identity); + + let state = DummyState::new(); + let app = DatabaseRoutes:: { + db_put: axum::routing::put(|| async { "publish" }), + db_delete: axum::routing::delete(|| async { "delete" }), + identity_get: axum::routing::get(|| async { "identity" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let delete_uri = format!("/{missing_identity}"); + let identity_uri = format!("/{missing_identity}/identity"); + for (method, uri, expected) in [ + (http::Method::PUT, "/unregistered-name", "publish"), + (http::Method::DELETE, delete_uri.as_str(), "delete"), + (http::Method::GET, identity_uri.as_str(), "identity"), + ] { + let response = app + .clone() + .oneshot(Request::builder().method(method).uri(uri).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.into_body().collect().await.unwrap().to_bytes(), expected); + } + + assert_eq!(collected_http_response_size_metric(missing_identity), None); + } + #[tokio::test] async fn http_response_egress_metric_does_not_count_subscribe() { let database_identity = test_identity(15); diff --git a/crates/client-api/src/routes/mcp.rs b/crates/client-api/src/routes/mcp.rs index 9f60644735b..fab0aa610af 100644 --- a/crates/client-api/src/routes/mcp.rs +++ b/crates/client-api/src/routes/mcp.rs @@ -1,23 +1,22 @@ use std::time::Duration; -use axum::extract::{Path, State}; +use axum::extract::State; use axum::response::{ErrorResponse, IntoResponse, Response}; use axum::{Extension, Json}; use http::StatusCode; -use serde::Deserialize; use serde_json::{json, Value}; use spacetimedb::auth::identity::ConnectionAuthCtx; use spacetimedb::host::{FunctionArgs, ReducerOutcome}; +use spacetimedb::messages::control_db::Database; use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9; use spacetimedb_lib::sats; use super::database::{ - client_connected_error_to_response, client_disconnected_error_to_response, find_leader_and_database, - find_module_and_database, map_reducer_error, sql_direct, SqlParams, SqlQueryParams, + client_connected_error_to_response, client_disconnected_error_to_response, find_database_leader, + find_database_module, map_reducer_error, sql_direct, ResolvedDatabase, SqlQueryParams, }; use crate::auth::SpacetimeAuth; use crate::routes::subscribe::generate_random_connection_id; -use crate::util::NameOrIdentity; use crate::{log_and_500, Authorization, ControlStateDelegate, NodeDelegate}; const PROTOCOL_VERSION: &str = "2025-06-18"; @@ -36,15 +35,10 @@ const MAX_ERROR_BODY_BYTES: usize = 64 * 1024; type RpcError = (i64, String); -#[derive(Deserialize)] -pub struct McpParams { - name_or_identity: NameOrIdentity, -} - /// handle MCP JSON-RPC request pub async fn mcp( State(ctx): State, - Path(McpParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, Json(request): Json, ) -> axum::response::Result @@ -65,7 +59,7 @@ where // protocol ping, distinct from the ping tool "ping" => jsonrpc_result(&id, json!({})), "tools/list" => jsonrpc_result(&id, tools_list()), - "tools/call" => match tools_call(&ctx, name_or_identity, auth, request.get("params")).await { + "tools/call" => match tools_call(&ctx, database, auth, request.get("params")).await { Ok(result) => jsonrpc_result(&id, result), Err((code, message)) => jsonrpc_error(&id, code, message), }, @@ -143,7 +137,7 @@ fn tools_list() -> Value { async fn tools_call( ctx: &S, - name_or_identity: NameOrIdentity, + database: Database, auth: SpacetimeAuth, params: Option<&Value>, ) -> Result @@ -163,20 +157,20 @@ where Some(message) => format!("pong: {message}"), None => "pong".to_owned(), }), - "get_schema" => tool_get_schema(ctx, name_or_identity).await, + "get_schema" => tool_get_schema(ctx, &database).await, "sql" => { let Some(sql) = arguments.and_then(|a| a.get("sql")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "sql argument must be a string".to_owned())); }; let confirmed = arguments.and_then(|a| a.get("confirmed")).and_then(Value::as_bool); - tool_sql(ctx, name_or_identity, auth, sql.to_owned(), confirmed).await + tool_sql(ctx, database, auth, sql.to_owned(), confirmed).await } "call" => { let Some(reducer) = arguments.and_then(|a| a.get("reducer")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "reducer argument must be a string".to_owned())); }; let args_json = reducer_args_json(arguments)?; - tool_call_reducer(ctx, name_or_identity, auth, reducer.to_owned(), args_json).await + tool_call_reducer(ctx, &database, auth, reducer.to_owned(), args_json).await } other => return Err((INVALID_PARAMS, format!("unknown tool: {other}"))), }; @@ -207,11 +201,11 @@ async fn execution_error_to_tool_result(err: ErrorResponse) -> Value { json!({ "content": [ { "type": "text", "text": text } ], "isError": true }) } -async fn tool_get_schema(ctx: &S, name_or_identity: NameOrIdentity) -> axum::response::Result +async fn tool_get_schema(ctx: &S, database: &Database) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate, { - let (leader, _) = find_leader_and_database(ctx, name_or_identity).await?; + let leader = find_database_leader(ctx, database).await?; let module = leader.wait_for_module(MODULE_WAIT_TIMEOUT).await.map_err(log_and_500)?; let raw = RawModuleDefV9::from(module.info.module_def.as_ref().clone()); let json = serde_json::to_string(&sats::serde::SerdeWrapper(raw)).map_err(log_and_500)?; @@ -220,7 +214,7 @@ where async fn tool_sql( ctx: &S, - name_or_identity: NameOrIdentity, + database: Database, auth: SpacetimeAuth, sql: String, confirmed: Option, @@ -232,7 +226,7 @@ where let caller_auth: ConnectionAuthCtx = auth.into(); let rows = sql_direct( ctx.clone(), - SqlParams { name_or_identity }, + database, SqlQueryParams { confirmed }, caller_identity, caller_auth, @@ -245,7 +239,7 @@ where async fn tool_call_reducer( ctx: &S, - name_or_identity: NameOrIdentity, + database: &Database, auth: SpacetimeAuth, reducer: String, args_json: String, @@ -255,7 +249,7 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let (module, _) = find_module_and_database(ctx, name_or_identity).await?; + let module = find_database_module(ctx, database).await?; let connection_id = generate_random_connection_id(); module diff --git a/crates/pg/src/pg_server.rs b/crates/pg/src/pg_server.rs index 15170a35866..1b3109d506b 100644 --- a/crates/pg/src/pg_server.rs +++ b/crates/pg/src/pg_server.rs @@ -25,7 +25,7 @@ use pgwire::tokio::process_socket; use spacetimedb_auth::identity::ConnectionAuthCtx; use spacetimedb_client_api::auth::validate_token; use spacetimedb_client_api::routes::database; -use spacetimedb_client_api::routes::database::{SqlParams, SqlQueryParams}; +use spacetimedb_client_api::routes::database::SqlQueryParams; use spacetimedb_client_api::{Authorization, ControlStateReadAccess, ControlStateWriteAccess, NodeDelegate}; use spacetimedb_client_api_messages::http::SqlStmtResult; use spacetimedb_client_api_messages::name::DatabaseName; @@ -154,14 +154,25 @@ where { async fn exe_sql(&self, query: String) -> PgWireResult> { let params = self.cached.lock().await.clone().unwrap(); - let db = SqlParams { - name_or_identity: database::NameOrIdentity::Name(DatabaseName(params.database.clone())), - }; + let name_or_identity = database::NameOrIdentity::Name(DatabaseName(params.database.clone())); + let database_identity = response(name_or_identity.resolve(&self.ctx).await, ¶ms.database).await?; + let database = response( + self.ctx + .get_database_by_identity(&database_identity) + .await + .map_err(|err| { + log::warn!("PG: unable to load database {database_identity}: {err:#}"); + (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error").into() + }) + .and_then(|database| database.ok_or_else(|| database::NO_SUCH_DATABASE.into())), + ¶ms.database, + ) + .await?; let sql = match response( database::sql_direct( self.ctx.clone(), - db, + database, SqlQueryParams { confirmed: Some(true) }, params.caller_identity, params.caller_auth.clone(), From aee413d09668763f5827bb1fbf661a99cd886bdf Mon Sep 17 00:00:00 2001 From: Phoebe Goldman Date: Fri, 7 Aug 2026 10:53:13 -0400 Subject: [PATCH 5/7] Move `POST mcp` out of the middleware See comments in mcp.rs --- crates/client-api/src/routes/database.rs | 18 ++++++++-- crates/client-api/src/routes/mcp.rs | 44 ++++++++++++++++-------- 2 files changed, 46 insertions(+), 16 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index a690458d0b0..c6ee663fd8e 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -654,6 +654,18 @@ pub(crate) async fn worker_ctx_find_database( .map_err(log_and_500) } +pub(crate) async fn find_database_or_404( + worker_ctx: &(impl ControlStateDelegate + ?Sized), + name_or_identity: NameOrIdentity, +) -> axum::response::Result { + let identity = name_or_identity.resolve(worker_ctx).await?; + let database = worker_ctx_find_database(worker_ctx, &identity).await?.ok_or_else(|| { + log::debug!("Identity {identity} in HTTP request does not refer to a database"); + NO_SUCH_DATABASE + })?; + Ok(database) +} + #[derive(Deserialize)] pub struct SqlQueryParams { /// If `true`, return the query result only after its transaction offset @@ -1537,7 +1549,6 @@ where .route("/schema", self.schema_get) .route("/logs", self.logs_get) .route("/sql", self.sql_post) - .route("/mcp", self.mcp_post) .route("/unstable/timestamp", self.timestamp_get) .route("/pre_publish", self.pre_publish) .route("/reset", self.db_reset) @@ -1555,7 +1566,10 @@ where // so we don't mind that we don't measure them. let db_router = db_router .route("/", self.db_delete) - .route("/identity", self.identity_get); + .route("/identity", self.identity_get) + // I (pgoldman 2026-08-07) am actually somewhat concerned that we do care about measuring egress for MCP requests, + // but the MCP handler's name resolution and error handling are significantly incompatible with the middleware. + .route("/mcp", self.mcp_post); // Add the subscribe route after `resolving_egress_metrics_middleware` // so that its egress bytes don't get counted into `http_response_size_bytes`; diff --git a/crates/client-api/src/routes/mcp.rs b/crates/client-api/src/routes/mcp.rs index fab0aa610af..73cef659123 100644 --- a/crates/client-api/src/routes/mcp.rs +++ b/crates/client-api/src/routes/mcp.rs @@ -1,22 +1,23 @@ use std::time::Duration; -use axum::extract::State; +use axum::extract::{Path, State}; use axum::response::{ErrorResponse, IntoResponse, Response}; use axum::{Extension, Json}; use http::StatusCode; +use serde::Deserialize; use serde_json::{json, Value}; use spacetimedb::auth::identity::ConnectionAuthCtx; use spacetimedb::host::{FunctionArgs, ReducerOutcome}; -use spacetimedb::messages::control_db::Database; use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9; use spacetimedb_lib::sats; use super::database::{ client_connected_error_to_response, client_disconnected_error_to_response, find_database_leader, - find_database_module, map_reducer_error, sql_direct, ResolvedDatabase, SqlQueryParams, + find_database_module, find_database_or_404, map_reducer_error, sql_direct, SqlQueryParams, }; use crate::auth::SpacetimeAuth; use crate::routes::subscribe::generate_random_connection_id; +use crate::util::NameOrIdentity; use crate::{log_and_500, Authorization, ControlStateDelegate, NodeDelegate}; const PROTOCOL_VERSION: &str = "2025-06-18"; @@ -35,10 +36,22 @@ const MAX_ERROR_BODY_BYTES: usize = 64 * 1024; type RpcError = (i64, String); +#[derive(Deserialize)] +pub struct McpParams { + name_or_identity: NameOrIdentity, +} + /// handle MCP JSON-RPC request +// +// Due to different name resolution and error handling behavior in different branches, +// this route handler does not use [`super::database::resolve_database_name_and_count_response_egress_middleware`]. +// This is unfortunate, as we probably would like to count egress bytes from MCP calls, +// but I (pgoldman 2026-08-07) do not have the wherewithal +// to significantly rewrite this file in order to make it compatible with the middleware, +// and do not know which of its error-handling behaviors are safe to change. pub async fn mcp( State(ctx): State, - Extension(ResolvedDatabase(database)): Extension, + Path(McpParams { name_or_identity }): Path, Extension(auth): Extension, Json(request): Json, ) -> axum::response::Result @@ -59,7 +72,7 @@ where // protocol ping, distinct from the ping tool "ping" => jsonrpc_result(&id, json!({})), "tools/list" => jsonrpc_result(&id, tools_list()), - "tools/call" => match tools_call(&ctx, database, auth, request.get("params")).await { + "tools/call" => match tools_call(&ctx, name_or_identity, auth, request.get("params")).await { Ok(result) => jsonrpc_result(&id, result), Err((code, message)) => jsonrpc_error(&id, code, message), }, @@ -137,7 +150,7 @@ fn tools_list() -> Value { async fn tools_call( ctx: &S, - database: Database, + name_or_identity: NameOrIdentity, auth: SpacetimeAuth, params: Option<&Value>, ) -> Result @@ -157,20 +170,20 @@ where Some(message) => format!("pong: {message}"), None => "pong".to_owned(), }), - "get_schema" => tool_get_schema(ctx, &database).await, + "get_schema" => tool_get_schema(ctx, name_or_identity).await, "sql" => { let Some(sql) = arguments.and_then(|a| a.get("sql")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "sql argument must be a string".to_owned())); }; let confirmed = arguments.and_then(|a| a.get("confirmed")).and_then(Value::as_bool); - tool_sql(ctx, database, auth, sql.to_owned(), confirmed).await + tool_sql(ctx, name_or_identity, auth, sql.to_owned(), confirmed).await } "call" => { let Some(reducer) = arguments.and_then(|a| a.get("reducer")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "reducer argument must be a string".to_owned())); }; let args_json = reducer_args_json(arguments)?; - tool_call_reducer(ctx, &database, auth, reducer.to_owned(), args_json).await + tool_call_reducer(ctx, name_or_identity, auth, reducer.to_owned(), args_json).await } other => return Err((INVALID_PARAMS, format!("unknown tool: {other}"))), }; @@ -201,11 +214,12 @@ async fn execution_error_to_tool_result(err: ErrorResponse) -> Value { json!({ "content": [ { "type": "text", "text": text } ], "isError": true }) } -async fn tool_get_schema(ctx: &S, database: &Database) -> axum::response::Result +async fn tool_get_schema(ctx: &S, name_or_identity: NameOrIdentity) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate, { - let leader = find_database_leader(ctx, database).await?; + let database = find_database_or_404(ctx, name_or_identity).await?; + let leader = find_database_leader(ctx, &database).await?; let module = leader.wait_for_module(MODULE_WAIT_TIMEOUT).await.map_err(log_and_500)?; let raw = RawModuleDefV9::from(module.info.module_def.as_ref().clone()); let json = serde_json::to_string(&sats::serde::SerdeWrapper(raw)).map_err(log_and_500)?; @@ -214,7 +228,7 @@ where async fn tool_sql( ctx: &S, - database: Database, + name_or_identity: NameOrIdentity, auth: SpacetimeAuth, sql: String, confirmed: Option, @@ -224,6 +238,7 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); + let database = find_database_or_404(ctx, name_or_identity).await?; let rows = sql_direct( ctx.clone(), database, @@ -239,7 +254,7 @@ where async fn tool_call_reducer( ctx: &S, - database: &Database, + name_or_identity: NameOrIdentity, auth: SpacetimeAuth, reducer: String, args_json: String, @@ -249,7 +264,8 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let module = find_database_module(ctx, database).await?; + let database = find_database_or_404(ctx, name_or_identity).await?; + let module = find_database_module(ctx, &database).await?; let connection_id = generate_random_connection_id(); module From bcc3d79600beca58baf4b579611af6f2ff3fba51 Mon Sep 17 00:00:00 2001 From: Phoebe Goldman Date: Fri, 7 Aug 2026 12:31:50 -0400 Subject: [PATCH 6/7] Tidy pre-publish auth check and middleware return type --- crates/client-api/src/routes/database.rs | 34 ++++++------------------ 1 file changed, 8 insertions(+), 26 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 738c693d162..8f74ad2688d 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -1207,8 +1207,12 @@ pub async fn pre_publish Extension(auth): Extension, program_bytes: Bytes, ) -> axum::response::Result> { + let database_identity = database.database_identity; + // User should not be able to print migration plans for a database that they do not own - let database_identity = resolve_and_authenticate(&ctx, database, &auth).await?; + ctx.authorize_action(auth.claims.identity, database_identity, Action::UpdateDatabase) + .await?; + let style = match style { PrettyPrintStyle::NoColor => AutoMigratePrettyPrintStyle::NoColor, PrettyPrintStyle::AnsiColor => AutoMigratePrettyPrintStyle::AnsiColor, @@ -1270,20 +1274,6 @@ pub async fn pre_publish .map(axum::Json) } -/// Checks if the `auth` identity owns the middleware-resolved database. -async fn resolve_and_authenticate( - ctx: &S, - database: Database, - auth: &SpacetimeAuth, -) -> axum::response::Result { - let database_identity = database.database_identity; - - ctx.authorize_action(auth.claims.identity, database.database_identity, Action::UpdateDatabase) - .await?; - - Ok(database_identity) -} - #[derive(Deserialize)] pub struct DeleteDatabaseParams { pub name_or_identity: NameOrIdentity, @@ -1617,19 +1607,11 @@ async fn resolve_database_name_and_count_response_egress_middleware( Path(DatabaseParam { name_or_identity }): Path, mut request: Request, next: axum::middleware::Next, -) -> axum::response::Response +) -> axum::response::Result where S: ControlStateDelegate + Clone + Send + Sync + 'static, { - let database_identity = match name_or_identity.resolve(&worker_ctx).await { - Ok(database_identity) => database_identity, - Err(response) => return Err::<(), ErrorResponse>(response).into_response(), - }; - let database = match worker_ctx_find_database(&worker_ctx, &database_identity).await { - Ok(Some(database)) => database, - Ok(None) => return NO_SUCH_DATABASE.into_response(), - Err(response) => return Err::<(), ErrorResponse>(response).into_response(), - }; + let database = find_database_or_404(&worker_ctx, name_or_identity).await?; request.extensions_mut().insert(ResolvedDatabase(database.clone())); let response = next.run(request).await; @@ -1660,7 +1642,7 @@ where frame }); - axum::response::Response::from_parts(parts, Body::new(body)) + Ok(axum::response::Response::from_parts(parts, Body::new(body))) } #[cfg(test)] From e0f8f127d347f08b665e57d5e5c47f9075b8662e Mon Sep 17 00:00:00 2001 From: Phoebe Goldman Date: Fri, 7 Aug 2026 13:16:50 -0400 Subject: [PATCH 7/7] Fully drain body before responding 404, to avoid broken pipe error on clients --- crates/client-api/src/routes/database.rs | 41 ++++++++++++++++++++++-- 1 file changed, 38 insertions(+), 3 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 8f74ad2688d..f45eee60cd1 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -1611,7 +1611,16 @@ async fn resolve_database_name_and_count_response_egress_middleware( where S: ControlStateDelegate + Clone + Send + Sync + 'static, { - let database = find_database_or_404(&worker_ctx, name_or_identity).await?; + let database = match find_database_or_404(&worker_ctx, name_or_identity).await { + Ok(database) => database, + Err(err) => { + // Fully drain the request before responding + // so clients uploading a body receive the HTTP error + // instead of a broken pipe from a connection closed mid-upload. + while let Some(Ok(_)) = request.body_mut().frame().await {} + return Err(err); + } + }; request.extensions_mut().insert(ResolvedDatabase(database.clone())); let response = next.run(request).await; @@ -1654,7 +1663,7 @@ mod tests { Action, Authorization, ControlStateReadAccess, ControlStateWriteAccess, MaybeMisdirected, Unauthorized, }; use async_trait::async_trait; - use axum::body::Body; + use axum::body::{Body, Bytes}; use http::Request; use spacetimedb::auth::identity::{JwtError, JwtErrorKind, SpacetimeIdentityClaims}; use spacetimedb::auth::token_validation::{TokenSigner, TokenValidationError, TokenValidator}; @@ -1670,7 +1679,8 @@ mod tests { use spacetimedb_paths::FromPathUnchecked; use spacetimedb_schema::auto_migrate::{MigrationPolicy, PrettyPrintStyle}; use std::collections::HashMap; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::convert::Infallible; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use tower::util::ServiceExt; #[derive(Clone, Default)] @@ -2159,6 +2169,31 @@ mod tests { remove_http_response_size_metric(arbitrary_identity); } + #[tokio::test] + async fn resolving_middleware_drains_request_body_before_not_found() { + let body_was_polled = std::sync::Arc::new(AtomicBool::new(false)); + let body_was_polled_by_stream = body_was_polled.clone(); + let body = Body::from_stream(futures::stream::once(async move { + body_was_polled_by_stream.store(true, Ordering::Relaxed); + Ok::<_, Infallible>(Bytes::from_static(b"module")) + })); + let state = DummyState::new(); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { "not reached" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot(Request::builder().uri("/unresolved-name").body(body).unwrap()) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert!(body_was_polled.load(Ordering::Relaxed)); + } + #[tokio::test] async fn resolving_middleware_attaches_database_and_resolves_a_name_once() { let database_identity = test_identity(17);