security: isolate non-pgwire connection failures

Author: farhan-syah

nodedb/src/control/server/http/routes/crdt.rs

@@ -115,7 +115,7 @@ pub async fn crdt_apply(
     // replication. A local-only dispatch would land it on the receiving node only
     // — lost to followers and entirely on leader failover. This handler is scoped
     // to the default database (matching its surrogate assignment above).
-    state.shared.tenant_request_start(identity.tenant_id);
+    let _request = state.shared.tenant_request_guard(identity.tenant_id);
     let policy = crate::control::crdt_post_image_policy::ExternalCrdtPostImagePolicy::from_identity(
         identity.tenant_id,
         crate::types::DatabaseId::DEFAULT,
@@ -138,7 +138,6 @@ pub async fn crdt_apply(
         },
     )
     .await;
-    state.shared.tenant_request_end(identity.tenant_id);
 
     result.map_err(ApiError::from)?;
 

nodedb/src/control/server/http/routes/query/materialized.rs

@@ -129,12 +129,12 @@ pub async fn query(
     }
 
     // Track active request for quota accounting.
-    state.shared.tenant_request_start(tenant_id);
+    let _request = state.shared.tenant_request_guard(tenant_id);
 
     // Execute each task via the SPSC bridge.
     let mut result_rows = Vec::new();
 
-    let result = async {
+    async {
         for (task, authorized_task) in tasks.into_iter().zip(authorized_tasks) {
             // `INSERT ... SELECT` is orchestrated on the Control Plane: the
             // source is scanned, each target row gets its OWN fresh, registered
@@ -308,10 +308,7 @@ pub async fn query(
 
         Ok(axum::Json(HttpQueryResponse::ok(result_rows)))
     }
-    .await;
-
-    state.shared.tenant_request_end(tenant_id);
-    result
+    .await
 }
 
 fn ddl_error_to_api(error: crate::control::server::shared::ddl::DdlError) -> ApiError {

nodedb/src/control/server/http/routes/query/ndjson.rs

@@ -152,7 +152,7 @@ pub async fn query_ndjson(
         Err(error) => return ApiError::from(crate::Error::from(error)).into_response(),
     };
 
-    state.shared.tenant_request_start(tenant_id);
+    let _request = state.shared.tenant_request_guard(tenant_id);
 
     let mut ndjson = String::new();
     for (task, authorized_task) in tasks.into_iter().zip(authorized_tasks) {
@@ -262,8 +262,6 @@ pub async fn query_ndjson(
         }
     }
 
-    state.shared.tenant_request_end(tenant_id);
-
     Response::builder()
         .header("Content-Type", "application/x-ndjson")
         .body(axum::body::Body::from(ndjson))

nodedb/src/control/server/http/routes/ws_rpc/execute_sql.rs

@@ -70,7 +70,7 @@ pub async fn execute_sql(
     )?
     .into_tasks();
 
-    shared.tenant_request_start(tenant_id);
+    let _request = shared.tenant_request_guard(tenant_id);
 
     let mut results = Vec::new();
     for (task, authorized_task) in tasks.into_iter().zip(authorized_tasks) {
@@ -98,10 +98,7 @@ pub async fn execute_sql(
                         }
                     }
                 }
-                Err(e) => {
-                    shared.tenant_request_end(tenant_id);
-                    return Err(e);
-                }
+                Err(e) => return Err(e),
             }
             continue;
         }
@@ -138,10 +135,7 @@ pub async fn execute_sql(
                         }
                     }
                 }
-                Err(e) => {
-                    shared.tenant_request_end(tenant_id);
-                    return Err(e);
-                }
+                Err(e) => return Err(e),
             }
             continue;
         }
@@ -182,10 +176,7 @@ pub async fn execute_sql(
                         }
                     }
                 }
-                Err(e) => {
-                    shared.tenant_request_end(tenant_id);
-                    return Err(e);
-                }
+                Err(e) => return Err(e),
             }
             continue;
         }
@@ -225,15 +216,10 @@ pub async fn execute_sql(
                     }
                 }
             }
-            Err(e) => {
-                shared.tenant_request_end(tenant_id);
-                return Err(e);
-            }
+            Err(e) => return Err(e),
         }
     }
 
-    shared.tenant_request_end(tenant_id);
-
     match results.len() {
         0 => Ok(serde_json::Value::Null),
         1 => Ok(results

nodedb/src/control/server/http/routes/ws_rpc/handler.rs

@@ -21,17 +21,20 @@
 
 use std::sync::Arc;
 
+use tokio::task::JoinHandle;
+
 use axum::extract::ws::{Message, WebSocket};
 use axum::extract::{State, WebSocketUpgrade};
 use axum::http::HeaderMap;
 use axum::response::IntoResponse;
 use futures::{SinkExt, StreamExt};
-use tracing::debug;
+use tracing::{debug, warn};
 
 use crate::control::change_stream::LiveSubscriptionSet;
 use crate::control::security::audit::ArcAuditEmitter;
 use crate::control::server::http::auth::{AppState, ResolvedIdentity};
 use crate::control::server::shared::authorization::authorize_database;
+use crate::control::server::shared::{ConnectionFutureOutcome, isolate_connection_future};
 use crate::control::state::SharedState;
 use crate::types::DatabaseId;
 
@@ -66,12 +69,64 @@ pub async fn ws_handler(
         .into_response();
     }
     let trace_id = crate::control::trace_context::extract_from_headers(&headers);
-    ws.on_upgrade(move |socket| {
-        handle_ws_connection(socket, state, identity, database_id, trace_id)
+    ws.on_upgrade(move |socket| async move {
+        match isolate_connection_future(handle_ws_connection(
+            socket,
+            state,
+            identity,
+            database_id,
+            trace_id,
+        ))
+        .await
+        {
+            ConnectionFutureOutcome::Completed(()) => {}
+            ConnectionFutureOutcome::Panicked => {
+                warn!("ws-rpc connection panicked");
+            }
+        }
     })
     .into_response()
 }
 
+/// Owns a sender task until normal completion or connection teardown.
+///
+/// Dropping the guard aborts the task synchronously. This ensures a panic or
+/// cancellation of the connection future cannot detach the sender task.
+struct AbortOnDropJoinHandle {
+    handle: Option<JoinHandle<()>>,
+}
+
+impl AbortOnDropJoinHandle {
+    fn new(handle: JoinHandle<()>) -> Self {
+        Self {
+            handle: Some(handle),
+        }
+    }
+
+    /// Join the sender after the channel is closed. A sender panic is observed
+    /// without inspecting its payload or logging its join error.
+    async fn finish(mut self) {
+        let terminated_unexpectedly = match self.handle.as_mut() {
+            Some(handle) => handle.await.is_err(),
+            None => false,
+        };
+        // Remove only after the await completes. If this future is cancelled
+        // while pending, `Drop` still owns and aborts the sender task.
+        self.handle.take();
+        if terminated_unexpectedly {
+            warn!("ws-rpc sender task terminated unexpectedly");
+        }
+    }
+}
+
+impl Drop for AbortOnDropJoinHandle {
+    fn drop(&mut self) {
+        if let Some(handle) = self.handle.take() {
+            handle.abort();
+        }
+    }
+}
+
 /// Handle a single WebSocket connection.
 async fn handle_ws_connection(
     socket: WebSocket,
@@ -87,7 +142,7 @@ async fn handle_ws_connection(
     // 256 messages provides ~10s of buffer at 25 events/sec.
     let (live_tx, mut live_rx) = tokio::sync::mpsc::channel::<String>(256);
 
-    let send_handle = tokio::spawn(async move {
+    let sender = AbortOnDropJoinHandle::new(tokio::spawn(async move {
         loop {
             tokio::select! {
                 Some(msg) = live_rx.recv() => {
@@ -99,7 +154,7 @@ async fn handle_ws_connection(
                 else => break,
             }
         }
-    });
+    }));
 
     // Session ID is set only after successful auth inside process_message.
     let mut authenticated_session_id: Option<String> = None;
@@ -157,7 +212,7 @@ async fn handle_ws_connection(
     drop(live_set);
 
     drop(live_tx);
-    let _ = send_handle.await;
+    sender.finish().await;
     debug!("WebSocket RPC connection closed");
 }
 
@@ -189,3 +244,88 @@ pub fn save_ws_session(shared: &SharedState, session_id: &str) {
         "WS session saved for reconnect"
     );
 }
+
+#[cfg(test)]
+mod tests {
+    use std::future::pending;
+
+    use tokio::sync::oneshot;
+
+    use super::{AbortOnDropJoinHandle, ConnectionFutureOutcome, isolate_connection_future};
+
+    struct AbortSignal(Option<oneshot::Sender<()>>);
+
+    impl Drop for AbortSignal {
+        fn drop(&mut self) {
+            if let Some(sender) = self.0.take() {
+                let _ = sender.send(());
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn sender_is_aborted_when_guard_is_dropped() {
+        let (started_tx, started_rx) = oneshot::channel();
+        let (dropped_tx, dropped_rx) = oneshot::channel();
+        let sender = AbortOnDropJoinHandle::new(tokio::spawn(async move {
+            let _signal = AbortSignal(Some(dropped_tx));
+            let _ = started_tx.send(());
+            pending::<()>().await;
+        }));
+
+        assert!(started_rx.await.is_ok());
+        drop(sender);
+
+        assert!(dropped_rx.await.is_ok());
+    }
+
+    #[tokio::test]
+    async fn cancellation_while_finishing_aborts_sender() {
+        let (started_tx, started_rx) = oneshot::channel();
+        let (dropped_tx, dropped_rx) = oneshot::channel();
+        let sender = AbortOnDropJoinHandle::new(tokio::spawn(async move {
+            let _signal = AbortSignal(Some(dropped_tx));
+            let _ = started_tx.send(());
+            pending::<()>().await;
+        }));
+        assert!(started_rx.await.is_ok());
+
+        let finishing = tokio::spawn(sender.finish());
+        tokio::task::yield_now().await;
+        finishing.abort();
+        let _ = finishing.await;
+
+        assert!(dropped_rx.await.is_ok());
+    }
+
+    #[tokio::test]
+    async fn sender_finishes_normally() {
+        let (completed_tx, completed_rx) = oneshot::channel();
+        let sender = AbortOnDropJoinHandle::new(tokio::spawn(async move {
+            let _ = completed_tx.send(());
+        }));
+
+        sender.finish().await;
+
+        assert!(completed_rx.await.is_ok());
+    }
+
+    #[tokio::test]
+    async fn isolation_catches_connection_panic_and_releases_sender() {
+        let (started_tx, started_rx) = oneshot::channel();
+        let (dropped_tx, dropped_rx) = oneshot::channel();
+        let outcome = isolate_connection_future(async move {
+            let _sender = AbortOnDropJoinHandle::new(tokio::spawn(async move {
+                let _signal = AbortSignal(Some(dropped_tx));
+                let _ = started_tx.send(());
+                pending::<()>().await;
+            }));
+            let _ = started_rx.await;
+            panic!("simulated websocket connection panic");
+        })
+        .await;
+
+        assert!(matches!(outcome, ConnectionFutureOutcome::Panicked));
+        assert!(dropped_rx.await.is_ok());
+    }
+}

nodedb/src/control/server/ilp_batch.rs

@@ -191,30 +191,11 @@ pub(crate) async fn flush_authenticated_ilp_batch(
     // all collection permissions have passed. The tenant is never caller input.
     let tenant_id = identity.tenant_id;
     state.check_tenant_quota(tenant_id)?;
-    let _request = TenantRequestAccounting::start(state, tenant_id);
+    let _request = state.tenant_request_guard(tenant_id);
 
     flush_ilp_batch_inner(state, identity, database_id, groups).await
 }
 
-/// Cancellation-safe tenant request accounting for one ILP batch.
-struct TenantRequestAccounting<'a> {
-    state: &'a SharedState,
-    tenant_id: TenantId,
-}
-
-impl<'a> TenantRequestAccounting<'a> {
-    fn start(state: &'a SharedState, tenant_id: TenantId) -> Self {
-        state.tenant_request_start(tenant_id);
-        Self { state, tenant_id }
-    }
-}
-
-impl Drop for TenantRequestAccounting<'_> {
-    fn drop(&mut self) {
-        self.state.tenant_request_end(self.tenant_id);
-    }
-}
-
 /// Inner dispatch logic for ILP batch (separated for clean quota bookkeeping).
 async fn flush_ilp_batch_inner(
     state: &SharedState,