mod.rs

#![allow(dead_code)]

use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use chrono::{DateTime, Duration, SecondsFormat, Utc};
use p256::ecdsa::{SigningKey, signature::Signer};
use p256::pkcs8::EncodePublicKey;
use rcgen::{
    BasicConstraints, CertificateParams, CertifiedIssuer, CustomExtension, DistinguishedName,
    DnType, ExtendedKeyUsagePurpose, IsCa, KeyPair, KeyUsagePurpose,
};
use rustls_pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
use sqlx::postgres::PgPoolOptions;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use x509_parser::prelude::*;

use agent_gateway::policy::PolicyEngine;
use agent_gateway::proxy::MakeProxyService;

const CLIENT_EXTENSION_OID: &[u64] = &[1, 3, 6, 1, 4, 1, 57264, 1, 1];
static TEST_ID: AtomicU64 = AtomicU64::new(1);
static TEST_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations");

// ---------------------------------------------------------------------------
// PKI generation
// ---------------------------------------------------------------------------

pub fn generate_ca() -> Result<CertifiedIssuer<'static, KeyPair>, rcgen::Error> {
    let mut params = CertificateParams::new(Vec::<String>::new())?;
    params.distinguished_name = DistinguishedName::new();
    params
        .distinguished_name
        .push(DnType::CommonName, "agent-gateway test CA");
    params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
    params.key_usages = vec![
        KeyUsagePurpose::DigitalSignature,
        KeyUsagePurpose::KeyCertSign,
        KeyUsagePurpose::CrlSign,
    ];

    let key_pair = KeyPair::generate()?;
    CertifiedIssuer::self_signed(params, key_pair)
}

pub fn generate_server_cert(
    issuer: &CertifiedIssuer<'_, impl rcgen::SigningKey>,
) -> Result<(rcgen::Certificate, KeyPair), rcgen::Error> {
    let mut params = CertificateParams::new(vec!["localhost".to_owned(), "127.0.0.1".to_owned()])?;
    params.distinguished_name = DistinguishedName::new();
    params
        .distinguished_name
        .push(DnType::CommonName, "localhost");
    params.key_usages = vec![
        KeyUsagePurpose::DigitalSignature,
        KeyUsagePurpose::KeyEncipherment,
    ];
    params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth];

    let key_pair = KeyPair::generate()?;
    let cert = params.signed_by(&key_pair, issuer)?;
    Ok((cert, key_pair))
}

pub fn generate_client_cert(
    issuer: &CertifiedIssuer<'_, impl rcgen::SigningKey>,
    extension_value: &str,
) -> Result<(rcgen::Certificate, KeyPair), rcgen::Error> {
    let mut params = CertificateParams::new(Vec::<String>::new())?;
    params.distinguished_name = DistinguishedName::new();
    params
        .distinguished_name
        .push(DnType::CommonName, "test client");
    params.key_usages = vec![KeyUsagePurpose::DigitalSignature];
    params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ClientAuth];
    params
        .custom_extensions
        .push(CustomExtension::from_oid_content(
            CLIENT_EXTENSION_OID,
            der_encode_utf8_string(extension_value),
        ));

    let key_pair = KeyPair::generate()?;
    let cert = params.signed_by(&key_pair, issuer)?;
    Ok((cert, key_pair))
}

/// Generate a client cert with no custom extension at all.
pub fn generate_client_cert_no_extension(
    issuer: &CertifiedIssuer<'_, impl rcgen::SigningKey>,
) -> Result<(rcgen::Certificate, KeyPair), rcgen::Error> {
    let mut params = CertificateParams::new(Vec::<String>::new())?;
    params.distinguished_name = DistinguishedName::new();
    params
        .distinguished_name
        .push(DnType::CommonName, "test client (no ext)");
    params.key_usages = vec![KeyUsagePurpose::DigitalSignature];
    params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ClientAuth];

    let key_pair = KeyPair::generate()?;
    let cert = params.signed_by(&key_pair, issuer)?;
    Ok((cert, key_pair))
}

pub struct TestPki {
    pub ca: CertifiedIssuer<'static, KeyPair>,
    pub server_cert: rcgen::Certificate,
    pub server_key: KeyPair,
    pub client_cert: rcgen::Certificate,
    pub client_key: KeyPair,
}

impl TestPki {
    pub fn new(client_extension_value: &str) -> Self {
        let ca = generate_ca().expect("generate CA");
        let (server_cert, server_key) = generate_server_cert(&ca).expect("generate server cert");
        let (client_cert, client_key) =
            generate_client_cert(&ca, client_extension_value).expect("generate client cert");

        Self {
            ca,
            server_cert,
            server_key,
            client_cert,
            client_key,
        }
    }

    pub fn ca_cert_der(&self) -> CertificateDer<'static> {
        CertificateDer::from(self.ca.der().to_vec())
    }

    pub fn server_cert_chain(&self) -> Vec<CertificateDer<'static>> {
        vec![CertificateDer::from(self.server_cert.der().to_vec())]
    }

    pub fn server_key_der(&self) -> PrivateKeyDer<'static> {
        PrivateKeyDer::from(PrivatePkcs8KeyDer::from(self.server_key.serialize_der()))
    }

    pub fn client_cert_chain(&self) -> Vec<CertificateDer<'static>> {
        vec![CertificateDer::from(self.client_cert.der().to_vec())]
    }

    pub fn client_key_der(&self) -> PrivateKeyDer<'static> {
        PrivateKeyDer::from(PrivatePkcs8KeyDer::from(self.client_key.serialize_der()))
    }

    pub fn client_spki_der(&self) -> Vec<u8> {
        certificate_spki_der(&self.client_cert.der().to_vec())
    }
}

pub fn certificate_spki_der(cert_der: &[u8]) -> Vec<u8> {
    let (remaining, cert) = X509Certificate::from_der(cert_der).expect("parse certificate");
    assert!(
        remaining.is_empty(),
        "certificate should not contain trailing bytes"
    );
    cert.tbs_certificate.subject_pki.raw.to_vec()
}

fn der_encode_utf8_string(value: &str) -> Vec<u8> {
    let value_bytes = value.as_bytes();
    let len = value_bytes.len();
    let mut encoded = Vec::with_capacity(2 + len);
    encoded.push(0x0c); // UTF8String tag
    if len < 128 {
        encoded.push(len as u8);
    } else {
        let mut len_bytes = Vec::new();
        let mut remaining = len;
        while remaining > 0 {
            len_bytes.push((remaining & 0xff) as u8);
            remaining >>= 8;
        }
        len_bytes.reverse();
        encoded.push(0x80 | (len_bytes.len() as u8));
        encoded.extend_from_slice(&len_bytes);
    }
    encoded.extend_from_slice(value_bytes);
    encoded
}

// ---------------------------------------------------------------------------
// Tracing capture layer
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct CapturedEvent {
    pub message: String,
    pub fields: std::collections::HashMap<String, String>,
}

pub type EventLog = Arc<Mutex<Vec<CapturedEvent>>>;

pub fn new_event_log() -> EventLog {
    Arc::new(Mutex::new(Vec::new()))
}

struct CaptureLayer {
    log: EventLog,
}

struct CaptureVisitor {
    message: Option<String>,
    fields: std::collections::HashMap<String, String>,
}

impl tracing::field::Visit for CaptureVisitor {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        if field.name() == "message" {
            self.message = Some(format!("{value:?}").trim_matches('"').to_owned());
        } else {
            self.fields
                .insert(field.name().to_owned(), format!("{value:?}"));
        }
    }

    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
        if field.name() == "message" {
            self.message = Some(value.to_owned());
        } else {
            self.fields
                .insert(field.name().to_owned(), value.to_owned());
        }
    }

    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
        self.fields
            .insert(field.name().to_owned(), value.to_string());
    }

    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
        self.fields
            .insert(field.name().to_owned(), value.to_string());
    }

    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
        self.fields
            .insert(field.name().to_owned(), value.to_string());
    }
}

impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for CaptureLayer {
    fn on_event(
        &self,
        event: &tracing::Event<'_>,
        _ctx: tracing_subscriber::layer::Context<'_, S>,
    ) {
        let mut visitor = CaptureVisitor {
            message: None,
            fields: std::collections::HashMap::new(),
        };
        event.record(&mut visitor);
        let captured = CapturedEvent {
            message: visitor.message.unwrap_or_default(),
            fields: visitor.fields,
        };
        self.log.lock().unwrap().push(captured);
    }
}

/// Install the capture layer as the global default subscriber for this test.
/// Returns the event log that accumulates all tracing events.
///
/// NOTE: `tracing::subscriber::set_default` is per-thread only. Since the
/// proxy runs spawned tasks on the tokio runtime, we must use `set_global_default`
/// which can only be called once per process. For test binaries that run
/// tests sequentially in a single process, this works. The `EventLog` is
/// shared across all tests in this binary.
pub fn init_tracing_capture() -> EventLog {
    use std::sync::Once;
    use tracing_subscriber::layer::SubscriberExt;

    static INIT: Once = Once::new();
    static LOG: std::sync::OnceLock<EventLog> = std::sync::OnceLock::new();

    INIT.call_once(|| {
        let log = new_event_log();
        LOG.set(log.clone()).unwrap();
        let subscriber = tracing_subscriber::registry().with(CaptureLayer { log });
        tracing::subscriber::set_global_default(subscriber)
            .expect("setting global tracing subscriber");
    });

    LOG.get().unwrap().clone()
}

/// Drain the event log and return all events captured so far.
pub fn drain_events(log: &EventLog) -> Vec<CapturedEvent> {
    let mut guard = log.lock().unwrap();
    guard.drain(..).collect()
}

/// Find the first event with the given message substring.
pub fn find_event<'a>(events: &'a [CapturedEvent], message: &str) -> Option<&'a CapturedEvent> {
    events.iter().find(|e| e.message.contains(message))
}

/// Poll the event log until an event matching `message` appears, or time out.
/// Returns all drained events (including any captured before the match).
pub async fn wait_for_event(
    log: &EventLog,
    message: &str,
    timeout: std::time::Duration,
) -> Vec<CapturedEvent> {
    let deadline = tokio::time::Instant::now() + timeout;
    let mut collected = Vec::new();
    loop {
        collected.extend(drain_events(log));
        if find_event(&collected, message).is_some() {
            return collected;
        }
        if tokio::time::Instant::now() >= deadline {
            return collected;
        }
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
    }
}

/// Acquire an exclusive lock for e2e tests that share a global tracing subscriber.
/// Hold the returned guard for the duration of the test.
static E2E_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());

pub fn serial_test_lock() -> std::sync::MutexGuard<'static, ()> {
    E2E_MUTEX.lock().unwrap_or_else(|e| e.into_inner())
}

pub fn install_test_crypto_provider() {
    static INIT: std::sync::Once = std::sync::Once::new();
    INIT.call_once(|| {
        rustls::crypto::aws_lc_rs::default_provider()
            .install_default()
            .ok();
    });
}

// ---------------------------------------------------------------------------
// Postgres authorization registry helpers
// ---------------------------------------------------------------------------

#[derive(Clone)]
pub struct TestAuthzRegistry {
    pub pool: sqlx::PgPool,
    signing_key: SigningKey,
    key_id: String,
}

#[derive(Debug, Clone)]
pub struct SeededPermission {
    pub permission_id: String,
    pub normalized_destination: String,
}

impl TestAuthzRegistry {
    pub async fn new() -> Self {
        let database_url =
            std::env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set");
        let pool = PgPoolOptions::new()
            .max_connections(5)
            .connect(&database_url)
            .await
            .expect("connect to TEST_DATABASE_URL");

        TEST_MIGRATOR
            .run(&pool)
            .await
            .expect("run registry migrations");
        let signing_key = test_signing_key();
        let public_key_spki_der = signing_key
            .verifying_key()
            .to_public_key_der()
            .expect("encode verifying key")
            .as_bytes()
            .to_vec();
        let key_id = unique_id("test-key");
        let (not_before, not_after) = active_window();

        sqlx::query(
            r#"
            INSERT INTO principal_signing_keys (
                key_id, algorithm, public_key_spki_der,
                not_before, not_after
            )
            VALUES ($1, 'ecdsa_p256_sha256', $2, $3, $4)
            "#,
        )
        .bind(&key_id)
        .bind(&public_key_spki_der)
        .bind(not_before)
        .bind(not_after)
        .execute(&pool)
        .await
        .expect("insert test signing key");

        Self {
            pool,
            signing_key,
            key_id,
        }
    }

    pub async fn engine(&self, client_ext_oid: &str) -> Arc<dyn PolicyEngine> {
        let database_url =
            std::env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set");
        let config = agent_gateway::config::PolicyConfig {
            client_ext_oid: client_ext_oid.to_owned(),
            database_url: Some(database_url),
            database_url_env: None,
            max_connections: Some(5),
            connect_timeout_ms: Some(5_000),
            pool_acquire_timeout_ms: Some(1_000),
            query_timeout_ms: Some(2_000),
        };
        agent_gateway::policy::build_engine(&config)
            .await
            .expect("build test policy engine")
    }

    pub async fn cleanup(&self) {
        sqlx::query("DELETE FROM permission_registry WHERE signing_key_id = $1")
            .bind(&self.key_id)
            .execute(&self.pool)
            .await
            .expect("delete test permissions");
        sqlx::query("DELETE FROM principal_key_permissions WHERE signing_key_id = $1")
            .bind(&self.key_id)
            .execute(&self.pool)
            .await
            .expect("delete test signer scopes");
        sqlx::query("DELETE FROM principal_signing_keys WHERE key_id = $1")
            .bind(&self.key_id)
            .execute(&self.pool)
            .await
            .expect("delete test signing key");
    }

    pub async fn allow(
        &self,
        subject_identity: &str,
        subject_public_key_spki_der: &[u8],
        destination: &str,
    ) -> SeededPermission {
        self.allow_inner(
            subject_identity,
            subject_public_key_spki_der,
            destination,
            true,
        )
        .await
    }

    pub async fn allow_for_pki(
        &self,
        pki: &TestPki,
        subject_identity: &str,
        destination: &str,
    ) -> SeededPermission {
        let subject_public_key_spki_der = pki.client_spki_der();
        self.allow(subject_identity, &subject_public_key_spki_der, destination)
            .await
    }

    pub async fn allow_without_signer_scope(
        &self,
        subject_identity: &str,
        subject_public_key_spki_der: &[u8],
        destination: &str,
    ) -> SeededPermission {
        self.allow_inner(
            subject_identity,
            subject_public_key_spki_der,
            destination,
            false,
        )
        .await
    }

    pub async fn allow_without_signer_scope_for_pki(
        &self,
        pki: &TestPki,
        subject_identity: &str,
        destination: &str,
    ) -> SeededPermission {
        let subject_public_key_spki_der = pki.client_spki_der();
        self.allow_without_signer_scope(subject_identity, &subject_public_key_spki_der, destination)
            .await
    }

    pub async fn revoke_permission(&self, permission_id: &str) {
        sqlx::query("UPDATE permission_registry SET revoked_at = now() WHERE permission_id = $1")
            .bind(permission_id)
            .execute(&self.pool)
            .await
            .expect("revoke permission");
    }

    pub async fn revoke_signer(&self) {
        sqlx::query("UPDATE principal_signing_keys SET revoked_at = now() WHERE key_id = $1")
            .bind(&self.key_id)
            .execute(&self.pool)
            .await
            .expect("revoke signer");
    }

    pub async fn tamper_permission_destination(&self, permission_id: &str, destination: &str) {
        sqlx::query("UPDATE permission_registry SET destination = $2 WHERE permission_id = $1")
            .bind(permission_id)
            .bind(destination)
            .execute(&self.pool)
            .await
            .expect("tamper permission destination");
    }

    async fn allow_inner(
        &self,
        subject_identity: &str,
        subject_public_key_spki_der: &[u8],
        destination: &str,
        include_scope: bool,
    ) -> SeededPermission {
        let permission_id = unique_id("test-permission");
        let (not_before, not_after) = active_window();

        if include_scope {
            sqlx::query(
                r#"
                INSERT INTO principal_key_permissions (
                    signing_key_id, destination, not_before, not_after
                )
                VALUES ($1, $2, $3, $4)
                "#,
            )
            .bind(&self.key_id)
            .bind(destination)
            .bind(not_before)
            .bind(not_after)
            .execute(&self.pool)
            .await
            .expect("insert signer scope");
        }

        let signed_bytes = test_canonical_permission_bytes(
            &permission_id,
            &self.key_id,
            subject_identity,
            subject_public_key_spki_der,
            destination,
            not_before,
            not_after,
        );
        let signature: p256::ecdsa::Signature = self.signing_key.sign(&signed_bytes);

        sqlx::query(
            r#"
            INSERT INTO permission_registry (
                permission_id, signing_key_id, subject_identity, subject_public_key_spki_der, destination,
                not_before, not_after, signature
            )
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
            "#,
        )
        .bind(&permission_id)
        .bind(&self.key_id)
        .bind(subject_identity)
        .bind(subject_public_key_spki_der)
        .bind(destination)
        .bind(not_before)
        .bind(not_after)
        .bind(signature.to_der().as_bytes())
        .execute(&self.pool)
        .await
        .expect("insert signed permission");

        SeededPermission {
            permission_id,
            normalized_destination: destination.to_owned(),
        }
    }
}

pub struct TestPolicyEngine {
    engine: Arc<dyn PolicyEngine>,
    registry: TestAuthzRegistry,
}

impl TestPolicyEngine {
    pub fn engine(&self) -> Arc<dyn PolicyEngine> {
        self.engine.clone()
    }

    pub async fn cleanup(&self) {
        self.registry.cleanup().await;
    }
}

pub async fn postgres_engine_allowing(
    client_ext_oid: &str,
    pki: &TestPki,
    subject_identity: &str,
    destinations: Vec<String>,
) -> TestPolicyEngine {
    let registry = TestAuthzRegistry::new().await;
    let subject_public_key_spki_der = pki.client_spki_der();
    for destination in destinations {
        registry
            .allow(subject_identity, &subject_public_key_spki_der, &destination)
            .await;
    }
    let engine = registry.engine(client_ext_oid).await;
    TestPolicyEngine { engine, registry }
}

fn test_signing_key() -> SigningKey {
    let id = TEST_ID.fetch_add(1, Ordering::Relaxed);
    let mut bytes = [0u8; 32];
    bytes[20..24].copy_from_slice(&std::process::id().to_be_bytes());
    bytes[24..].copy_from_slice(&id.to_be_bytes());
    SigningKey::from_slice(&bytes).expect("build deterministic test signing key")
}

fn unique_id(prefix: &str) -> String {
    let id = TEST_ID.fetch_add(1, Ordering::Relaxed);
    format!("{prefix}-{}-{id}", std::process::id())
}

pub fn unique_test_identity(prefix: &str) -> String {
    unique_id(prefix)
}

fn active_window() -> (DateTime<Utc>, DateTime<Utc>) {
    let now = DateTime::<Utc>::from_timestamp(Utc::now().timestamp(), 0).unwrap();
    (now - Duration::hours(1), now + Duration::hours(1))
}

fn test_canonical_permission_bytes(
    permission_id: &str,
    signing_key_id: &str,
    subject_identity: &str,
    subject_public_key_spki_der: &[u8],
    destination: &str,
    not_before: DateTime<Utc>,
    not_after: DateTime<Utc>,
) -> Vec<u8> {
    format!(
        "agent-gateway-permission-v1\npermission_id={permission_id}\nsigning_key_id={signing_key_id}\nsubject_identity={subject_identity}\nsubject_public_key_spki_der={}\ndestination={destination}\nnot_before={}\nnot_after={}\n",
        lower_hex(subject_public_key_spki_der),
        not_before.to_rfc3339_opts(SecondsFormat::Micros, true),
        not_after.to_rfc3339_opts(SecondsFormat::Micros, true),
    )
    .into_bytes()
}

fn lower_hex(bytes: &[u8]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut out = String::with_capacity(bytes.len() * 2);
    for &byte in bytes {
        out.push(HEX[(byte >> 4) as usize] as char);
        out.push(HEX[(byte & 0x0f) as usize] as char);
    }
    out
}

// ---------------------------------------------------------------------------
// E2E harness: proxy, client, echo server
// ---------------------------------------------------------------------------

/// A handle that aborts the background accept loop when dropped.
pub struct ServerGuard {
    task: tokio::task::JoinHandle<()>,
}

impl Drop for ServerGuard {
    fn drop(&mut self) {
        self.task.abort();
    }
}

/// Start the real proxy in-process. Returns the bound address and a guard
/// that stops the server when dropped or explicitly shut down.
pub async fn start_proxy(
    pki: &TestPki,
    policy_engine: Arc<dyn PolicyEngine>,
) -> (SocketAddr, ServerGuard) {
    install_test_crypto_provider();

    let mut server_config = rustls::ServerConfig::builder()
        .with_client_cert_verifier(agent_gateway::tls::db_rooted_client_cert_verifier())
        .with_single_cert(pki.server_cert_chain(), pki.server_key_der())
        .unwrap();
    server_config.alpn_protocols = vec![b"h2".to_vec()];

    let tls_acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(server_config));

    let make_service = Arc::new(MakeProxyService::new(policy_engine));

    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();

    let task = tokio::spawn(async move {
        loop {
            let (tcp, peer) = match listener.accept().await {
                Ok(c) => c,
                Err(_) => continue,
            };
            let acceptor = tls_acceptor.clone();
            let svc = make_service.clone();
            tokio::spawn(async move {
                let tls_stream = match acceptor.accept(tcp).await {
                    Ok(s) => s,
                    Err(e) => {
                        tracing::error!(
                            source_peer_addr = %peer,
                            error = %e,
                            "TLS handshake failed"
                        );
                        return;
                    }
                };
                let peer_certs = agent_gateway::proxy::extract_peer_certs(tls_stream.get_ref().1);
                let service = svc.make_service(peer_certs, peer);
                let io = hyper_util::rt::TokioIo::new(tls_stream);
                let _ = hyper_util::server::conn::auto::Builder::new(
                    hyper_util::rt::TokioExecutor::new(),
                )
                .http2_only()
                .serve_connection_with_upgrades(io, service)
                .await;
            });
        }
    });

    (addr, ServerGuard { task })
}

/// Connect an HTTP/2 mTLS client to the proxy. Returns a SendRequest handle.
pub async fn connect_client(
    proxy_addr: SocketAddr,
    pki: &TestPki,
) -> hyper::client::conn::http2::SendRequest<http_body_util::Empty<bytes::Bytes>> {
    install_test_crypto_provider();

    let mut ca_store = rustls::RootCertStore::empty();
    ca_store.add(pki.ca_cert_der()).unwrap();

    let mut client_config = rustls::ClientConfig::builder()
        .with_root_certificates(ca_store)
        .with_client_auth_cert(pki.client_cert_chain(), pki.client_key_der())
        .unwrap();
    client_config.alpn_protocols = vec![b"h2".to_vec()];

    let connector = tokio_rustls::TlsConnector::from(Arc::new(client_config));
    let tcp = tokio::net::TcpStream::connect(proxy_addr).await.unwrap();
    let server_name = rustls_pki_types::ServerName::try_from("localhost").unwrap();
    let tls_stream = connector.connect(server_name, tcp).await.unwrap();

    let io = hyper_util::rt::TokioIo::new(tls_stream);
    let (send_request, connection) =
        hyper::client::conn::http2::handshake(hyper_util::rt::TokioExecutor::new(), io)
            .await
            .unwrap();

    tokio::spawn(async move {
        let _ = connection.await;
    });

    send_request
}

/// Try to connect, perform HTTP/2 handshake, and send a CONNECT request using
/// the given client TLS config. Returns `Ok(response)` if everything succeeds,
/// or `Err` if any step (TLS, h2, request) fails. Used for negative mTLS tests.
pub async fn try_request_with_tls_config(
    proxy_addr: SocketAddr,
    client_config: rustls::ClientConfig,
    dest_authority: &str,
) -> Result<hyper::Response<hyper::body::Incoming>, Box<dyn std::error::Error + Send + Sync>> {
    install_test_crypto_provider();

    let connector = tokio_rustls::TlsConnector::from(Arc::new(client_config));
    let tcp = tokio::net::TcpStream::connect(proxy_addr).await?;
    let server_name = rustls_pki_types::ServerName::try_from("localhost").unwrap();
    let tls_stream = connector.connect(server_name, tcp).await?;

    let io = hyper_util::rt::TokioIo::new(tls_stream);
    let (mut send_request, connection) =
        hyper::client::conn::http2::handshake(hyper_util::rt::TokioExecutor::new(), io).await?;

    tokio::spawn(async move {
        let _ = connection.await;
    });

    let req = hyper::Request::connect(dest_authority)
        .body(http_body_util::Empty::<bytes::Bytes>::new())
        .unwrap();
    let resp = send_request.send_request(req).await?;
    Ok(resp)
}

/// Connect to the proxy with a custom client cert/key (for extension mismatch tests).
pub async fn connect_client_with_cert(
    proxy_addr: SocketAddr,
    ca_cert: CertificateDer<'static>,
    client_cert_chain: Vec<CertificateDer<'static>>,
    client_key: PrivateKeyDer<'static>,
) -> hyper::client::conn::http2::SendRequest<http_body_util::Empty<bytes::Bytes>> {
    install_test_crypto_provider();

    let mut ca_store = rustls::RootCertStore::empty();
    ca_store.add(ca_cert).unwrap();

    let mut client_config = rustls::ClientConfig::builder()
        .with_root_certificates(ca_store)
        .with_client_auth_cert(client_cert_chain, client_key)
        .unwrap();
    client_config.alpn_protocols = vec![b"h2".to_vec()];

    let connector = tokio_rustls::TlsConnector::from(Arc::new(client_config));
    let tcp = tokio::net::TcpStream::connect(proxy_addr).await.unwrap();
    let server_name = rustls_pki_types::ServerName::try_from("localhost").unwrap();
    let tls_stream = connector.connect(server_name, tcp).await.unwrap();

    let io = hyper_util::rt::TokioIo::new(tls_stream);
    let (send_request, connection) =
        hyper::client::conn::http2::handshake(hyper_util::rt::TokioExecutor::new(), io)
            .await
            .unwrap();

    tokio::spawn(async move {
        let _ = connection.await;
    });

    send_request
}

/// Start a TCP echo server. Returns its address and a guard that stops it.
pub async fn start_echo_server() -> (SocketAddr, ServerGuard) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();

    let task = tokio::spawn(async move {
        while let Ok((mut stream, _)) = listener.accept().await {
            tokio::spawn(async move {
                let mut buf = [0u8; 4096];
                loop {
                    let n = match stream.read(&mut buf).await {
                        Ok(0) | Err(_) => break,
                        Ok(n) => n,
                    };
                    if stream.write_all(&buf[..n]).await.is_err() {
                        break;
                    }
                }
            });
        }
    });

    (addr, ServerGuard { task })
}

/// Bind a TCP listener and immediately drop it, returning the port.
/// This guarantees the port is unoccupied and nothing is listening on it.
pub async fn allocate_closed_port() -> u16 {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    drop(listener);
    port
}