diff --git a/proxy/src/console_redirect_proxy.rs b/proxy/src/console_redirect_proxy.rs index ffa95241b9..e836863193 100644 --- a/proxy/src/console_redirect_proxy.rs +++ b/proxy/src/console_redirect_proxy.rs @@ -1,21 +1,19 @@ use crate::auth::backend::ConsoleRedirectBackend; -use crate::config::{ProxyConfig, ProxyProtocolV2}; -use crate::proxy::{ - prepare_client_connection, run_until_cancelled, ClientRequestError, ErrorSource, -}; +use crate::config::ProxyConfig; +use crate::metrics::Protocol; +use crate::proxy::{prepare_client_connection, transition_connection, ClientRequestError}; use crate::{ cancellation::CancellationHandlerMain, context::RequestMonitoring, - error::ReportableError, metrics::{Metrics, NumClientConnectionsGuard}, - protocol2::read_proxy_protocol, proxy::handshake::{handshake, HandshakeData}, }; use futures::TryFutureExt; +use std::net::IpAddr; use std::sync::Arc; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use tokio_util::sync::CancellationToken; -use tracing::{error, info, Instrument}; +use tracing::{info, Instrument}; use crate::proxy::{ connect_compute::{connect_to_compute, TcpMechanism}, @@ -33,107 +31,53 @@ pub async fn task_main( info!("proxy has shut down"); } - // When set for the server socket, the keepalive setting - // will be inherited by all accepted client sockets. - socket2::SockRef::from(&listener).set_keepalive(true)?; + super::connection_loop( + config, + listener, + cancellation_token, + Protocol::Tcp, + C { + config, + backend, + cancellation_handler, + }, + ) + .await +} - let connections = tokio_util::task::task_tracker::TaskTracker::new(); +#[derive(Clone)] +struct C { + config: &'static ProxyConfig, + backend: &'static ConsoleRedirectBackend, + cancellation_handler: Arc, +} - while let Some(accept_result) = - run_until_cancelled(listener.accept(), &cancellation_token).await - { - let (socket, peer_addr) = accept_result?; +impl super::ConnHandler for C { + async fn handle( + self, + session_id: uuid::Uuid, + peer_addr: IpAddr, + socket: crate::protocol2::ChainRW, + conn_gauge: crate::metrics::NumClientConnectionsGuard<'static>, + ) { + let ctx = RequestMonitoring::new(session_id, peer_addr, Protocol::Tcp, &self.config.region); + let span = ctx.span(); - let conn_gauge = Metrics::get() - .proxy - .client_connections - .guard(crate::metrics::Protocol::Tcp); + let startup = Box::pin( + handle_client( + self.config, + self.backend, + &ctx, + self.cancellation_handler, + socket, + conn_gauge, + ) + .instrument(span.clone()), + ); - let session_id = uuid::Uuid::new_v4(); - let cancellation_handler = Arc::clone(&cancellation_handler); - - tracing::info!(protocol = "tcp", %session_id, "accepted new TCP connection"); - - connections.spawn(async move { - let (socket, peer_addr) = match read_proxy_protocol(socket).await { - Err(e) => { - error!("per-client task finished with an error: {e:#}"); - return; - } - Ok((_socket, None)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { - error!("missing required proxy protocol header"); - return; - } - Ok((_socket, Some(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { - error!("proxy protocol header not supported"); - return; - } - Ok((socket, Some(addr))) => (socket, addr.ip()), - Ok((socket, None)) => (socket, peer_addr.ip()), - }; - - match socket.inner.set_nodelay(true) { - Ok(()) => {} - Err(e) => { - error!("per-client task finished with an error: failed to set socket option: {e:#}"); - return; - } - }; - - let ctx = RequestMonitoring::new( - session_id, - peer_addr, - crate::metrics::Protocol::Tcp, - &config.region, - ); - let span = ctx.span(); - - let startup = Box::pin( - handle_client( - config, - backend, - &ctx, - cancellation_handler, - socket, - conn_gauge, - ) - .instrument(span.clone()), - ); - let res = startup.await; - - match res { - Err(e) => { - // todo: log and push to ctx the error kind - ctx.set_error_kind(e.get_error_kind()); - error!(parent: &span, "per-client task finished with an error: {e:#}"); - } - Ok(None) => { - ctx.set_success(); - } - Ok(Some(p)) => { - ctx.set_success(); - ctx.log_connect(); - match p.proxy_pass().instrument(span.clone()).await { - Ok(()) => {} - Err(ErrorSource::Client(e)) => { - error!(parent: &span, "per-client task finished with an IO error from the client: {e:#}"); - } - Err(ErrorSource::Compute(e)) => { - error!(parent: &span, "per-client task finished with an IO error from the compute: {e:#}"); - } - } - } - } - }); + let res = startup.await; + transition_connection(ctx, res).await; } - - connections.close(); - drop(listener); - - // Drain connections - connections.wait().await; - - Ok(()) } pub(crate) async fn handle_client( diff --git a/proxy/src/lib.rs b/proxy/src/lib.rs index 333e09ab72..8dd86ce556 100644 --- a/proxy/src/lib.rs +++ b/proxy/src/lib.rs @@ -82,13 +82,16 @@ impl_trait_overcaptures, )] -use std::convert::Infallible; +use std::{convert::Infallible, future::Future, net::IpAddr}; use anyhow::{bail, Context}; use intern::{EndpointIdInt, EndpointIdTag, InternId}; -use tokio::task::JoinError; +use protocol2::{get_client_conn_info, ChainRW}; +use proxy::run_until_cancelled; +use tokio::{net::TcpStream, task::JoinError}; use tokio_util::sync::CancellationToken; -use tracing::warn; +use tracing::{error, warn}; +use uuid::Uuid; extern crate hyper0 as hyper; @@ -276,3 +279,81 @@ impl EndpointId { ProjectId(self.0.clone()) } } + +pub(crate) trait ConnHandler: Clone + Send + 'static { + fn handle( + self, + session_id: Uuid, + peer_addr: IpAddr, + stream: ChainRW, + conn_gauge: metrics::NumClientConnectionsGuard<'static>, + ) -> impl Future + Send; +} + +/// Accept connections, parse the proxy-protocol v2 header and spawn a tracked connection task. +pub(crate) async fn connection_loop( + config: &'static config::ProxyConfig, + listener: tokio::net::TcpListener, + cancellation_token: CancellationToken, + protocol: metrics::Protocol, + conn_handler: C, +) -> anyhow::Result<()> +where + C: ConnHandler, +{ + // When set for the server socket, the keepalive setting + // will be inherited by all accepted client sockets. + socket2::SockRef::from(&listener).set_keepalive(true)?; + + let connections = tokio_util::task::task_tracker::TaskTracker::new(); + + while let Some(accept_result) = + run_until_cancelled(listener.accept(), &cancellation_token).await + { + let (socket, peer_addr) = accept_result?; + + let conn_gauge = metrics::Metrics::get() + .proxy + .client_connections + .guard(protocol); + + let session_id = uuid::Uuid::new_v4(); + let conn_handler = conn_handler.clone(); + + tracing::info!(protocol = protocol.as_str(), %session_id, "accepted new TCP connection"); + + connections.spawn(async move { + let (socket, peer_addr) = match get_client_conn_info(socket, config.proxy_protocol_v2).await { + Err(e) => { + error!("per-client task finished with an error: {e:#}"); + return; + } + Ok((socket, Some(addr))) => (socket, addr), + Ok((socket, None)) => (socket, peer_addr.ip()), + }; + + match socket.inner.set_nodelay(true) { + Ok(()) => {} + Err(e) => { + error!("per-client task finished with an error: failed to set socket option: {e:#}"); + return; + } + }; + + conn_handler.handle( + session_id, + peer_addr, + socket, + conn_gauge, + ).await; + }); + } + + connections.close(); + drop(listener); + + // Drain connections + connections.wait().await; + + Ok(()) +} diff --git a/proxy/src/protocol2.rs b/proxy/src/protocol2.rs index 17764f78d1..39bb98db7b 100644 --- a/proxy/src/protocol2.rs +++ b/proxy/src/protocol2.rs @@ -2,15 +2,18 @@ use std::{ io, - net::SocketAddr, + net::{IpAddr, SocketAddr}, pin::Pin, task::{Context, Poll}, }; +use anyhow::bail; use bytes::BytesMut; use pin_project_lite::pin_project; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf}; +use crate::config::ProxyProtocolV2; + pin_project! { /// A chained [`AsyncRead`] with [`AsyncWrite`] passthrough pub(crate) struct ChainRW { @@ -60,7 +63,23 @@ const HEADER: [u8; 12] = [ 0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, ]; -pub(crate) async fn read_proxy_protocol( +pub(crate) async fn get_client_conn_info( + socket: T, + proxy_protocol_v2: ProxyProtocolV2, +) -> anyhow::Result<(ChainRW, Option)> { + match read_proxy_protocol(socket).await? { + (_socket, None) if proxy_protocol_v2 == ProxyProtocolV2::Required => { + bail!("missing required proxy protocol header"); + } + (_socket, Some(_)) if proxy_protocol_v2 == ProxyProtocolV2::Rejected => { + bail!("proxy protocol header not supported"); + } + (socket, Some(addr)) => Ok((socket, Some(addr.ip()))), + (socket, None) => Ok((socket, None)), + } +} + +async fn read_proxy_protocol( mut read: T, ) -> std::io::Result<(ChainRW, Option)> { let mut buf = BytesMut::with_capacity(128); diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs index 69d32f11b2..c7dbbca43d 100644 --- a/proxy/src/proxy/mod.rs +++ b/proxy/src/proxy/mod.rs @@ -10,8 +10,8 @@ pub(crate) mod wake_compute; pub use copy_bidirectional::copy_bidirectional_client_compute; pub use copy_bidirectional::ErrorSource; -use crate::config::ProxyProtocolV2; use crate::control_plane::provider::ControlPlaneBackend; +use crate::metrics::Protocol; use crate::{ auth, cancellation::{self, CancellationHandlerMain}, @@ -20,7 +20,6 @@ use crate::{ context::RequestMonitoring, error::ReportableError, metrics::{Metrics, NumClientConnectionsGuard}, - protocol2::read_proxy_protocol, proxy::handshake::{handshake, HandshakeData}, rate_limiter::EndpointRateLimiter, stream::{PqStream, Stream}, @@ -32,6 +31,7 @@ use once_cell::sync::OnceCell; use pq_proto::{BeMessage as Be, StartupMessageParams}; use regex::Regex; use smol_str::{format_smolstr, SmolStr}; +use std::net::IpAddr; use std::sync::Arc; use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; @@ -72,110 +72,91 @@ pub async fn task_main( info!("proxy has shut down"); } - // When set for the server socket, the keepalive setting - // will be inherited by all accepted client sockets. - socket2::SockRef::from(&listener).set_keepalive(true)?; + super::connection_loop( + config, + listener, + cancellation_token, + Protocol::Tcp, + C { + config, + auth_backend, + cancellation_handler, + endpoint_rate_limiter, + }, + ) + .await +} - let connections = tokio_util::task::task_tracker::TaskTracker::new(); +#[derive(Clone)] +struct C { + config: &'static ProxyConfig, + auth_backend: &'static ControlPlaneBackend, + cancellation_handler: Arc, + endpoint_rate_limiter: Arc, +} - while let Some(accept_result) = - run_until_cancelled(listener.accept(), &cancellation_token).await - { - let (socket, peer_addr) = accept_result?; +impl super::ConnHandler for C { + async fn handle( + self, + session_id: uuid::Uuid, + peer_addr: IpAddr, + socket: crate::protocol2::ChainRW, + conn_gauge: crate::metrics::NumClientConnectionsGuard<'static>, + ) { + let ctx = RequestMonitoring::new( + session_id, + peer_addr, + crate::metrics::Protocol::Tcp, + &self.config.region, + ); + let span = ctx.span(); - let conn_gauge = Metrics::get() - .proxy - .client_connections - .guard(crate::metrics::Protocol::Tcp); + let startup = Box::pin( + handle_client( + self.config, + self.auth_backend, + &ctx, + self.cancellation_handler, + socket, + ClientMode::Tcp, + self.endpoint_rate_limiter, + conn_gauge, + ) + .instrument(span.clone()), + ); - let session_id = uuid::Uuid::new_v4(); - let cancellation_handler = Arc::clone(&cancellation_handler); + let res = startup.await; + transition_connection(ctx, res).await; + } +} - tracing::info!(protocol = "tcp", %session_id, "accepted new TCP connection"); - let endpoint_rate_limiter2 = endpoint_rate_limiter.clone(); - - connections.spawn(async move { - let (socket, peer_addr) = match read_proxy_protocol(socket).await { - Err(e) => { - error!("per-client task finished with an error: {e:#}"); - return; - } - Ok((_socket, None)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { - error!("missing required proxy protocol header"); - return; - } - Ok((_socket, Some(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { - error!("proxy protocol header not supported"); - return; - } - Ok((socket, Some(addr))) => (socket, addr.ip()), - Ok((socket, None)) => (socket, peer_addr.ip()), - }; - - match socket.inner.set_nodelay(true) { +pub(crate) async fn transition_connection( + ctx: RequestMonitoring, + res: Result>, ClientRequestError>, +) { + let span = ctx.span(); + match res { + Err(e) => { + ctx.set_error_kind(e.get_error_kind()); + error!(parent: &span, "per-client task finished with an error: {e:#}"); + } + Ok(None) => { + ctx.set_success(); + } + Ok(Some(p)) => { + ctx.set_success(); + ctx.log_connect(); + match p.proxy_pass().instrument(span.clone()).await { Ok(()) => {} - Err(e) => { - error!("per-client task finished with an error: failed to set socket option: {e:#}"); - return; + Err(ErrorSource::Client(e)) => { + error!(parent: &span, "per-client task finished with an IO error from the client: {e:#}"); } - }; - - let ctx = RequestMonitoring::new( - session_id, - peer_addr, - crate::metrics::Protocol::Tcp, - &config.region, - ); - let span = ctx.span(); - - let startup = Box::pin( - handle_client( - config, - auth_backend, - &ctx, - cancellation_handler, - socket, - ClientMode::Tcp, - endpoint_rate_limiter2, - conn_gauge, - ) - .instrument(span.clone()), - ); - let res = startup.await; - - match res { - Err(e) => { - // todo: log and push to ctx the error kind - ctx.set_error_kind(e.get_error_kind()); - error!(parent: &span, "per-client task finished with an error: {e:#}"); - } - Ok(None) => { - ctx.set_success(); - } - Ok(Some(p)) => { - ctx.set_success(); - ctx.log_connect(); - match p.proxy_pass().instrument(span.clone()).await { - Ok(()) => {} - Err(ErrorSource::Client(e)) => { - error!(parent: &span, "per-client task finished with an IO error from the client: {e:#}"); - } - Err(ErrorSource::Compute(e)) => { - error!(parent: &span, "per-client task finished with an IO error from the compute: {e:#}"); - } - } + Err(ErrorSource::Compute(e)) => { + error!(parent: &span, "per-client task finished with an IO error from the compute: {e:#}"); } } - }); + } } - - connections.close(); - drop(listener); - - // Drain connections - connections.wait().await; - - Ok(()) } pub(crate) enum ClientMode { diff --git a/proxy/src/proxy/tests/mod.rs b/proxy/src/proxy/tests/mod.rs index 01bda948da..3ef00033ca 100644 --- a/proxy/src/proxy/tests/mod.rs +++ b/proxy/src/proxy/tests/mod.rs @@ -10,13 +10,14 @@ use super::*; use crate::auth::backend::{ ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, TestBackend, }; -use crate::config::{CertResolver, RetryConfig}; +use crate::config::{CertResolver, ProxyProtocolV2, RetryConfig}; use crate::control_plane::messages::{ControlPlaneError, Details, MetricsAuxInfo, Status}; use crate::control_plane::provider::{ CachedAllowedIps, CachedRoleSecret, ControlPlaneBackend, NodeInfoCache, }; use crate::control_plane::{self, CachedNodeInfo, NodeInfo}; use crate::error::ErrorKind; +use crate::protocol2::get_client_conn_info; use crate::{sasl, scram, BranchId, EndpointId, ProjectId}; use anyhow::{bail, Context}; use async_trait::async_trait; @@ -177,7 +178,7 @@ async fn dummy_proxy( tls: Option, auth: impl TestAuth + Send, ) -> anyhow::Result<()> { - let (client, _) = read_proxy_protocol(client).await?; + let (client, _) = get_client_conn_info(client, ProxyProtocolV2::Supported).await?; let mut stream = match handshake(&RequestMonitoring::test(), client, tls.as_ref(), false).await? { HandshakeData::Startup(stream, _) => stream, diff --git a/proxy/src/serverless/mod.rs b/proxy/src/serverless/mod.rs index 8b25f542d8..b2eafd7ce4 100644 --- a/proxy/src/serverless/mod.rs +++ b/proxy/src/serverless/mod.rs @@ -16,7 +16,6 @@ use atomic_take::AtomicTake; use bytes::Bytes; pub use conn_pool::GlobalConnPoolOptions; -use anyhow::Context; use futures::future::{select, Either}; use futures::TryFutureExt; use http::{Method, Response, StatusCode}; @@ -36,19 +35,18 @@ use crate::auth::ServerlessBackend; use crate::cancellation::CancellationHandlerMain; use crate::config::ProxyConfig; use crate::context::RequestMonitoring; -use crate::metrics::Metrics; -use crate::protocol2::{read_proxy_protocol, ChainRW}; -use crate::proxy::run_until_cancelled; +use crate::metrics::{Metrics, Protocol}; +use crate::protocol2::ChainRW; use crate::rate_limiter::EndpointRateLimiter; use crate::serverless::backend::PoolingBackend; use crate::serverless::http_util::{api_error_into_response, json_response}; -use std::net::{IpAddr, SocketAddr}; +use std::net::IpAddr; use std::pin::{pin, Pin}; use std::sync::Arc; use tokio::net::{TcpListener, TcpStream}; use tokio_util::sync::CancellationToken; -use tracing::{error, info, warn, Instrument}; +use tracing::{error, info, instrument, warn, Instrument}; use utils::http::error::ApiError; pub(crate) const SERVERLESS_DRIVER_SNI: &str = "api"; @@ -125,81 +123,100 @@ pub async fn task_main( } }; - let connections = tokio_util::task::task_tracker::TaskTracker::new(); - connections.close(); // allows `connections.wait to complete` + let requests = TaskTracker::new(); + requests.close(); // allows `requests.wait to complete` - while let Some(res) = run_until_cancelled(ws_listener.accept(), &cancellation_token).await { - let (conn, peer_addr) = res.context("could not accept TCP stream")?; - if let Err(e) = conn.set_nodelay(true) { - tracing::error!("could not set nodelay: {e}"); - continue; - } - let conn_id = uuid::Uuid::new_v4(); - let http_conn_span = tracing::info_span!("http_conn", ?conn_id); + crate::connection_loop( + config, + ws_listener, + cancellation_token.clone(), + Protocol::Http, + C { + config, + backend, + cancellation_handler, + endpoint_rate_limiter, + tls_acceptor, + requests: requests.clone(), + cancellation_token, + }, + ) + .await?; + requests.wait().await; + + Ok(()) +} + +#[derive(Clone)] +struct C { + config: &'static ProxyConfig, + backend: Arc, + cancellation_handler: Arc, + endpoint_rate_limiter: Arc, + tls_acceptor: Arc, + requests: TaskTracker, + cancellation_token: CancellationToken, +} + +impl super::ConnHandler for C { + #[instrument(name = "http_conn", skip_all, fields(conn_id))] + async fn handle( + self, + conn_id: uuid::Uuid, + peer_addr: IpAddr, + stream: ChainRW, + conn_gauge: crate::metrics::NumClientConnectionsGuard<'static>, + ) { + // try and close an old HTTP connection. + // picked at random let n_connections = Metrics::get() .proxy .client_connections .sample(crate::metrics::Protocol::Http); - tracing::trace!(?n_connections, threshold = ?config.http_config.client_conn_threshold, "check"); - if n_connections > config.http_config.client_conn_threshold { + tracing::trace!(?n_connections, threshold = ?self.config.http_config.client_conn_threshold, "check"); + if n_connections > self.config.http_config.client_conn_threshold { tracing::trace!("attempting to cancel a random connection"); - if let Some(token) = config.http_config.cancel_set.take() { + if let Some(token) = self.config.http_config.cancel_set.take() { tracing::debug!("cancelling a random connection"); token.cancel(); } } - let conn_token = cancellation_token.child_token(); - let tls_acceptor = tls_acceptor.clone(); - let backend = backend.clone(); - let connections2 = connections.clone(); - let cancellation_handler = cancellation_handler.clone(); - let endpoint_rate_limiter = endpoint_rate_limiter.clone(); - connections.spawn( - async move { - let conn_token2 = conn_token.clone(); - let _cancel_guard = config.http_config.cancel_set.insert(conn_id, conn_token2); + let conn_token = self.cancellation_token.child_token(); + let _cancel_guard = self + .config + .http_config + .cancel_set + .insert(conn_id, conn_token.clone()); - let session_id = uuid::Uuid::new_v4(); + let startup_result = Box::pin(connection_startup( + self.config, + self.tls_acceptor, + conn_id, + stream, + peer_addr, + )) + .await; + let Some((conn, peer_addr)) = startup_result else { + return; + }; - let _gauge = Metrics::get() - .proxy - .client_connections - .guard(crate::metrics::Protocol::Http); + Box::pin(connection_handler( + self.config, + self.backend, + self.requests, + self.cancellation_handler, + self.endpoint_rate_limiter, + conn_token, + conn, + peer_addr, + conn_id, + )) + .await; - let startup_result = Box::pin(connection_startup( - config, - tls_acceptor, - session_id, - conn, - peer_addr, - )) - .await; - let Some((conn, peer_addr)) = startup_result else { - return; - }; - - Box::pin(connection_handler( - config, - backend, - connections2, - cancellation_handler, - endpoint_rate_limiter, - conn_token, - conn, - peer_addr, - session_id, - )) - .await; - } - .instrument(http_conn_span), - ); + drop(conn_gauge); } - - connections.wait().await; - - Ok(()) } pub(crate) trait AsyncReadWrite: AsyncRead + AsyncWrite + Send + 'static {} @@ -227,26 +244,14 @@ impl MaybeTlsAcceptor for NoTls { } } -/// Handles the TCP startup lifecycle. -/// 1. Parses PROXY protocol V2 -/// 2. Handles TLS handshake +/// Handles the TLS startup handshake. async fn connection_startup( config: &ProxyConfig, tls_acceptor: Arc, session_id: uuid::Uuid, - conn: TcpStream, - peer_addr: SocketAddr, + conn: ChainRW, + peer_addr: IpAddr, ) -> Option<(AsyncRW, IpAddr)> { - // handle PROXY protocol - let (conn, peer) = match read_proxy_protocol(conn).await { - Ok(c) => c, - Err(e) => { - tracing::error!(?session_id, %peer_addr, "failed to accept TCP connection: invalid PROXY protocol V2 header: {e:#}"); - return None; - } - }; - - let peer_addr = peer.unwrap_or(peer_addr).ip(); let has_private_peer_addr = match peer_addr { IpAddr::V4(ip) => ip.is_private(), IpAddr::V6(_) => false,