mz_balancerd/

lib.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! The balancerd service is a horizontally scalable, stateless, multi-tenant ingress router for
//! pgwire and HTTPS connections.
//!
//! It listens on pgwire and HTTPS ports. When a new pgwire connection starts, the requested user is
//! authenticated with frontegg from which a tenant id is returned. From that a target internal
//! hostname is resolved to an IP address, and the connection is proxied to that address which has a
//! running environmentd's pgwire port. When a new HTTPS connection starts, its SNI hostname is used
//! to generate an internal hostname that is resolved to an IP address, which is similarly proxied.

mod codec;
mod dyncfgs;

use std::collections::BTreeMap;
use std::net::{IpAddr, SocketAddr};
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::time::{Duration, Instant};

use anyhow::Context;
use axum::response::IntoResponse;
use axum::{Router, routing};
use bytes::BytesMut;
use futures::TryFutureExt;
use futures::stream::BoxStream;
use hickory_resolver::config::LookupIpStrategy;
use hickory_resolver::lookup_ip::LookupIp;
use hickory_resolver::net::runtime::TokioRuntimeProvider;
use hickory_resolver::proto::rr::{RData, RecordType};
use hickory_resolver::system_conf::read_system_conf;
use hickory_resolver::{Resolver, TokioResolver};
use hyper::StatusCode;
use hyper_util::rt::TokioIo;
use launchdarkly_server_sdk as ld;
use mz_build_info::{BuildInfo, build_info};
use mz_dyncfg::ConfigSet;
use mz_frontegg_auth::Authenticator as FronteggAuthentication;
use mz_ore::cast::CastFrom;
use mz_ore::id_gen::conn_id_org_uuid;
use mz_ore::metrics::{ComputedGauge, ComputedUIntGauge, IntCounter, IntGauge, MetricsRegistry};
use mz_ore::netio::AsyncReady;
use mz_ore::now::{NowFn, SYSTEM_TIME, epoch_to_uuid_v7};
use mz_ore::task::{JoinSetExt, spawn};
use mz_ore::tracing::TracingHandle;
use mz_ore::{metric, netio};
use mz_pgwire_common::{
    ACCEPT_SSL_ENCRYPTION, CONN_UUID_KEY, Conn, ErrorResponse, FrontendMessage,
    FrontendStartupMessage, MAX_STARTUP_FRAME_SIZE, MZ_FORWARDED_FOR_KEY, REJECT_ENCRYPTION,
    VERSION_3, decode_startup,
};
use mz_server_core::{
    Connection, ConnectionStream, ListenerHandle, ReloadTrigger, ReloadingSslContext,
    ReloadingTlsConfig, ServeConfig, ServeDyncfg, TlsCertConfig, TlsMode, listen,
};
use openssl::error::ErrorStack;
use openssl::ssl::{NameType, Ssl, SslConnector, SslMethod, SslVerifyMode};
use prometheus::{IntCounterVec, IntGaugeVec};
use proxy_header::{ProxiedAddress, ProxyHeader};
use semver::Version;
use tokio::io::{self, AsyncRead, AsyncWrite, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::oneshot;
use tokio::task::JoinSet;
use tokio_metrics::TaskMetrics;
use tokio_openssl::SslStream;
use tokio_postgres::error::SqlState;
use tower::Service;
use tracing::{debug, error, info, warn};
use uuid::Uuid;

use crate::codec::{BackendMessage, FramedConn};
use crate::dyncfgs::{
    INJECT_PROXY_PROTOCOL_HEADER_HTTP, MAX_CONNECTIONS, PRE_RESOLVED_TIMEOUT,
    SIGTERM_CONNECTION_WAIT, SIGTERM_LISTEN_WAIT, has_tracing_config_update, tracing_config,
};

/// Balancer build information.
pub const BUILD_INFO: BuildInfo = build_info!();

pub struct BalancerConfig {
    /// Info about which version of the code is running.
    build_version: Version,
    /// Listen address for internal HTTP health and metrics server.
    internal_http_listen_addr: SocketAddr,
    /// Listen address for pgwire connections.
    pgwire_listen_addr: SocketAddr,
    /// Listen address for HTTPS connections.
    https_listen_addr: SocketAddr,
    /// DNS resolver for pgwire cancellation requests
    cancellation_resolver: CancellationResolver,
    /// DNS resolver.
    resolver: BalancerResolver,
    https_sni_addr_template: String,
    tls: Option<TlsCertConfig>,
    internal_tls: bool,
    metrics_registry: MetricsRegistry,
    reload_certs: BoxStream<'static, Option<oneshot::Sender<Result<(), anyhow::Error>>>>,
    launchdarkly_sdk_key: Option<String>,
    config_sync_file_path: Option<PathBuf>,
    config_sync_timeout: Duration,
    config_sync_loop_interval: Option<Duration>,
    cloud_provider: Option<String>,
    cloud_provider_region: Option<String>,
    tracing_handle: TracingHandle,
    default_configs: Vec<(String, String)>,
}

impl BalancerConfig {
    pub fn new(
        build_info: &BuildInfo,
        internal_http_listen_addr: SocketAddr,
        pgwire_listen_addr: SocketAddr,
        https_listen_addr: SocketAddr,
        cancellation_resolver: CancellationResolver,
        resolver: BalancerResolver,
        https_sni_addr_template: String,
        tls: Option<TlsCertConfig>,
        internal_tls: bool,
        metrics_registry: MetricsRegistry,
        reload_certs: ReloadTrigger,
        launchdarkly_sdk_key: Option<String>,
        config_sync_file: Option<PathBuf>,
        config_sync_timeout: Duration,
        config_sync_loop_interval: Option<Duration>,
        cloud_provider: Option<String>,
        cloud_provider_region: Option<String>,
        tracing_handle: TracingHandle,
        default_configs: Vec<(String, String)>,
    ) -> Self {
        Self {
            build_version: build_info.semver_version(),
            internal_http_listen_addr,
            pgwire_listen_addr,
            https_listen_addr,
            cancellation_resolver,
            resolver,
            https_sni_addr_template,
            tls,
            internal_tls,
            metrics_registry,
            reload_certs,
            launchdarkly_sdk_key,
            config_sync_file_path: config_sync_file,
            config_sync_timeout,
            config_sync_loop_interval,
            cloud_provider,
            cloud_provider_region,
            tracing_handle,
            default_configs,
        }
    }
}

/// Prometheus monitoring metrics.
#[derive(Debug)]
pub struct BalancerMetrics {
    _uptime: ComputedGauge,
}

impl BalancerMetrics {
    /// Returns a new [BalancerMetrics] instance connected to the registry in cfg.
    pub fn new(cfg: &BalancerConfig) -> Self {
        let start = Instant::now();
        let uptime = cfg.metrics_registry.register_computed_gauge(
            metric!(
                name: "mz_balancer_metadata_seconds",
                help: "server uptime, labels are build metadata",
                const_labels: {
                    "version" => cfg.build_version,
                    "build_type" => if cfg!(release) { "release" } else { "debug" }
                },
            ),
            move || start.elapsed().as_secs_f64(),
        );
        BalancerMetrics { _uptime: uptime }
    }
}

pub struct BalancerService {
    cfg: BalancerConfig,
    pub pgwire: (ListenerHandle, Pin<Box<dyn ConnectionStream>>),
    pub https: (ListenerHandle, Pin<Box<dyn ConnectionStream>>),
    pub internal_http: (ListenerHandle, Pin<Box<dyn ConnectionStream>>),
    _metrics: BalancerMetrics,
    configs: ConfigSet,
}

impl BalancerService {
    pub async fn new(cfg: BalancerConfig) -> Result<Self, anyhow::Error> {
        let pgwire = listen(&cfg.pgwire_listen_addr).await?;
        let https = listen(&cfg.https_listen_addr).await?;
        let internal_http = listen(&cfg.internal_http_listen_addr).await?;
        let metrics = BalancerMetrics::new(&cfg);
        let mut configs = ConfigSet::default();
        configs = dyncfgs::all_dyncfgs(configs);
        dyncfgs::set_defaults(&configs, cfg.default_configs.clone())?;
        let tracing_handle = cfg.tracing_handle.clone();
        // Configure dyncfg sync
        match (
            cfg.launchdarkly_sdk_key.as_deref(),
            cfg.config_sync_file_path.as_deref(),
        ) {
            (Some(key), None) => {
                let _ = mz_dyncfg_launchdarkly::sync_launchdarkly_to_configset(
                    configs.clone(),
                    &BUILD_INFO,
                    |builder| {
                        let region = cfg
                            .cloud_provider_region
                            .clone()
                            .unwrap_or_else(|| String::from("unknown"));
                        if let Some(provider) = cfg.cloud_provider.clone() {
                            builder.add_context(
                                ld::ContextBuilder::new(format!(
                                    "{}/{}/{}",
                                    provider, region, cfg.build_version
                                ))
                                .kind("balancer")
                                .set_string("provider", provider)
                                .set_string("region", region)
                                .set_string("version", cfg.build_version.to_string())
                                .build()
                                .map_err(|e| anyhow::anyhow!(e))?,
                            );
                        } else {
                            builder.add_context(
                                ld::ContextBuilder::new(format!(
                                    "{}/{}/{}",
                                    "unknown", region, cfg.build_version
                                ))
                                .anonymous(true) // exclude this user from the dashboard
                                .kind("balancer")
                                .set_string("provider", "unknown")
                                .set_string("region", region)
                                .set_string("version", cfg.build_version.to_string())
                                .build()
                                .map_err(|e| anyhow::anyhow!(e))?,
                            );
                        }
                        Ok(())
                    },
                    Some(key),
                    cfg.config_sync_timeout,
                    cfg.config_sync_loop_interval,
                    move |updates, configs| {
                        if has_tracing_config_update(updates) {
                            match tracing_config(configs) {
                                Ok(parameters) => parameters.apply(&tracing_handle),
                                Err(err) => warn!("unable to update tracing: {err}"),
                            }
                        }
                    },
                )
                .await
                .inspect_err(|e| warn!("LaunchDarkly sync error: {e}"));
            }
            (None, Some(path)) => {
                let _ = mz_dyncfg_file::sync_file_to_configset(
                    configs.clone(),
                    path,
                    cfg.config_sync_timeout,
                    cfg.config_sync_loop_interval,
                    move |updates, configs| {
                        if has_tracing_config_update(updates) {
                            match tracing_config(configs) {
                                Ok(parameters) => parameters.apply(&tracing_handle),
                                Err(err) => warn!("unable to update tracing: {err}"),
                            }
                        }
                    },
                )
                .await
                // If there's an Error, log but continue anyway. If LD is down
                // we have no way of fetching the previous value of the flag
                // (unlike the adapter, but it has a durable catalog). The
                // ConfigSet defaults have been chosen to be good enough if this
                // is the case.
                .inspect_err(|e| warn!("File config sync error: {e}"));
            }
            (Some(_), Some(_)) => panic!(
                "must provide either config_sync_file_path or launchdarkly_sdk_key for config syncing",
            ),
            (None, None) => {}
        };
        Ok(Self {
            cfg,
            pgwire,
            https,
            internal_http,
            _metrics: metrics,
            configs,
        })
    }

    pub async fn serve(self) -> Result<(), anyhow::Error> {
        let (pgwire_tls, https_tls) = match &self.cfg.tls {
            Some(tls) => {
                let context = tls.reloading_context(self.cfg.reload_certs)?;
                (
                    Some(ReloadingTlsConfig {
                        context: context.clone(),
                        mode: TlsMode::Require,
                    }),
                    Some(context),
                )
            }
            None => (None, None),
        };

        // Built once and shared by every upstream connection, since building a connector loads
        // and parses the system CA bundle.
        let internal_tls = self
            .cfg
            .internal_tls
            .then(internal_tls_connector)
            .transpose()
            .context("building internal TLS connector")?;

        let metrics = ServerMetricsConfig::register_into(&self.cfg.metrics_registry);
        let limiter = ConnectionLimiter::new(&self.cfg.metrics_registry, self.configs.clone());

        let mut set = JoinSet::new();
        let mut server_handles = Vec::new();
        let pgwire_addr = self.pgwire.0.local_addr();
        let https_addr = self.https.0.local_addr();
        let internal_http_addr = self.internal_http.0.local_addr();

        // The HTTPS balancer always resolves through a TenantDnsResolver. In
        // multi-tenant mode it shares the pgwire resolver so both listeners use
        // one resolver. In static mode pgwire does not resolve through it, so
        // HTTPS gets its own.
        let shared_dns = match self.cfg.resolver.shared_dns() {
            Some(dns) => dns,
            None => Arc::new(TenantDnsResolver::new()?),
        };

        {
            let pgwire = PgwireBalancer {
                resolver: Arc::new(self.cfg.resolver),
                cancellation_resolver: Arc::new(self.cfg.cancellation_resolver),
                tls: pgwire_tls,
                internal_tls: internal_tls.clone(),
                metrics: ServerMetrics::new(metrics.clone(), "pgwire"),
                limiter: Arc::clone(&limiter),
                configs: self.configs.clone(),
                now: SYSTEM_TIME.clone(),
            };
            let (handle, stream) = self.pgwire;
            server_handles.push(handle);
            set.spawn_named(|| "pgwire_stream", {
                let config_set = self.configs.clone();
                async move {
                    mz_server_core::serve(ServeConfig {
                        server: pgwire,
                        conns: stream,
                        dyncfg: Some(ServeDyncfg {
                            config_set,
                            sigterm_wait_config: &SIGTERM_CONNECTION_WAIT,
                        }),
                    })
                    .await;
                    warn!("pgwire server exited");
                }
            });
        }
        {
            let Some((addr, port)) = self.cfg.https_sni_addr_template.split_once(':') else {
                panic!("expected port in https_addr_template");
            };
            let port: u16 = port.parse().expect("unexpected port");

            let https = HttpsBalancer {
                resolver: shared_dns,
                tls: https_tls,
                resolve_template: Arc::from(addr),
                port,
                metrics: Arc::from(ServerMetrics::new(metrics, "https")),
                limiter,
                configs: self.configs.clone(),
                internal_tls,
            };
            let (handle, stream) = self.https;
            server_handles.push(handle);
            set.spawn_named(|| "https_stream", {
                let config_set = self.configs.clone();
                async move {
                    mz_server_core::serve(ServeConfig {
                        server: https,
                        conns: stream,
                        dyncfg: Some(ServeDyncfg {
                            config_set,
                            sigterm_wait_config: &SIGTERM_CONNECTION_WAIT,
                        }),
                    })
                    .await;
                    warn!("https server exited");
                }
            });
        }
        {
            let router = Router::new()
                .route(
                    "/metrics",
                    routing::get(move |headers: axum::http::HeaderMap| async move {
                        mz_http_util::handle_prometheus(&self.cfg.metrics_registry, headers).await
                    }),
                )
                .route(
                    "/api/livez",
                    routing::get(mz_http_util::handle_liveness_check),
                )
                .route("/api/readyz", routing::get(handle_readiness_check));
            let internal_http = InternalHttpServer { router };
            let (handle, stream) = self.internal_http;
            server_handles.push(handle);
            set.spawn_named(|| "internal_http_stream", async move {
                mz_server_core::serve(ServeConfig {
                    server: internal_http,
                    conns: stream,
                    // Disable graceful termination because our internal
                    // monitoring keeps persistent HTTP connections open.
                    dyncfg: None,
                })
                .await;
                warn!("internal_http server exited");
            });
        }
        #[cfg(unix)]
        {
            let mut sigterm =
                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
            set.spawn_named(|| "sigterm_handler", async move {
                sigterm.recv().await;
                let wait = SIGTERM_LISTEN_WAIT.get(&self.configs);
                warn!("received signal TERM - delaying for {:?}!", wait);
                tokio::time::sleep(wait).await;
                warn!("sigterm delay complete, dropping server handles");
                drop(server_handles);
            });
        }

        println!("balancerd {} listening...", BUILD_INFO.human_version(None));
        println!(" TLS enabled: {}", self.cfg.tls.is_some());
        println!(" pgwire address: {}", pgwire_addr);
        println!(" HTTPS address: {}", https_addr);
        println!(" internal HTTP address: {}", internal_http_addr);

        // Wait for all tasks to exit, which can happen on SIGTERM.
        while let Some(res) = set.join_next().await {
            if let Err(err) = res {
                error!("serving task failed: {err}")
            }
        }
        Ok(())
    }
}

#[allow(clippy::unused_async)]
async fn handle_readiness_check() -> impl IntoResponse {
    (StatusCode::OK, "ready")
}

struct InternalHttpServer {
    router: Router,
}

impl mz_server_core::Server for InternalHttpServer {
    const NAME: &'static str = "internal_http";

    // TODO(jkosh44) consider forwarding the connection UUID to the adapter.
    fn handle_connection(
        &self,
        conn: Connection,
        _tokio_metrics_intervals: impl Iterator<Item = TaskMetrics> + Send + 'static,
    ) -> mz_server_core::ConnectionHandler {
        let router = self.router.clone();
        let service = hyper::service::service_fn(move |req| router.clone().call(req));
        let conn = TokioIo::new(conn);

        Box::pin(async {
            let http = hyper::server::conn::http1::Builder::new();
            http.serve_connection(conn, service).err_into().await
        })
    }
}

/// Builds the connector for TLS to environmentd.
fn internal_tls_connector() -> Result<SslConnector, ErrorStack> {
    let mut builder = SslConnector::builder(SslMethod::tls())?;
    // environmentd doesn't yet have a cert we trust, so for now disable verification.
    builder.set_verify(SslVerifyMode::NONE);
    Ok(builder.build())
}

/// Wraps an IntGauge and automatically `inc`s on init and `drop`s on drop. Callers should not call
/// `inc().`. Useful for handling multiple task exit points, for example in the case of a panic.
struct GaugeGuard {
    gauge: IntGauge,
}

impl From<IntGauge> for GaugeGuard {
    fn from(gauge: IntGauge) -> Self {
        let _self = Self { gauge };
        _self.gauge.inc();
        _self
    }
}

impl Drop for GaugeGuard {
    fn drop(&mut self) {
        self.gauge.dec();
    }
}

#[derive(Clone, Debug)]
struct ServerMetricsConfig {
    connection_status: IntCounterVec,
    active_connections: IntGaugeVec,
    pre_resolved_connections: IntGaugeVec,
    pre_resolved_timeouts: IntCounterVec,
    tenant_connections: IntGaugeVec,
    tenant_connection_rx: IntCounterVec,
    tenant_connection_tx: IntCounterVec,
    tenant_pgwire_sni_count: IntCounterVec,
}

impl ServerMetricsConfig {
    fn register_into(registry: &MetricsRegistry) -> Self {
        let connection_status = registry.register(metric!(
            name: "mz_balancer_connection_status",
            help: "Count of completed network connections, by status",
            var_labels: ["source", "status"],
        ));
        let active_connections = registry.register(metric!(
            name: "mz_balancer_connection_active",
            help: "Count of currently open network connections.",
            var_labels: ["source"],
        ));
        let pre_resolved_connections = registry.register(metric!(
            name: "mz_balancer_pre_resolved_connection_active",
            help: "Count of open network connections that have not yet resolved a backend.",
            var_labels: ["source"],
        ));
        let pre_resolved_timeouts = registry.register(metric!(
            name: "mz_balancer_pre_resolved_timeout_total",
            help: "Count of connections closed for not resolving a backend within the pre-resolved timeout.",
            var_labels: ["source"],
        ));
        let tenant_connections = registry.register(metric!(
            name: "mz_balancer_tenant_connection_active",
            help: "Count of opened network connections by tenant.",
            var_labels: ["source",  "tenant"]
        ));
        let tenant_connection_rx = registry.register(metric!(
            name: "mz_balancer_tenant_connection_rx",
            help: "Number of bytes received from a client for a tenant.",
            var_labels: ["source", "tenant"],
        ));
        let tenant_connection_tx = registry.register(metric!(
            name: "mz_balancer_tenant_connection_tx",
            help: "Number of bytes sent to a client for a tenant.",
            var_labels: ["source", "tenant"],
        ));
        let tenant_pgwire_sni_count = registry.register(metric!(
            name: "mz_balancer_tenant_pgwire_sni_count",
            help: "Count of pgwire connections that have and do not have SNI available per tenant.",
            var_labels: ["tenant", "has_sni"],
        ));
        Self {
            connection_status,
            active_connections,
            pre_resolved_connections,
            pre_resolved_timeouts,
            tenant_connections,
            tenant_connection_rx,
            tenant_connection_tx,
            tenant_pgwire_sni_count,
        }
    }
}

#[derive(Clone, Debug)]
struct ServerMetrics {
    inner: ServerMetricsConfig,
    source: &'static str,
}

impl ServerMetrics {
    fn new(inner: ServerMetricsConfig, source: &'static str) -> Self {
        let self_ = Self { inner, source };

        // Pre-initialize labels we are planning to use to ensure they are all always emitted as
        // time series.
        self_.connection_status(false);
        self_.connection_status(true);
        drop(self_.active_connections());

        self_
    }

    fn connection_status(&self, is_ok: bool) -> IntCounter {
        self.inner
            .connection_status
            .with_label_values(&[self.source, Self::status_label(is_ok)])
    }

    fn active_connections(&self) -> GaugeGuard {
        self.inner
            .active_connections
            .with_label_values(&[self.source])
            .into()
    }

    fn pre_resolved_connections(&self) -> GaugeGuard {
        self.inner
            .pre_resolved_connections
            .with_label_values(&[self.source])
            .into()
    }

    fn pre_resolved_timeout(&self) -> IntCounter {
        self.inner
            .pre_resolved_timeouts
            .with_label_values(&[self.source])
    }

    fn tenant_connections(&self, tenant: &str) -> GaugeGuard {
        self.inner
            .tenant_connections
            .with_label_values(&[self.source, tenant])
            .into()
    }

    fn tenant_connections_rx(&self, tenant: &str) -> IntCounter {
        self.inner
            .tenant_connection_rx
            .with_label_values(&[self.source, tenant])
    }

    fn tenant_connections_tx(&self, tenant: &str) -> IntCounter {
        self.inner
            .tenant_connection_tx
            .with_label_values(&[self.source, tenant])
    }

    fn tenant_pgwire_sni_count(&self, tenant: &str, has_sni: bool) -> IntCounter {
        self.inner
            .tenant_pgwire_sni_count
            .with_label_values(&[tenant, &has_sni.to_string()])
    }

    fn status_label(is_ok: bool) -> &'static str {
        if is_ok { "success" } else { "error" }
    }
}

/// Ceiling on the number of client connections proxied at once, shared by the pgwire and HTTPS
/// listeners.
///
/// balancerd's memory use scales with the number of connections it proxies. Without a ceiling it
/// keeps accepting until the container is OOM killed, which drops every established connection.
/// Refusing new connections instead sheds only the load we cannot serve.
#[derive(Debug)]
struct ConnectionLimiter {
    configs: ConfigSet,
    active: AtomicU32,
    /// Whether connections are currently being refused, so that reaching and clearing the limit
    /// are logged once each rather than once per connection.
    limited: AtomicBool,
    rejected: IntCounter,
    _limit: ComputedUIntGauge,
}

impl ConnectionLimiter {
    fn new(registry: &MetricsRegistry, configs: ConfigSet) -> Arc<Self> {
        let rejected = registry.register(metric!(
            name: "mz_balancer_connection_rejected_total",
            help: "Count of connections refused because the connection limit was reached.",
        ));
        let limit = registry.register_computed_gauge(
            metric!(
                name: "mz_balancer_connection_limit",
                help: "Maximum number of connections proxied at once, 0 if unlimited.",
            ),
            {
                let configs = configs.clone();
                move || u64::from(MAX_CONNECTIONS.get(&configs))
            },
        );
        Arc::new(ConnectionLimiter {
            configs,
            active: AtomicU32::new(0),
            limited: AtomicBool::new(false),
            rejected,
            _limit: limit,
        })
    }

    /// Reserves capacity for one connection, or returns `None` if the limit has been reached, in
    /// which case the caller must refuse the connection.
    fn acquire(self: &Arc<Self>) -> Option<ConnectionGuard> {
        let limit = MAX_CONNECTIONS.get(&self.configs);
        let unlimited = limit == 0;
        match self
            .active
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |active| {
                (unlimited || active < limit).then(|| active + 1)
            }) {
            Ok(prev) => {
                // Only clear the limited state once comfortably below the limit, so that
                // connections churning at the ceiling do not flap the log lines.
                if unlimited || prev + 1 <= limit - limit / 10 {
                    if self.limited.swap(false, Ordering::Relaxed) {
                        info!("accepting new connections again, limit is {limit}");
                    }
                }
                Some(ConnectionGuard(Arc::clone(self)))
            }
            Err(active) => {
                self.rejected.inc();
                if !self.limited.swap(true, Ordering::Relaxed) {
                    warn!(
                        "refusing new connections: at the limit of {limit} connections \
                        ({active} active)"
                    );
                }
                None
            }
        }
    }
}

/// Releases the connection reserved by [`ConnectionLimiter::acquire`] when dropped.
struct ConnectionGuard(Arc<ConnectionLimiter>);

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.0.active.fetch_sub(1, Ordering::SeqCst);
    }
}

/// Answers the client with a final message and flushes it.
///
/// [`FramedConn::send`] only enqueues into a buffer, so a rejection that is not flushed reaches
/// the client as a bare socket close rather than the error it was given. Every path that answers
/// the client and then stops has to flush, because the connection is dropped as soon as it
/// returns.
async fn reject<A>(conn: &mut FramedConn<A>, err: ErrorResponse) -> Result<(), io::Error>
where
    A: AsyncRead + AsyncWrite + Unpin,
{
    conn.send(err).await?;
    conn.flush().await
}

/// Runs the pre-resolved phase under [`PRE_RESOLVED_TIMEOUT`], counting an overrun.
///
/// `fut` must cover TLS negotiation, startup, authentication, and backend resolution.
/// Proxying must run after it completes, outside the deadline. A zero timeout leaves
/// `fut` unbounded.
async fn under_pre_resolved_timeout<F, T, E>(
    timeout: Duration,
    metrics: &ServerMetrics,
    fut: F,
) -> Result<T, E>
where
    F: Future<Output = Result<T, E>>,
    E: From<io::Error>,
{
    if timeout.is_zero() {
        return fut.await;
    }
    match tokio::time::timeout(timeout, fut).await {
        Ok(result) => result,
        Err(_) => {
            metrics.pre_resolved_timeout().inc();
            Err(io::Error::new(
                io::ErrorKind::TimedOut,
                "timed out before resolving a backend",
            )
            .into())
        }
    }
}

pub enum CancellationResolver {
    Directory(PathBuf),
    Static(String),
}

struct PgwireBalancer {
    /// Read for [`PRE_RESOLVED_TIMEOUT`] at connection time.
    configs: ConfigSet,
    tls: Option<ReloadingTlsConfig>,
    internal_tls: Option<SslConnector>,
    cancellation_resolver: Arc<CancellationResolver>,
    resolver: Arc<BalancerResolver>,
    metrics: ServerMetrics,
    limiter: Arc<ConnectionLimiter>,
    now: NowFn,
}

impl PgwireBalancer {
    /// Negotiates pgwire startup and resolves a backend.
    ///
    /// Returns `None` after a client disconnect, a flushed rejection, or a dispatched
    /// cancellation request.
    async fn pre_resolve(
        conn: Connection,
        tls: Option<ReloadingTlsConfig>,
        resolver: &BalancerResolver,
        cancellation_resolver: Arc<CancellationResolver>,
        conn_uuid: Uuid,
        metrics: &ServerMetrics,
    ) -> Result<
        Option<(
            FramedConn<Connection>,
            ResolvedAddr,
            BTreeMap<String, String>,
        )>,
        anyhow::Error,
    > {
        let peer_addr = conn.peer_addr();
        let mut conn = Conn::Unencrypted(conn);
        loop {
            let message = decode_startup(&mut conn, MAX_STARTUP_FRAME_SIZE).await?;
            conn = match message {
                // Clients sometimes hang up during the startup sequence, e.g.
                // because they receive an unacceptable response to an
                // `SslRequest`. This is considered a graceful termination.
                None => return Ok(None),

                Some(FrontendStartupMessage::Startup {
                    version,
                    mut params,
                }) => {
                    let mut conn = FramedConn::new(conn);
                    let rejected = SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION;
                    let peer_addr = match peer_addr {
                        Ok(addr) => addr.ip(),
                        Err(e) => {
                            error!("Invalid peer_addr {:?}", e);
                            reject(
                                &mut conn,
                                ErrorResponse::fatal(rejected, "invalid peer address"),
                            )
                            .await?;
                            return Ok(None);
                        }
                    };
                    debug!(
                        %conn_uuid, %peer_addr,
                        "starting new pgwire connection in balancer",
                    );
                    let prev = params.insert(CONN_UUID_KEY.to_string(), conn_uuid.to_string());
                    if prev.is_some() {
                        reject(
                            &mut conn,
                            ErrorResponse::fatal(
                                rejected,
                                format!("invalid parameter '{CONN_UUID_KEY}'"),
                            ),
                        )
                        .await?;
                        return Ok(None);
                    }

                    let forwarded_for = params.insert(
                        MZ_FORWARDED_FOR_KEY.to_string(),
                        peer_addr.to_string().clone(),
                    );
                    if let Some(_) = forwarded_for {
                        reject(
                            &mut conn,
                            ErrorResponse::fatal(
                                rejected,
                                format!("invalid parameter '{MZ_FORWARDED_FOR_KEY}'"),
                            ),
                        )
                        .await?;
                        return Ok(None);
                    };

                    let Some(resolved) = Self::resolve_destination(
                        &mut conn,
                        version,
                        &params,
                        resolver,
                        tls.map(|tls| tls.mode),
                        metrics,
                    )
                    .await?
                    else {
                        return Ok(None);
                    };
                    return Ok(Some((conn, resolved, params)));
                }

                Some(FrontendStartupMessage::CancelRequest {
                    conn_id,
                    secret_key,
                }) => {
                    spawn(|| "cancel request", async move {
                        cancel_request(conn_id, secret_key, &cancellation_resolver).await;
                    });
                    // Do not wait on cancel requests to return because cancellation is best
                    // effort.
                    return Ok(None);
                }

                Some(FrontendStartupMessage::SslRequest) => match (conn, &tls) {
                    (Conn::Unencrypted(mut conn), Some(tls)) => {
                        conn.write_all(&[ACCEPT_SSL_ENCRYPTION]).await?;
                        let mut ssl_stream = SslStream::new(Ssl::new(&tls.context.get())?, conn)?;
                        if let Err(e) = Pin::new(&mut ssl_stream).accept().await {
                            let _ = ssl_stream.get_mut().shutdown().await;
                            return Err(e.into());
                        }
                        Conn::Ssl(ssl_stream)
                    }
                    (mut conn, _) => {
                        conn.write_all(&[REJECT_ENCRYPTION]).await?;
                        conn
                    }
                },

                Some(FrontendStartupMessage::GssEncRequest) => {
                    conn.write_all(&[REJECT_ENCRYPTION]).await?;
                    conn
                }
            }
        }
    }

    #[mz_ore::instrument(level = "debug")]
    /// Resolves a destination for the connection, or answers the client saying why it cannot.
    ///
    /// `Ok(None)` means the rejection has been flushed and the connection is finished.
    async fn resolve_destination<'a, A>(
        conn: &'a mut FramedConn<A>,
        version: i32,
        params: &BTreeMap<String, String>,
        resolver: &BalancerResolver,
        tls_mode: Option<TlsMode>,
        metrics: &ServerMetrics,
    ) -> Result<Option<ResolvedAddr>, io::Error>
    where
        A: AsyncRead + AsyncWrite + AsyncReady + Send + Sync + Unpin,
    {
        if version != VERSION_3 {
            reject(
                conn,
                ErrorResponse::fatal(
                    SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION,
                    "server does not support the client's requested protocol version",
                ),
            )
            .await?;
            return Ok(None);
        }

        let Some(user) = params.get("user") else {
            reject(
                conn,
                ErrorResponse::fatal(
                    SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION,
                    "user parameter required",
                ),
            )
            .await?;
            return Ok(None);
        };

        if let Err(err) = conn.inner().ensure_tls_compatibility(&tls_mode) {
            reject(conn, err).await?;
            return Ok(None);
        }

        let resolved = match resolver.resolve(conn, user, metrics).await {
            Ok(v) => v,
            Err(err) => {
                let sql_state = match &err {
                    ResolveError::InvalidPassword => SqlState::INVALID_PASSWORD,
                    ResolveError::Client(details) => {
                        warn!("client-caused connection failure: {details:#}");
                        SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION
                    }
                    ResolveError::Upstream(details) => {
                        warn!("upstream not available: {details:#}");
                        SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION
                    }
                    ResolveError::Internal(details) => {
                        error!("resolving connection destination: {details:#}");
                        SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION
                    }
                };
                reject(conn, ErrorResponse::fatal(sql_state, err.to_string())).await?;
                return Ok(None);
            }
        };

        Ok(Some(resolved))
    }

    /// Proxies a resolved connection until either side closes.
    async fn proxy<'a, A>(
        conn: &'a mut FramedConn<A>,
        resolved: ResolvedAddr,
        params: BTreeMap<String, String>,
        internal_tls: Option<&SslConnector>,
        metrics: &ServerMetrics,
    ) -> Result<(), io::Error>
    where
        A: AsyncRead + AsyncWrite + AsyncReady + Send + Sync + Unpin,
    {
        let _active_guard = resolved
            .tenant
            .as_ref()
            .map(|tenant| metrics.tenant_connections(tenant));
        let mut mz_stream =
            match Self::init_stream(conn, resolved.addr, resolved.password, params, internal_tls)
                .await
            {
                Ok(stream) => stream,
                Err(e) => {
                    error!("failed to connect to upstream server: {e}");
                    return conn
                        .send(ErrorResponse::fatal(
                            SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION,
                            "upstream server not available",
                        ))
                        .await;
                }
            };

        let mut client_counter = CountingConn::new(conn.inner_mut());

        // Now blindly shuffle bytes back and forth until closed.
        // TODO: Limit total memory use.
        let res = tokio::io::copy_bidirectional(&mut client_counter, &mut mz_stream).await;
        if let Some(tenant) = &resolved.tenant {
            metrics
                .tenant_connections_tx(tenant)
                .inc_by(u64::cast_from(client_counter.written));
            metrics
                .tenant_connections_rx(tenant)
                .inc_by(u64::cast_from(client_counter.read));
        }
        res?;

        Ok(())
    }

    #[mz_ore::instrument(level = "debug")]
    async fn init_stream<'a, A>(
        conn: &'a mut FramedConn<A>,
        envd_addr: SocketAddr,
        password: Option<String>,
        params: BTreeMap<String, String>,
        internal_tls: Option<&SslConnector>,
    ) -> Result<Conn<TcpStream>, anyhow::Error>
    where
        A: AsyncRead + AsyncWrite + AsyncReady + Send + Sync + Unpin,
    {
        let mut mz_stream = TcpStream::connect(envd_addr).await?;
        let mut buf = BytesMut::new();

        let mut mz_stream = if let Some(internal_tls) = internal_tls {
            FrontendStartupMessage::SslRequest.encode(&mut buf)?;
            mz_stream.write_all(&buf).await?;
            buf.clear();
            let mut maybe_ssl_request_response = [0u8; 1];
            let nread =
                netio::read_exact_or_eof(&mut mz_stream, &mut maybe_ssl_request_response).await?;
            if nread == 1 && maybe_ssl_request_response == [ACCEPT_SSL_ENCRYPTION] {
                // do a TLS handshake
                let mut ssl = internal_tls.configure()?.into_ssl(&envd_addr.to_string())?;
                ssl.set_connect_state();
                Conn::Ssl(SslStream::new(ssl, mz_stream)?)
            } else {
                Conn::Unencrypted(mz_stream)
            }
        } else {
            Conn::Unencrypted(mz_stream)
        };

        // Send initial startup and password messages.
        let startup = FrontendStartupMessage::Startup {
            version: VERSION_3,
            params,
        };
        startup.encode(&mut buf)?;
        mz_stream.write_all(&buf).await?;
        let client_stream = conn.inner_mut();

        // This early return is important in self managed with SASL mode.
        // The below code specifically looks for cleartext password requests, but in SASL mode
        // the server will send a different message type (SASLInitialResponse) that we should
        // not try to interpret or respond to.
        // "Why not? That code looks like it should fall back fine?" You may ask.
        // The below block unconditionally reads 9 bytes from the server. If we don't have
        // a password or the message isn't a cleartext password request, we forward those 9 bytes
        // to the client. Then we return the stream to the caller, who will continue shuffling bytes.
        // The problem is that with TLS enabled between balancerd <-> client, flushing the first 9 bytes
        // before copying bidirectionally will have the side effect of splitting the auth handshake into
        // two SSL records. Pgbouncer misbehaves in this scenario, and fails the connection.
        // PGbouncer shouldn't do this! It's a common footgun of protocols over TLS.
        // So common in fact that PGbouncer already hit and fixed this issue on the bouncer <-> client side:
        // once before: https://github.com/pgbouncer/pgbouncer/pull/1058.
        // We will work to upstream a fix, but in the meantime, this early return avoids the issue entirely.
        if password.is_none() {
            return Ok(mz_stream);
        }

        // Read a single backend message, which may be a password request. Send ours if so.
        // Otherwise start shuffling bytes. message type (len 1, 'R') + message len (len 4, 8_i32) +
        // auth type (len 4, 3_i32).
        let mut maybe_auth_frame = [0; 1 + 4 + 4];
        let nread = netio::read_exact_or_eof(&mut mz_stream, &mut maybe_auth_frame).await?;
        // 'R' for auth message, 0008 for message length, 0003 for password cleartext variant.
        // See: https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-AUTHENTICATIONCLEARTEXTPASSWORD
        const AUTH_PASSWORD_CLEARTEXT: [u8; 9] = [b'R', 0, 0, 0, 8, 0, 0, 0, 3];
        if nread == AUTH_PASSWORD_CLEARTEXT.len()
            && maybe_auth_frame == AUTH_PASSWORD_CLEARTEXT
            && password.is_some()
        {
            // If we got exactly a cleartext password request and have one, send it.
            let Some(password) = password else {
                unreachable!("verified some above");
            };
            let password = FrontendMessage::Password { password };
            buf.clear();
            password.encode(&mut buf)?;
            mz_stream.write_all(&buf).await?;
            mz_stream.flush().await?;
        } else {
            // Otherwise pass on the bytes we just got. This *might* even be a password request, but
            // we don't have a password. In which case it can be forwarded up to the client.
            client_stream.write_all(&maybe_auth_frame[0..nread]).await?;
        }

        Ok(mz_stream)
    }
}

impl mz_server_core::Server for PgwireBalancer {
    const NAME: &'static str = "pgwire_balancer";

    fn handle_connection(
        &self,
        conn: Connection,
        _tokio_metrics_intervals: impl Iterator<Item = TaskMetrics> + Send + 'static,
    ) -> mz_server_core::ConnectionHandler {
        let tls = self.tls.clone();
        let internal_tls = self.internal_tls.clone();
        let resolver = Arc::clone(&self.resolver);
        let inner_metrics = self.metrics.clone();
        let outer_metrics = self.metrics.clone();
        let limiter = Arc::clone(&self.limiter);
        let pre_resolved_timeout = PRE_RESOLVED_TIMEOUT.get(&self.configs);
        let cancellation_resolver = Arc::clone(&self.cancellation_resolver);
        let conn_uuid = epoch_to_uuid_v7(&(self.now)());
        conn.uuid_handle().set(conn_uuid);
        Box::pin(async move {
            // TODO: Try to merge this with pgwire/server.rs to avoid the duplication. May not be
            // worth it.
            let active_guard = outer_metrics.active_connections();
            let result: Result<(), anyhow::Error> = async move {
                let Some(_conn_guard) = limiter.acquire() else {
                    return Ok(());
                };
                let pre_resolved = inner_metrics.pre_resolved_connections();

                let destination = under_pre_resolved_timeout(
                    pre_resolved_timeout,
                    &inner_metrics,
                    Self::pre_resolve(
                        conn,
                        tls,
                        &resolver,
                        cancellation_resolver,
                        conn_uuid,
                        &inner_metrics,
                    ),
                )
                .await?;

                let Some((mut conn, resolved, params)) = destination else {
                    return Ok(());
                };
                drop(pre_resolved);

                PgwireBalancer::proxy(
                    &mut conn,
                    resolved,
                    params,
                    internal_tls.as_ref(),
                    &inner_metrics,
                )
                .await?;
                conn.flush().await?;
                Ok(())
            }
            .await;
            drop(active_guard);
            outer_metrics.connection_status(result.is_ok()).inc();
            Ok(())
        })
    }
}

// A struct that counts bytes exchanged.
struct CountingConn<C> {
    inner: C,
    read: usize,
    written: usize,
}

impl<C> CountingConn<C> {
    fn new(inner: C) -> Self {
        CountingConn {
            inner,
            read: 0,
            written: 0,
        }
    }
}

impl<C> AsyncRead for CountingConn<C>
where
    C: AsyncRead + Unpin,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        let counter = self.get_mut();
        let pin = Pin::new(&mut counter.inner);
        let bytes = buf.filled().len();
        let poll = pin.poll_read(cx, buf);
        let bytes = buf.filled().len() - bytes;
        if let std::task::Poll::Ready(Ok(())) = poll {
            counter.read += bytes
        }
        poll
    }
}

impl<C> AsyncWrite for CountingConn<C>
where
    C: AsyncWrite + Unpin,
{
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<Result<usize, std::io::Error>> {
        let counter = self.get_mut();
        let pin = Pin::new(&mut counter.inner);
        let poll = pin.poll_write(cx, buf);
        if let std::task::Poll::Ready(Ok(bytes)) = poll {
            counter.written += bytes
        }
        poll
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        let counter = self.get_mut();
        let pin = Pin::new(&mut counter.inner);
        pin.poll_flush(cx)
    }

    fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        let counter = self.get_mut();
        let pin = Pin::new(&mut counter.inner);
        pin.poll_shutdown(cx)
    }
}

/// Broadcasts cancellation to all matching environmentds. `conn_id`'s bits [31..20] are the lower
/// 12 bits of a UUID for an environmentd/organization. Using that and the template in
/// `cancellation_resolver` we generate a hostname. That hostname resolves to all IPs of envds that
/// match the UUID (cloud k8s infrastructure maintains that mapping). This function creates a new
/// task for each envd and relays the cancellation message to it, broadcasting it to any envd that
/// might match the connection.
///
/// This function returns after it has spawned the tasks, and does not wait for them to complete.
/// This is acceptable because cancellation in the Postgres protocol is best effort and has no
/// guarantees.
///
/// The safety of broadcasting this is due to the various randomness in the connection id and secret
/// key, which must match exactly in order to execute a query cancellation. The connection id has 19
/// bits of randomness, and the secret key the full 32, for a total of 51 bits. That is more than
/// 2e15 combinations, enough to nearly certainly prevent two different envds generating identical
/// combinations.
async fn cancel_request(
    conn_id: u32,
    secret_key: u32,
    cancellation_resolver: &CancellationResolver,
) {
    let suffix = conn_id_org_uuid(conn_id);
    let contents = match cancellation_resolver {
        CancellationResolver::Directory(dir) => {
            let path = dir.join(&suffix);
            match std::fs::read_to_string(&path) {
                Ok(contents) => contents,
                Err(err) => {
                    error!("could not read cancel file {path:?}: {err}");
                    return;
                }
            }
        }
        CancellationResolver::Static(addr) => addr.to_owned(),
    };
    let mut all_ips = Vec::new();
    for addr in contents.lines() {
        let addr = addr.trim();
        if addr.is_empty() {
            continue;
        }
        match tokio::net::lookup_host(addr).await {
            Ok(ips) => all_ips.extend(ips),
            Err(err) => {
                error!("{addr} failed resolution: {err}");
            }
        }
    }
    let mut buf = BytesMut::with_capacity(16);
    let msg = FrontendStartupMessage::CancelRequest {
        conn_id,
        secret_key,
    };
    msg.encode(&mut buf).expect("must encode");
    let buf = buf.freeze();
    for ip in all_ips {
        debug!("cancelling {suffix} to {ip}");
        let buf = buf.clone();
        spawn(|| "cancel request for ip", async move {
            let send = async {
                let mut stream = TcpStream::connect(&ip).await?;
                stream.write_all(&buf).await?;
                stream.shutdown().await?;
                Ok::<_, io::Error>(())
            };
            if let Err(err) = send.await {
                error!("error mirroring cancel to {ip}: {err}");
            }
        });
    }
}

/// Writes an HTTP error response to a client and closes the connection.
///
/// The proxied connections are raw TCP streams, so the HTTP framing has to be written by hand.
/// Errors are ignored: the connection is going away regardless.
async fn send_http_error(client_stream: &mut Box<dyn ClientStream>, status: &str, body: &str) {
    let response = format!(
        "HTTP/1.1 {status}\r\n\
         Content-Type: text/plain\r\n\
         Content-Length: {}\r\n\
         Connection: close\r\n\
         \r\n\
         {body}",
        body.len(),
    );
    let _ = client_stream.write_all(response.as_bytes()).await;
    let _ = client_stream.shutdown().await;
}

struct HttpsBalancer {
    resolver: Arc<TenantDnsResolver>,
    tls: Option<ReloadingSslContext>,
    resolve_template: Arc<str>,
    port: u16,
    metrics: Arc<ServerMetrics>,
    limiter: Arc<ConnectionLimiter>,
    configs: ConfigSet,
    internal_tls: Option<SslConnector>,
}

impl HttpsBalancer {
    /// Negotiates TLS and resolves a backend using the client's SNI hostname.
    async fn pre_resolve(
        conn: Connection,
        tls_context: Option<ReloadingSslContext>,
        resolver: &TenantDnsResolver,
        resolve_template: &str,
        port: u16,
    ) -> Result<(Box<dyn ClientStream>, ResolvedAddr, SocketAddr), anyhow::Error> {
        let peer_addr = conn.peer_addr().context("fetching peer addr")?;
        let (client_stream, servername): (Box<dyn ClientStream>, Option<String>) = match tls_context
        {
            Some(tls_context) => {
                let mut ssl_stream = SslStream::new(Ssl::new(&tls_context.get())?, conn)?;
                if let Err(e) = Pin::new(&mut ssl_stream).accept().await {
                    let _ = ssl_stream.get_mut().shutdown().await;
                    return Err(e.into());
                }
                let servername: Option<String> =
                    ssl_stream.ssl().servername(NameType::HOST_NAME).map(|sn| {
                        match sn.split_once('.') {
                            Some((left, _right)) => left,
                            None => sn,
                        }
                        .into()
                    });
                debug!("Found sni servername: {servername:?} (https)");
                (Box::new(ssl_stream), servername)
            }
            _ => (Box::new(conn), None),
        };
        let resolved =
            Self::resolve(resolver, resolve_template, port, servername.as_deref()).await?;
        Ok((client_stream, resolved, peer_addr))
    }

    async fn resolve(
        resolver: &TenantDnsResolver,
        resolve_template: &str,
        port: u16,
        servername: Option<&str>,
    ) -> Result<ResolvedAddr, anyhow::Error> {
        let (addr, tenant) = match servername {
            Some(sni) => resolver.resolve_sni(resolve_template, port, sni).await?,
            None => {
                // Without SNI, resolve the template as is. Not expected for
                // HTTPS in practice.
                debug!("https hostname (no SNI): {}:{}", resolve_template, port);
                resolver.resolve(resolve_template, port).await?
            }
        };

        Ok(ResolvedAddr {
            addr,
            password: None,
            tenant,
        })
    }
}

/// Extracts the tenant ID from an environmentd CNAME target.
///
/// The CNAME points at the environmentd service, of the form
/// `<service>.<namespace>.svc.cluster.local`, e.g.
/// `environmentd.environment-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-0.svc.cluster.local`.
/// The `<namespace>` is the environment name `environment-<tenant_id>-<index>`,
/// where `<tenant_id>` is the tenant's UUID and `<index>` is the environment
/// generation.
///
/// NOTE: `<index>` is currently always 0, since a tenant has one environment
/// per region, but this does not rely on that so that multiple environments
/// per tenant can be supported later.
fn extract_tenant_from_cname(cname: &str) -> Option<String> {
    let mut parts = cname.split('.');
    let _service = parts.next();
    let Some(namespace) = parts.next() else {
        return None;
    };
    // Trim off the starting `environment-`.
    let Some((_, namespace)) = namespace.split_once('-') else {
        return None;
    };
    // Trim off the ending `-<index>`.
    let Some((tenant, _)) = namespace.rsplit_once('-') else {
        return None;
    };
    // Convert to a Uuid so that this tenant matches the frontegg resolver exactly, because it
    // also uses Uuid::to_string.
    let Ok(tenant) = Uuid::parse_str(tenant) else {
        error!("cname tenant not a uuid: {tenant}");
        return None;
    };
    Some(tenant.to_string())
}

/// Strips the surrounding brackets from an IPv6 host literal, e.g. `[::1]`
/// becomes `::1`. Leaves other hosts unchanged. This lets a bracketed IPv6
/// literal in an address template parse as an `IpAddr` and resolve, matching
/// what `tokio::net::lookup_host` accepts.
fn strip_ipv6_brackets(host: &str) -> &str {
    host.strip_prefix('[')
        .and_then(|h| h.strip_suffix(']'))
        .unwrap_or(host)
}

impl mz_server_core::Server for HttpsBalancer {
    const NAME: &'static str = "https_balancer";

    // TODO(jkosh44) consider forwarding the connection UUID to the adapter.
    fn handle_connection(
        &self,
        conn: Connection,
        _tokio_metrics_intervals: impl Iterator<Item = TaskMetrics> + Send + 'static,
    ) -> mz_server_core::ConnectionHandler {
        let tls_context = self.tls.clone();
        let internal_tls = self.internal_tls.clone();
        let resolver = Arc::clone(&self.resolver);
        let resolve_template = Arc::clone(&self.resolve_template);
        let port = self.port;
        let inner_metrics = Arc::clone(&self.metrics);
        let outer_metrics = Arc::clone(&self.metrics);
        let limiter = Arc::clone(&self.limiter);
        let pre_resolved_timeout = PRE_RESOLVED_TIMEOUT.get(&self.configs);
        let inject_proxy_headers = INJECT_PROXY_PROTOCOL_HEADER_HTTP.get(&self.configs);
        Box::pin(async move {
            let active_guard = inner_metrics.active_connections();
            let result: Result<_, anyhow::Error> = Box::pin(async move {
                let Some(_conn_guard) = limiter.acquire() else {
                    return Ok(());
                };
                let pre_resolved = inner_metrics.pre_resolved_connections();

                let (mut client_stream, resolved, peer_addr) = under_pre_resolved_timeout(
                    pre_resolved_timeout,
                    &inner_metrics,
                    Self::pre_resolve(conn, tls_context, &resolver, &resolve_template, port),
                )
                .await?;
                drop(pre_resolved);
                let inner_active_guard = resolved
                    .tenant
                    .as_ref()
                    .map(|tenant| inner_metrics.tenant_connections(tenant));
                let mut mz_stream = match TcpStream::connect(resolved.addr).await {
                    Ok(stream) => stream,
                    Err(e) => {
                        error!("failed to connect to upstream server: {e}");
                        send_http_error(
                            &mut client_stream,
                            "502 Bad Gateway",
                            "upstream server not available",
                        )
                        .await;
                        return Ok(());
                    }
                };

                if inject_proxy_headers {
                    // Write the tcp proxy header
                    let addrs = ProxiedAddress::stream(peer_addr, resolved.addr);
                    let header = ProxyHeader::with_address(addrs);
                    let mut buf = [0u8; 1024];
                    let len = header.encode_to_slice_v2(&mut buf)?;
                    mz_stream.write_all(&buf[..len]).await?;
                }

                let mut mz_stream = if let Some(internal_tls) = &internal_tls {
                    // do a TLS handshake
                    let mut ssl = internal_tls
                        .configure()?
                        .into_ssl(&resolved.addr.to_string())?;
                    ssl.set_connect_state();
                    Conn::Ssl(SslStream::new(ssl, mz_stream)?)
                } else {
                    Conn::Unencrypted(mz_stream)
                };

                let mut client_counter = CountingConn::new(client_stream);

                // Now blindly shuffle bytes back and forth until closed.
                // TODO: Limit total memory use.
                // See corresponding comment in pgwire implementation about ignoring the error.
                let _ = tokio::io::copy_bidirectional(&mut client_counter, &mut mz_stream).await;
                if let Some(tenant) = &resolved.tenant {
                    inner_metrics
                        .tenant_connections_tx(tenant)
                        .inc_by(u64::cast_from(client_counter.written));
                    inner_metrics
                        .tenant_connections_rx(tenant)
                        .inc_by(u64::cast_from(client_counter.read));
                }
                drop(inner_active_guard);
                Ok(())
            })
            .await;
            drop(active_guard);
            outer_metrics.connection_status(result.is_ok()).inc();
            if let Err(e) = result {
                debug!("connection error: {e}");
            }
            Ok(())
        })
    }
}

/// Template for constructing the destination hostname from a TLS SNI
/// servername. `{}` is replaced with the first label of the servername.
#[derive(Debug)]
pub struct SniTemplate {
    pub template: String,
    pub port: u16,
}

trait ClientStream: AsyncRead + AsyncWrite + Unpin + Send {}
impl<T: AsyncRead + AsyncWrite + Unpin + Send> ClientStream for T {}

#[derive(Debug)]
pub enum BalancerResolver {
    Static(String),
    MultiTenant {
        dns: Arc<TenantDnsResolver>,
        frontegg: FronteggResolver,
        sni: Option<SniTemplate>,
    },
}

/// An error resolving a connection's destination.
///
/// The `Display` of this error is sent to unauthenticated clients, so it must
/// not contain internal details such as hostnames. Those belong in the source
/// error attached to each variant, which is only logged.
#[derive(Debug, thiserror::Error)]
enum ResolveError {
    #[error("invalid password")]
    InvalidPassword,
    /// A client protocol violation, e.g. sending the wrong message during
    /// startup. An unauthenticated client can trigger these at will, so they
    /// are logged at `warn!`, not `error!`, to avoid making client noise look
    /// like server faults and spamming the error log.
    #[error("internal error")]
    Client(#[source] anyhow::Error),
    /// The tenant's upstream backend could not be reached, e.g. its hostname
    /// did not resolve. Reported to the client with a distinct, non-leaking
    /// message rather than a generic internal error, so a down environment is
    /// not mistaken for a balancerd bug. Logged at `warn!`: a bogus SNI reaches
    /// this from an unauthenticated client, and a genuinely down environment is
    /// an operational condition, not a balancerd fault.
    #[error("upstream server not available")]
    Upstream(#[source] anyhow::Error),
    /// A server-side fault.
    #[error("internal error")]
    Internal(#[from] anyhow::Error),
}

impl From<io::Error> for ResolveError {
    fn from(e: io::Error) -> Self {
        ResolveError::Internal(e.into())
    }
}

impl BalancerResolver {
    /// Returns a clone of the shared DNS resolver if in multi-tenant mode.
    /// This allows sharing the resolver with other components like HttpsBalancer.
    pub fn shared_dns(&self) -> Option<Arc<TenantDnsResolver>> {
        match self {
            BalancerResolver::Static(_) => None,
            BalancerResolver::MultiTenant { dns, .. } => Some(Arc::clone(dns)),
        }
    }

    async fn resolve<A>(
        &self,
        conn: &mut FramedConn<A>,
        user: &str,
        metrics: &ServerMetrics,
    ) -> Result<ResolvedAddr, ResolveError>
    where
        A: AsyncRead + AsyncWrite + Unpin,
    {
        match self {
            BalancerResolver::MultiTenant {
                dns: dns_resolver,
                frontegg:
                    FronteggResolver {
                        auth,
                        addr_template,
                    },
                sni: sni_resolver,
            } => {
                let servername = match conn.inner() {
                    Conn::Ssl(ssl_stream) => {
                        ssl_stream.ssl().servername(NameType::HOST_NAME).map(|sn| {
                            match sn.split_once('.') {
                                Some((left, _right)) => left,
                                None => sn,
                            }
                        })
                    }
                    Conn::Unencrypted(_) => None,
                };
                let has_sni = servername.is_some();
                let resolved_addr = match (servername, sni_resolver.as_ref()) {
                    (Some(servername), Some(SniTemplate { template, port })) => {
                        // A resolution failure here means the tenant's backend
                        // is unreachable (or the client sent a bogus SNI). Not
                        // a server fault.
                        let (addr, tenant) = dns_resolver
                            .resolve_sni(template, *port, servername)
                            .await
                            .map_err(ResolveError::Upstream)?;
                        debug!("pgwire SNI resolved tenant: {:?}", tenant);
                        ResolvedAddr {
                            addr,
                            password: None,
                            tenant,
                        }
                    }
                    _ => {
                        conn.send(BackendMessage::AuthenticationCleartextPassword)
                            .await?;
                        conn.flush().await?;
                        let password = match conn.recv().await? {
                            Some(FrontendMessage::Password { password }) => password,
                            _ => {
                                return Err(ResolveError::Client(anyhow::anyhow!(
                                    "expected Password message"
                                )));
                            }
                        };

                        // balancerd only needs the validated tenant_id to route
                        // the connection; group extraction happens in
                        // environmentd, so skip it here.
                        let auth_response = auth.authenticate(user, &password, None).await;
                        let auth_session = match auth_response {
                            Ok((auth_session, _)) => auth_session,
                            Err(e) => {
                                warn!("pgwire connection failed authentication: {}", e);
                                return Err(ResolveError::InvalidPassword);
                            }
                        };

                        let hostname_with_port =
                            addr_template.replace("{}", &auth_session.tenant_id().to_string());
                        let (hostname, port_str) = hostname_with_port
                            .rsplit_once(':')
                            .ok_or_else(|| anyhow::anyhow!("port required in addr_template"))?;
                        let port: u16 = port_str.parse().with_context(|| {
                            format!("invalid port in addr_template: {}", port_str)
                        })?;
                        // The tenant is already known from authentication, so
                        // skip the CNAME lookup that resolve() would do. A
                        // failure here means the tenant's backend is unreachable.
                        let addr = dns_resolver
                            .resolve_addr(hostname, port)
                            .await
                            .map_err(ResolveError::Upstream)?;
                        let tenant = auth_session.tenant_id().to_string();
                        debug!("Frontegg resolved tenant: {}", tenant);
                        ResolvedAddr {
                            addr,
                            password: Some(password),
                            tenant: Some(tenant),
                        }
                    }
                };
                metrics
                    .tenant_pgwire_sni_count(
                        resolved_addr.tenant.as_deref().unwrap_or("unknown"),
                        has_sni,
                    )
                    .inc();

                Ok(resolved_addr)
            }
            BalancerResolver::Static(addr) => {
                // We don't want any caching here so we just use the standard
                // tokio resolver.
                let Some(addr) = tokio::net::lookup_host(addr).await?.next() else {
                    return Err(anyhow::anyhow!("{addr} did not resolve to any addresses").into());
                };

                Ok(ResolvedAddr {
                    addr,
                    password: None,
                    tenant: None,
                })
            }
        }
    }
}

/// Creates a resolver from the system DNS configuration.
///
/// Caching is delegated to the infrastructure (node-local DNS), so this
/// resolver does no caching of its own. Fails if the system DNS configuration
/// cannot be read. We must not fall back to hickory's default config (Google
/// public DNS) here, that would leak internal hostnames to an external party
/// and could not resolve them anyway.
fn create_resolver() -> Result<TokioResolver, anyhow::Error> {
    let (config, mut opts) = read_system_conf().context("reading system DNS configuration")?;
    opts.cache_size = 0;
    // Query A records first and AAAA only on failure, rather than both.
    opts.ip_strategy = LookupIpStrategy::Ipv4thenIpv6;

    Ok(
        Resolver::builder_with_config(config, TokioRuntimeProvider::default())
            .with_options(opts)
            .build()?,
    )
}

/// Resolves tenant hostnames for pgwire and HTTPS routing.
///
/// Caching is delegated to the infrastructure (node-local DNS), so every
/// lookup issues a query. CNAMEs are resolved separately from A records only
/// because the CNAME carries the tenant, not for caching reasons.
#[derive(Debug)]
pub struct TenantDnsResolver {
    resolver: TokioResolver,
}

impl TenantDnsResolver {
    /// Creates a new resolver. Fails if the system DNS configuration cannot be
    /// read.
    pub fn new() -> Result<Self, anyhow::Error> {
        Ok(Self {
            resolver: create_resolver()?,
        })
    }

    /// Resolves the CNAME a hostname points at, if any.
    async fn resolve_cname(&self, hostname: &str) -> Option<String> {
        match self.resolver.lookup(hostname, RecordType::CNAME).await {
            Ok(cname_response) => {
                if let Some(cname_record) = cname_response.answers().first() {
                    if let RData::CNAME(cname_data) = &cname_record.data {
                        let cname = cname_data.to_string();
                        debug!("CNAME for {}: {}", hostname, cname);
                        return Some(cname);
                    }
                }
                None
            }
            Err(e) => {
                debug!("CNAME lookup failed for {}: {}", hostname, e);
                None
            }
        }
    }

    /// Resolves the A records for a hostname.
    async fn resolve_a(&self, hostname: &str) -> Result<LookupIp, anyhow::Error> {
        self.resolver
            .lookup_ip(hostname)
            .await
            .with_context(|| format!("resolving A records for {}", hostname))
    }

    /// Resolves the environment address for a TLS SNI servername.
    ///
    /// `servername` is the first label of the SNI host, e.g.
    /// `3dl07g8zmj91pntk4eo9cfvwe`. Substituting it into `template` (e.g.
    /// `blncr-{}`) yields a Kubernetes hostname like
    /// `blncr-3dl07g8zmj91pntk4eo9cfvwe`, which resolves via a CNAME to the
    /// environmentd service, e.g.
    /// `environmentd.environment-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-0.svc.cluster.local`.
    /// The tenant is extracted from that CNAME. See
    /// `extract_tenant_from_cname`.
    pub async fn resolve_sni(
        &self,
        template: &str,
        port: u16,
        servername: &str,
    ) -> Result<(SocketAddr, Option<String>), anyhow::Error> {
        let hostname = template.replace("{}", servername);
        debug!("SNI hostname: {}", hostname);
        self.resolve(&hostname, port).await
    }

    /// Resolves the address for a hostname, skipping CNAME resolution and
    /// tenant extraction. Use when the tenant is already known.
    async fn resolve_addr(&self, host: &str, port: u16) -> Result<SocketAddr, anyhow::Error> {
        let host = strip_ipv6_brackets(host);
        // IP literals need no resolution. Resolving them through hickory
        // would walk the search domain list first when ndots is large, as it
        // is in Kubernetes.
        if let Ok(ip) = host.parse::<IpAddr>() {
            return Ok(SocketAddr::new(ip, port));
        }
        Self::first_addr(self.resolve_a(host).await?, port)
    }

    /// Resolves the address and tenant from a hostname and port.
    ///
    /// The tenant is extracted from the CNAME if the hostname points at one.
    async fn resolve(
        &self,
        host: &str,
        port: u16,
    ) -> Result<(SocketAddr, Option<String>), anyhow::Error> {
        let host = strip_ipv6_brackets(host);
        // IP literals need no resolution and carry no tenant CNAME.
        if let Ok(ip) = host.parse::<IpAddr>() {
            return Ok((SocketAddr::new(ip, port), None));
        }

        // The CNAME carries the tenant, so resolve it separately to extract it.
        let (ips, tenant) = if let Some(cname) = self.resolve_cname(host).await {
            let tenant = extract_tenant_from_cname(&cname);
            (self.resolve_a(&cname).await?, tenant)
        } else {
            (self.resolve_a(host).await?, None)
        };

        Ok((Self::first_addr(ips, port)?, tenant))
    }

    /// Returns the first resolved IP as a socket address.
    fn first_addr(ips: LookupIp, port: u16) -> Result<SocketAddr, anyhow::Error> {
        ips.iter()
            .next()
            .map(|ip| SocketAddr::new(ip, port))
            .ok_or_else(|| anyhow::anyhow!("no A records found in DNS response"))
    }
}

#[derive(Debug)]
pub struct FronteggResolver {
    pub auth: FronteggAuthentication,
    pub addr_template: String,
}

#[derive(Debug)]
struct ResolvedAddr {
    addr: SocketAddr,
    password: Option<String>,
    tenant: Option<String>,
}

#[cfg(test)]
mod tests {
    use mz_dyncfg::ConfigUpdates;

    use super::*;

    #[mz_ore::test]
    fn test_tenant() {
        let tests = vec![
            ("", None),
            (
                "environmentd.environment-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-0.svc.cluster.local",
                Some("58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3"),
            ),
            (
                // Trailing dot from an absolute DNS name, as returned by the
                // resolver.
                "environmentd.environment-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-0.svc.cluster.local.",
                Some("58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3"),
            ),
            (
                // Variously named parts.
                "service.something-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-0.ssvvcc.cloister.faraway",
                Some("58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3"),
            ),
            (
                // No dashes in uuid.
                "environmentd.environment-58cd23ffa4d74bd0ad85a6ff29cc86c3-0.svc.cluster.local",
                Some("58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3"),
            ),
            (
                // -1234 suffix.
                "environmentd.environment-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-1234.svc.cluster.local",
                Some("58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3"),
            ),
            (
                // Uppercase.
                "environmentd.environment-58CD23FF-A4D7-4BD0-AD85-A6FF29CC86C3-0.svc.cluster.local",
                Some("58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3"),
            ),
            (
                // No -number suffix.
                "environmentd.environment-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3.svc.cluster.local",
                None,
            ),
            (
                // No service name.
                "environment-58cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-0.svc.cluster.local",
                None,
            ),
            (
                // Invalid UUID.
                "environmentd.environment-8cd23ff-a4d7-4bd0-ad85-a6ff29cc86c3-0.svc.cluster.local",
                None,
            ),
        ];
        for (name, expect) in tests {
            let cname = extract_tenant_from_cname(name);
            assert_eq!(
                cname.as_deref(),
                expect,
                "{name} got {cname:?} expected {expect:?}"
            );
        }
    }

    #[mz_ore::test]
    fn test_strip_ipv6_brackets() {
        assert_eq!(strip_ipv6_brackets("[::1]"), "::1");
        assert_eq!(strip_ipv6_brackets("[2001:db8::1]"), "2001:db8::1");
        // Hosts without a matched bracket pair are left untouched.
        assert_eq!(strip_ipv6_brackets("127.0.0.1"), "127.0.0.1");
        assert_eq!(strip_ipv6_brackets("host.example.com"), "host.example.com");
        assert_eq!(strip_ipv6_brackets("[unclosed"), "[unclosed");
        assert_eq!(strip_ipv6_brackets("unopened]"), "unopened]");
    }

    #[mz_ore::test]
    fn test_connection_limiter() {
        let configs = dyncfgs::all_dyncfgs(ConfigSet::default());
        let set_max = |max: u32| {
            let mut updates = ConfigUpdates::default();
            updates.add(&MAX_CONNECTIONS, max);
            updates.apply(&configs);
        };

        set_max(2);
        let limiter = ConnectionLimiter::new(&MetricsRegistry::new(), configs.clone());
        let first = limiter.acquire().expect("under the limit");
        let second = limiter.acquire().expect("at the limit");
        assert!(limiter.acquire().is_none());
        assert_eq!(limiter.rejected.get(), 1);

        // A connection closing frees capacity for a new one.
        drop(second);
        let _third = limiter.acquire().expect("capacity freed");
        drop(first);

        // Lowering the limit below the current count refuses new connections but leaves the
        // established ones alone.
        set_max(1);
        assert!(limiter.acquire().is_none());

        // Zero disables the limit.
        set_max(0);
        assert!(limiter.acquire().is_some());
    }

    /// A zero timeout disables the deadline, which is what leaves the phase unbounded.
    #[mz_ore::test(tokio::test(start_paused = true))]
    async fn test_pre_resolved_timeout_zero_disables() {
        let metrics = ServerMetrics::new(
            ServerMetricsConfig::register_into(&MetricsRegistry::new()),
            "pgwire",
        );

        // Zero lets the phase run as long as it likes. The paused clock makes the hour free.
        let slow = async {
            tokio::time::sleep(Duration::from_secs(60 * 60)).await;
            Ok::<(), io::Error>(())
        };
        assert!(
            under_pre_resolved_timeout(Duration::ZERO, &metrics, slow)
                .await
                .is_ok()
        );
        assert_eq!(metrics.pre_resolved_timeout().get(), 0);

        // Anything else bounds it, and the overrun is counted.
        let stalled = std::future::pending::<Result<(), io::Error>>();
        let err = under_pre_resolved_timeout(Duration::from_secs(30), &metrics, stalled)
            .await
            .expect_err("a phase that never finishes must be cut short");
        assert_eq!(err.kind(), io::ErrorKind::TimedOut);
        assert_eq!(metrics.pre_resolved_timeout().get(), 1);
    }
}