diff --git a/Justfile b/Justfile index ec5187d84e..1bd021a0a6 100644 --- a/Justfile +++ b/Justfile @@ -508,7 +508,8 @@ test-unit: # #[ignore]d Postgres suites — so these non-ignored tests ran in no lane # and a red one could ship green (exactly how a broken admin test slipped # past every gate once). Scoped to api::admin, not the whole buzz-relay - # --lib, because api::media has non-ignored tests that require Postgres. + # --lib, because much of handlers::req has non-ignored tests that require + # Postgres. # DB-backed api::admin tests are #[ignore]d and run in the PostgreSQL # lane; the non-ignored ones reject before touching the database. Any # new non-ignored test here must stay DB-free: without a database, a diff --git a/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests/failure_postgres_tests.rs b/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests/failure_postgres_tests.rs index c7f7e7bf92..d795917cb6 100644 --- a/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests/failure_postgres_tests.rs +++ b/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests/failure_postgres_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::test_support::PausingProxy; use std::{ future::Future, sync::atomic::{AtomicUsize, Ordering}, @@ -9,23 +10,31 @@ use tracing_subscriber::{layer::Context, prelude::*, registry::LookupSpan, Layer // Pause the HTTP future after the first real auxiliary page completes. This // observes a production span rather than replacing the database or closure. -struct AuxPages(Arc); +// With a proxy, the first page's span close also pauses it. That runs inside +// the request's poll, before the second page's query is written, so the second +// page stays pending however fast the database answers. +struct AuxPages(Arc, Option>); impl Layer for AuxPages where S: tracing::Subscriber + for<'a> LookupSpan<'a>, { fn on_close(&self, id: tracing::Id, ctx: Context<'_, S>) { - if ctx.span(&id).unwrap().metadata().name() == "thread_window_aux" { - self.0.fetch_add(1, Ordering::SeqCst); + if ctx.span(&id).unwrap().metadata().name() == "thread_window_aux" + && self.0.fetch_add(1, Ordering::SeqCst) == 0 + { + if let Some(proxy) = &self.1 { + proxy.set_paused(true); + } } } } async fn pause_after_aux_page( request: F, + proxy: Option>, ) -> std::pin::Pin>> { let pages = Arc::new(AtomicUsize::new(0)); - let subscriber = tracing_subscriber::registry().with(AuxPages(pages.clone())); + let subscriber = tracing_subscriber::registry().with(AuxPages(pages.clone(), proxy)); let mut request = Box::pin(request.with_subscriber(subscriber)); tokio::time::timeout( Duration::from_secs(5), @@ -63,9 +72,17 @@ async fn thread_window_bridge_restarts_both_aux_hops_after_replica_failure() { let reader_name = format!("tw-reader-{}", Uuid::new_v4()); let options: sqlx::postgres::PgConnectOptions = crate::test_support::database_url().parse().unwrap(); + // The replica reads through a proxy that holds the second aux page. + let proxy = + Arc::new(PausingProxy::spawn(options.get_host().to_owned(), options.get_port()).await); let reader = sqlx::postgres::PgPoolOptions::new() .max_connections(1) - .connect_with(options.application_name(&reader_name)) + .connect_with( + options + .host("127.0.0.1") + .port(proxy.port) + .application_name(&reader_name), + ) .await .unwrap(); let db = buzz_db::Db::from_pools(f.pool.clone(), reader.clone()); @@ -75,7 +92,11 @@ async fn thread_window_bridge_restarts_both_aux_hops_after_replica_failure() { let mut filter = f.filter(); filter["until"] = json!(f.root.created_at.as_secs() + 1); filter["before_id"] = json!("00".repeat(32)); - let request = pause_after_aux_page(f.post(&f.keys, "/query", json!([filter]))).await; + let request = pause_after_aux_page( + f.post(&f.keys, "/query", json!([filter])), + Some(proxy.clone()), + ) + .await; let held: i64 = sqlx::query_scalar("SELECT count(*) FROM pg_stat_activity WHERE application_name=$1 AND xact_start IS NOT NULL") .bind(&reader_name).fetch_one(&f.pool).await.unwrap(); assert_eq!(held, 1, "must actually hold the proved replica transaction"); @@ -87,6 +108,9 @@ async fn thread_window_bridge_restarts_both_aux_hops_after_replica_failure() { .execute(&f.pool) .await .unwrap(); + // Deliver the termination, so the held page fails and the closure + // restarts on the writer. + proxy.set_paused(false); let (status, body) = request.await; assert_eq!(status, StatusCode::OK, "{body}"); assert_eq!(f.bounds(&body, &filter)["has_more"], false); @@ -194,7 +218,7 @@ async fn assert_aux_access_change(grant: bool) { "pre-transition visibility control" ); - let request = pause_after_aux_page(f.post(&f.keys, "/query", json!([filter]))).await; + let request = pause_after_aux_page(f.post(&f.keys, "/query", json!([filter])), None).await; set_access(&f, aux_channel, grant).await; let current = f .state @@ -331,7 +355,7 @@ async fn thread_window_batch_discards_earlier_windows_after_revocation() { let (status, before) = f.post(&f.keys, "/query", batch.clone()).await; assert_eq!(status, StatusCode::OK, "{before}"); assert_eq!(ids(&before, Some(39007)).len(), 2); - let request = pause_after_aux_page(f.post(&f.keys, "/query", batch.clone())).await; + let request = pause_after_aux_page(f.post(&f.keys, "/query", batch.clone()), None).await; set_access(&f, f.channel, false).await; let (status, body) = request.await; assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE, "{body}"); diff --git a/crates/buzz-relay/src/api/media.rs b/crates/buzz-relay/src/api/media.rs index 2592574b9b..e86f37a497 100644 --- a/crates/buzz-relay/src/api/media.rs +++ b/crates/buzz-relay/src/api/media.rs @@ -1276,14 +1276,12 @@ fn blossom_strictness_from_state(state: &AppState) -> BlossomStrictness { #[cfg(test)] mod tests { use super::*; - use std::sync::Arc; use axum::{ - body::{to_bytes, Body}, - http::{header, Request, StatusCode}, + body::to_bytes, + http::{header, StatusCode}, }; - use nostr::{EventBuilder, JsonUtil, Keys, Kind, Tag, Timestamp}; - use tower::ServiceExt; + use nostr::{JsonUtil, Timestamp}; use uuid::Uuid; const VALID_HASH: &str = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789"; @@ -1702,216 +1700,6 @@ mod tests { assert!(should_stream_as_video(bytes)); } - async fn test_state() -> Arc { - let mut config = crate::config::Config::for_test(); // [FI-TRACE-ENV-RACE] - config.require_relay_membership = false; - config.redis_url = "redis://127.0.0.1:1".to_string(); - config.media_uploads_per_minute = 1; - config.media_max_concurrent_uploads = 2; - config.media_max_concurrent_uploads_per_pubkey = 1; - - let pool = sqlx::PgPool::connect_lazy(&config.database_url).expect("lazy pg pool"); - let db = buzz_db::Db::from_pool(pool.clone()); - db.ensure_configured_community("relay.example") - .await - .expect("seed relay.example community for host-bound media tests"); - let redis_pool = deadpool_redis::Config::from_url(&config.redis_url) - .create_pool(Some(deadpool_redis::Runtime::Tokio1)) - .expect("redis pool"); - let pubsub = Arc::new( - buzz_pubsub::PubSubManager::new(&config.redis_url, redis_pool.clone()) - .await - .expect("pubsub manager"), - ); - let audit = buzz_audit::AuditService::new(pool.clone()); - let auth = buzz_auth::AuthService::new(config.auth.clone()); - let search = buzz_search::SearchService::new(pool.clone()); - let workflow_engine = Arc::new(buzz_workflow::WorkflowEngine::new( - db.clone(), - buzz_workflow::WorkflowConfig::default(), - )); - let media_storage = buzz_media::MediaStorage::new(&config.media).expect("media storage"); - let (state, _audit_shutdown) = AppState::new( - config, - db, - redis_pool, - audit, - pubsub, - auth, - search, - workflow_engine, - nostr::Keys::generate(), - media_storage, - ); - Arc::new(state) - } - - async fn media_get_auth_router() -> axum::Router { - let state = test_state().await; - axum::Router::new() - .route( - "/media/{sha256_ext}", - axum::routing::get(get_blob).head(head_blob), - ) - .with_state(state) - } - - fn media_get_auth_header(keys: &Keys, tags: Vec) -> String { - let event = EventBuilder::new(Kind::from(24242), "Get media") - .tags(tags) - .sign_with_keys(keys) - .expect("sign get auth"); - format!( - "Nostr {}", - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(event.as_json().as_bytes()) - ) - } - - fn media_get_tags_for(host: &str, sha256: Option<&str>) -> Vec { - let now = Timestamp::now().as_secs(); - let expiration = (now + 55).to_string(); - let mut tags = vec![ - Tag::parse(["t", "get"]).expect("t tag"), - Tag::parse(["expiration", &expiration]).expect("expiration tag"), - Tag::parse(["server", host]).expect("server tag"), - ]; - if let Some(sha256) = sha256 { - tags.push(Tag::parse(["x", sha256]).expect("x tag")); - } - tags - } - - fn media_request(method: &str, auth: Option) -> Request { - let mut builder = Request::builder() - .method(method) - .uri(format!("/media/{VALID_HASH}.jpg")) - .header(header::HOST, "relay.example"); - if let Some(auth) = auth { - builder = builder.header(header::AUTHORIZATION, auth); - } - builder.body(Body::empty()).expect("request") - } - - #[tokio::test] - async fn media_reads_reject_unauthenticated_get_and_head_before_sidecar_gate() { - for method in ["GET", "HEAD"] { - let response = media_get_auth_router() - .await - .oneshot(media_request(method, None)) - .await - .expect("response"); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED, "{method}"); - } - } - - #[tokio::test] - async fn media_read_with_valid_server_scoped_token_reaches_sidecar_gate() { - let keys = Keys::generate(); - let auth = media_get_auth_header(&keys, media_get_tags_for("relay.example", None)); - let response = media_get_auth_router() - .await - .oneshot(media_request("GET", Some(auth))) - .await - .expect("response"); - - assert_eq!(response.status(), StatusCode::NOT_FOUND); - } - - #[tokio::test] - async fn media_read_rejects_upload_verb_wrong_server_and_wrong_x() { - let keys = Keys::generate(); - let now = Timestamp::now().as_secs(); - let expiration = (now + 55).to_string(); - let cases = [ - vec![ - Tag::parse(["t", "upload"]).expect("t tag"), - Tag::parse(["expiration", &expiration]).expect("expiration tag"), - Tag::parse(["server", "relay.example"]).expect("server tag"), - ], - vec![ - Tag::parse(["t", "get"]).expect("t tag"), - Tag::parse(["expiration", &expiration]).expect("expiration tag"), - Tag::parse(["server", "evil.example"]).expect("server tag"), - ], - vec![ - Tag::parse(["t", "get"]).expect("t tag"), - Tag::parse(["expiration", &expiration]).expect("expiration tag"), - Tag::parse(["x", &"f".repeat(64)]).expect("x tag"), - ], - ]; - - for tags in cases { - let auth = media_get_auth_header(&keys, tags); - let response = media_get_auth_router() - .await - .oneshot(media_request("GET", Some(auth))) - .await - .expect("response"); - - assert_ne!(response.status(), StatusCode::NOT_FOUND); - assert!( - response.status() == StatusCode::UNAUTHORIZED - || response.status() == StatusCode::FORBIDDEN, - "unexpected status {}", - response.status() - ); - } - } - - #[tokio::test] - async fn media_read_accepts_range_header_only_after_auth() { - let keys = Keys::generate(); - let auth = media_get_auth_header(&keys, media_get_tags_for("relay.example", None)); - let mut request = media_request("GET", Some(auth)); - request - .headers_mut() - .insert(header::RANGE, "bytes=0-0".parse().expect("range header")); - - let response = media_get_auth_router() - .await - .oneshot(request) - .await - .expect("response"); - - assert_eq!(response.status(), StatusCode::NOT_FOUND); - } - - #[tokio::test] - async fn upload_rate_limiter_is_scoped_by_community() { - let state = test_state().await; - let pubkey = nostr::Keys::generate().public_key(); - let community_a = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xAAAA)); - let community_b = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xBBBB)); - - assert!(!upload_rate_limited(&state, community_a, &pubkey)); - assert!(upload_rate_limited(&state, community_a, &pubkey)); - assert!( - !upload_rate_limited(&state, community_b, &pubkey), - "A's exhausted upload budget must not rate-limit the same key in B" - ); - } - - #[tokio::test] - async fn upload_concurrency_limit_is_scoped_by_community() { - let state = test_state().await; - let pubkey = nostr::Keys::generate().public_key(); - let community_a = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xAAAA)); - let community_b = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xBBBB)); - - let permit_a = - acquire_upload_permit(&state, community_a, &pubkey).expect("first A upload allowed"); - assert!(matches!( - acquire_upload_permit(&state, community_a, &pubkey), - Err(MediaError::UploadConcurrencyLimitReached) - )); - let permit_b = acquire_upload_permit(&state, community_b, &pubkey) - .expect("A's in-flight upload must not block B"); - - drop(permit_b); - drop(permit_a); - } - #[test] fn test_validate_media_path_bare_hash() { assert!(validate_media_path(VALID_HASH).is_ok()); @@ -2114,7 +1902,7 @@ mod tests { // ── NIP-FI media regression tests ──────────────────────────────────────── // // These postgres-backed tests prove the NIP-FI admission gate is wired into - // the media upload and read routes at the router level. They all require a + // the media upload and read routes at the router level. They require a // live Postgres + Redis and are tagged #[ignore = "requires Postgres"]. // // Verified contract rows (NIP-FI.md rejection table): @@ -2166,11 +1954,7 @@ mod tests { const TEST_ISSUER: &str = "https://issuer.example"; const TEST_AUDIENCE: &str = "https://relay.example"; const TEST_KID: &str = "test-key-1"; - const TEST_EC_PKCS8_PEM: &str = "-----BEGIN PRIVATE KEY-----\n\ - MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgcnxDM4EiirH9dHUE\n\ - WZc759TX4s5PAn8kO5ovXSnGxCWhRANCAARFb6ZnsfkqOOXyEhj3KBQphGKF4vTa\n\ - zhebbavbZ1ZoklqkF1cGg+jTO7rONAVEzXvXUWtV6CdDV+rybiVmFP2w\n\ - -----END PRIVATE KEY-----\n"; + const TEST_EC_PKCS8_PEM: &str = "-----BEGIN PRIVATE KEY-----\nMIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgcnxDM4EiirH9dHUE\nWZc759TX4s5PAn8kO5ovXSnGxCWhRANCAARFb6ZnsfkqOOXyEhj3KBQphGKF4vTa\nzhebbavbZ1ZoklqkF1cGg+jTO7rONAVEzXvXUWtV6CdDV+rybiVmFP2w\n-----END PRIVATE KEY-----\n"; // sadscan:disable np.pem.1 kingfisher.privkey.2 -- public ES256 test fixture /// Build an AppState with NIP-FI Enforce and a real injected P-256 verifier. async fn media_enforce_test_state() -> Option> { @@ -4415,5 +4199,228 @@ mod tests { ); } } + + // ── Host-bound media reads and per-community upload limits ───────────── + // + // `test_state` seeds the relay.example community, so these need Postgres + // too. Its Redis URL is unreachable on purpose: none of them use Redis. + + async fn test_state() -> Arc { + let mut config = crate::config::Config::for_test(); // [FI-TRACE-ENV-RACE] + config.database_url = crate::test_support::database_url(); + config.require_relay_membership = false; + config.redis_url = "redis://127.0.0.1:1".to_string(); + config.media_uploads_per_minute = 1; + config.media_max_concurrent_uploads = 2; + config.media_max_concurrent_uploads_per_pubkey = 1; + + let pool = sqlx::PgPool::connect_lazy(&config.database_url).expect("lazy pg pool"); + let db = buzz_db::Db::from_pool(pool.clone()); + db.ensure_configured_community("relay.example") + .await + .expect("seed relay.example community for host-bound media tests"); + let redis_pool = deadpool_redis::Config::from_url(&config.redis_url) + .create_pool(Some(deadpool_redis::Runtime::Tokio1)) + .expect("redis pool"); + let pubsub = Arc::new( + buzz_pubsub::PubSubManager::new(&config.redis_url, redis_pool.clone()) + .await + .expect("pubsub manager"), + ); + let audit = buzz_audit::AuditService::new(pool.clone()); + let auth = buzz_auth::AuthService::new(config.auth.clone()); + let search = buzz_search::SearchService::new(pool.clone()); + let workflow_engine = Arc::new(buzz_workflow::WorkflowEngine::new( + db.clone(), + buzz_workflow::WorkflowConfig::default(), + )); + let media_storage = + buzz_media::MediaStorage::new(&config.media).expect("media storage"); + let (state, _audit_shutdown) = AppState::new( + config, + db, + redis_pool, + audit, + pubsub, + auth, + search, + workflow_engine, + nostr::Keys::generate(), + media_storage, + ); + Arc::new(state) + } + + async fn media_get_auth_router() -> axum::Router { + let state = test_state().await; + axum::Router::new() + .route( + "/media/{sha256_ext}", + axum::routing::get(get_blob).head(head_blob), + ) + .with_state(state) + } + + fn media_get_auth_header(keys: &Keys, tags: Vec) -> String { + let event = EventBuilder::new(Kind::from(24242), "Get media") + .tags(tags) + .sign_with_keys(keys) + .expect("sign get auth"); + format!( + "Nostr {}", + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(event.as_json().as_bytes()) + ) + } + + fn media_get_tags_for(host: &str, sha256: Option<&str>) -> Vec { + let now = Timestamp::now().as_secs(); + let expiration = (now + 55).to_string(); + let mut tags = vec![ + Tag::parse(["t", "get"]).expect("t tag"), + Tag::parse(["expiration", &expiration]).expect("expiration tag"), + Tag::parse(["server", host]).expect("server tag"), + ]; + if let Some(sha256) = sha256 { + tags.push(Tag::parse(["x", sha256]).expect("x tag")); + } + tags + } + + fn media_request(method: &str, auth: Option) -> Request { + let mut builder = Request::builder() + .method(method) + .uri(format!("/media/{VALID_HASH}.jpg")) + .header(header::HOST, "relay.example"); + if let Some(auth) = auth { + builder = builder.header(header::AUTHORIZATION, auth); + } + builder.body(Body::empty()).expect("request") + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn media_reads_reject_unauthenticated_get_and_head_before_sidecar_gate() { + for method in ["GET", "HEAD"] { + let response = media_get_auth_router() + .await + .oneshot(media_request(method, None)) + .await + .expect("response"); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED, "{method}"); + } + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn media_read_with_valid_server_scoped_token_reaches_sidecar_gate() { + let keys = Keys::generate(); + let auth = media_get_auth_header(&keys, media_get_tags_for("relay.example", None)); + let response = media_get_auth_router() + .await + .oneshot(media_request("GET", Some(auth))) + .await + .expect("response"); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn media_read_rejects_upload_verb_wrong_server_and_wrong_x() { + let keys = Keys::generate(); + let now = Timestamp::now().as_secs(); + let expiration = (now + 55).to_string(); + let cases = [ + vec![ + Tag::parse(["t", "upload"]).expect("t tag"), + Tag::parse(["expiration", &expiration]).expect("expiration tag"), + Tag::parse(["server", "relay.example"]).expect("server tag"), + ], + vec![ + Tag::parse(["t", "get"]).expect("t tag"), + Tag::parse(["expiration", &expiration]).expect("expiration tag"), + Tag::parse(["server", "evil.example"]).expect("server tag"), + ], + vec![ + Tag::parse(["t", "get"]).expect("t tag"), + Tag::parse(["expiration", &expiration]).expect("expiration tag"), + Tag::parse(["x", &"f".repeat(64)]).expect("x tag"), + ], + ]; + + for tags in cases { + let auth = media_get_auth_header(&keys, tags); + let response = media_get_auth_router() + .await + .oneshot(media_request("GET", Some(auth))) + .await + .expect("response"); + + assert_ne!(response.status(), StatusCode::NOT_FOUND); + assert!( + response.status() == StatusCode::UNAUTHORIZED + || response.status() == StatusCode::FORBIDDEN, + "unexpected status {}", + response.status() + ); + } + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn media_read_accepts_range_header_only_after_auth() { + let keys = Keys::generate(); + let auth = media_get_auth_header(&keys, media_get_tags_for("relay.example", None)); + let mut request = media_request("GET", Some(auth)); + request + .headers_mut() + .insert(header::RANGE, "bytes=0-0".parse().expect("range header")); + + let response = media_get_auth_router() + .await + .oneshot(request) + .await + .expect("response"); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn upload_rate_limiter_is_scoped_by_community() { + let state = test_state().await; + let pubkey = nostr::Keys::generate().public_key(); + let community_a = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xAAAA)); + let community_b = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xBBBB)); + + assert!(!upload_rate_limited(&state, community_a, &pubkey)); + assert!(upload_rate_limited(&state, community_a, &pubkey)); + assert!( + !upload_rate_limited(&state, community_b, &pubkey), + "A's exhausted upload budget must not rate-limit the same key in B" + ); + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn upload_concurrency_limit_is_scoped_by_community() { + let state = test_state().await; + let pubkey = nostr::Keys::generate().public_key(); + let community_a = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xAAAA)); + let community_b = buzz_core::CommunityId::from_uuid(Uuid::from_u128(0xBBBB)); + + let permit_a = acquire_upload_permit(&state, community_a, &pubkey) + .expect("first A upload allowed"); + assert!(matches!( + acquire_upload_permit(&state, community_a, &pubkey), + Err(MediaError::UploadConcurrencyLimitReached) + )); + let permit_b = acquire_upload_permit(&state, community_b, &pubkey) + .expect("A's in-flight upload must not block B"); + + drop(permit_b); + drop(permit_a); + } } } diff --git a/crates/buzz-relay/src/test_support.rs b/crates/buzz-relay/src/test_support.rs index 8d0dc6cacd..9856eff636 100644 --- a/crates/buzz-relay/src/test_support.rs +++ b/crates/buzz-relay/src/test_support.rs @@ -122,6 +122,72 @@ pub(crate) fn run_exact_test_child(test_name: &str, child_env: &str) { ); } +/// Loopback TCP proxy that can pause the upstream database. While paused, +/// every session holds the bytes it receives with both sockets open; on +/// resume the held bytes are delivered in order and the sessions carry on. +/// A trimmed copy of buzz-db's test-only `PausingProxy`, which is private to +/// that crate's unit tests. +#[cfg(test)] +pub(crate) struct PausingProxy { + pub(crate) port: u16, + paused: tokio::sync::watch::Sender, +} + +#[cfg(test)] +impl PausingProxy { + pub(crate) async fn spawn(upstream_host: String, upstream_port: u16) -> Self { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + async fn pump( + mut from: tokio::net::tcp::OwnedReadHalf, + mut to: tokio::net::tcp::OwnedWriteHalf, + mut paused: tokio::sync::watch::Receiver, + ) { + let mut buf = [0u8; 8192]; + loop { + let Ok(n) = from.read(&mut buf).await else { + return; + }; + if n == 0 { + return; + } + if paused.wait_for(|paused| !paused).await.is_err() { + return; + } + if to.write_all(&buf[..n]).await.is_err() { + return; + } + } + } + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind proxy"); + let port = listener.local_addr().expect("proxy addr").port(); + let (paused, accept_paused) = tokio::sync::watch::channel(false); + tokio::spawn(async move { + while let Ok((client, _)) = listener.accept().await { + let Ok(server) = + tokio::net::TcpStream::connect((upstream_host.as_str(), upstream_port)).await + else { + continue; + }; + let (client_read, client_write) = client.into_split(); + let (server_read, server_write) = server.into_split(); + tokio::spawn(pump(client_read, server_write, accept_paused.clone())); + tokio::spawn(pump(server_read, client_write, accept_paused.clone())); + } + }); + Self { port, paused } + } + + /// Hold or release every session's bytes. A pause applies to all bytes + /// the proxy reads after this call returns, on any runtime thread. + pub(crate) fn set_paused(&self, paused: bool) { + self.paused.send_replace(paused); + } +} + /// A database handle on which only restriction-state reads fail: its /// `search_path` puts a schema first whose `community_bans` is a view that /// errors when read, so every other table resolves to `public` as normal.