cellos-supervisor 0.5.1

//! SEC-21 Phase 3h.2 / T2.B Slot **A6** — DoT/DoH/DoQ upstream support.
//!
//! Phase 1's now-deleted `super::forward_upstream` only spoke `do53-udp`. This
//! module generalises that single hot path to the full set of upstream transports the
//! contract layer (`DnsAuthority::resolvers[].protocol`) admits:
//!
//! | Variant       | Wire        | Status (DNS-DOH-1..3 / DNS-DOQ-1..2)              |
//! |---------------|-------------|----------------------------------------------------|
//! | `Do53Udp`     | UDP/53      | Default. Byte-for-byte the legacy forward path.    |
//! | `Do53Tcp`     | TCP/53      | Length-prefixed two-byte framing over plain TCP.   |
//! | `Dot`         | TLS/853     | Length-prefixed framing over rustls 0.23 TLS.      |
//! | `Doh`         | HTTPS/443   | reqwest POST `application/dns-message` (RFC 8484). |
//! | `Doq`         | QUIC/853    | quinn 0.11 bidi stream (RFC 9250).                 |
//!
//! ## DoH / DoQ configuration shape
//!
//! The original DNS-DOH-1 ticket proposed widening `UpstreamTransport` into a
//! data-carrying enum (`Doh { url }`, `Doq { server, port }`). That would
//! ripple through ~10 call sites in `supervisor.rs` and the test surface that
//! pattern-match the enum as `Copy` unit variants. Instead, DoH/DoQ
//! configuration rides on [`UpstreamExtras`] alongside the existing DoT
//! overrides — the enum stays `Copy`, the production composition root keeps
//! the same shape, and the env-var protocol (`CELLOS_DNS_UPSTREAM_DOH_URL`,
//! `CELLOS_DNS_UPSTREAM_DOQ_SERVER`, `CELLOS_DNS_UPSTREAM_DOQ_PORT`) wires
//! through unchanged. The downside is that operators selecting DoH/DoQ with
//! no extras get a sensible default (Cloudflare `1.1.1.1`) rather than an
//! eager refusal; the upside is that all transports share one config seam.
//!
//! ## Why DoT/DoH/DoQ are split into typed variants now
//!
//! The supervisor's contract layer already accepts a `protocol` field on each
//! declared resolver — there's no schema migration in this slot, only a
//! widening of the dataplane forward path. Operators who declare a non-`do53`
//! protocol today get a SERVFAIL with an explicit `TransportNotEnabled`
//! discriminator (D9 in the dispatch ledger) instead of silently downgrading
//! to plain UDP. The toggle parses, the audit event tells the truth.
//!
//! ## DoT (the actually-implemented non-do53 path)
//!
//! Hickory's DoT support exists but is feature-gated and pulls in the full
//! `dns-over-rustls` stack — for one transport that's heavyweight. We
//! implement DoT directly on rustls 0.23 + tokio-rustls 0.26:
//!
//! 1. TCP connect to the resolver's `host:port` (port defaulted to 853 if the
//!    operator did not specify one — RFC 7858).
//! 2. TLS handshake using a `RootCertStore` populated from
//!    `webpki-roots = "1"` (Mozilla's bundled CA set; same trust set hickory
//!    uses internally).
//! 3. RFC 7858 framing: each query is prefixed with its 2-byte big-endian
//!    length, and the response is parsed in the same shape.
//!
//! The `forward()` entry point is synchronous (the proxy hot path runs on a
//! blocking thread; see [`super::run_one_shot`]) so we drive the async DoT /
//! DoH / DoQ paths via `tokio::runtime::Handle::block_on` wrapped in
//! `tokio::task::block_in_place`. If no runtime is in scope we return
//! [`UpstreamError::NoTokioRuntime`] and the caller fail-safes to SERVFAIL.
//!
//! ## DoH (RFC 8484, DNS-DOH-1..3)
//!
//! `forward_doh()` POSTs the raw DNS wire-format query to the operator-
//! supplied URL (default `https://1.1.1.1/dns-query`) with
//! `Content-Type: application/dns-message` + `Accept: application/dns-message`,
//! then reads the response body as the raw DNS wire response. reqwest's
//! client is constructed per call (no pool reuse) — keeps the path stateless
//! to match DoT and the rest of the proxy.
//!
//! ## DoQ (RFC 9250, DNS-DOQ-1..2)
//!
//! `forward_doq()` opens a QUIC connection to the operator-supplied
//! `server:port` (default `1.1.1.1:853`), opens a bidirectional stream, writes
//! the RFC 9250 2-byte-length-prefixed query, finishes the send side, then
//! reads the length-prefixed response. ALPN is `doq` per RFC 9250 §4.1.1.
//! rustls 0.23 + webpki-roots is the trust set (same as DoT).

use std::io;
use std::net::{SocketAddr, UdpSocket};
use std::sync::Arc;
use std::time::Duration;

use rustls::pki_types::ServerName;
use rustls::{ClientConfig, RootCertStore};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio_rustls::TlsConnector;

/// Maximum DNS payload the proxy will pass through. Mirrors the
/// `super::MAX_UDP_PAYLOAD` ceiling so callers can safely hand us a 1500-byte
/// scratch buffer.
const MAX_PAYLOAD: usize = 1500;

/// Default DoT port (RFC 7858).
const DEFAULT_DOT_PORT: u16 = 853;

/// Default DoH endpoint (Cloudflare). Picked when an operator selects DoH
/// without supplying `CELLOS_DNS_UPSTREAM_DOH_URL`. RFC 8484 §4.1 documents
/// the canonical `/dns-query` path; Cloudflare answers on both the apex IP
/// and `cloudflare-dns.com`.
const DEFAULT_DOH_URL: &str = "https://1.1.1.1/dns-query";

/// Default DoQ host + port (Cloudflare on 1.1.1.1:853, RFC 9250 §4).
const DEFAULT_DOQ_SERVER: &str = "1.1.1.1";
const DEFAULT_DOQ_PORT: u16 = 853;

/// Upstream transport selector. Production callers populate this from
/// `DnsAuthority::resolvers[].protocol`; operators in dev / CI override via
/// the `CELLOS_DNS_UPSTREAM_TRANSPORT` env var (see [`UpstreamTransport::from_env`]).
///
/// `Default = Do53Udp` so the existing UDP-only behaviour is the no-op
/// upgrade path — adding the `transport` field to [`super::DnsProxyConfig`]
/// does not change behaviour for cells that don't explicitly opt into a
/// non-default value.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum UpstreamTransport {
    /// Plain UDP/53 (legacy default, byte-identical to pre-A6 forward path).
    #[default]
    Do53Udp,
    /// Plain TCP/53 with the conventional 2-byte length prefix.
    Do53Tcp,
    /// DNS-over-TLS (RFC 7858). Length-prefixed framing inside a TLS 1.2/1.3 tunnel.
    Dot,
    /// DNS-over-HTTPS (RFC 8484). reqwest POST with
    /// `application/dns-message` framing. URL configured via
    /// [`UpstreamExtras::doh_url`] (default Cloudflare).
    Doh,
    /// DNS-over-QUIC (RFC 9250). quinn 0.11 bidi stream with the same
    /// 2-byte length prefix as DoT. Server + port configured via
    /// [`UpstreamExtras::doq_server`] / [`UpstreamExtras::doq_port`]
    /// (default `1.1.1.1:853`).
    Doq,
}

impl UpstreamTransport {
    /// Parse a textual transport selector. Case-insensitive; accepts the
    /// canonical contract names (`do53-udp`, `do53-tcp`, `dot`, `doh`, `doq`)
    /// AND the common aliases operators reach for (`udp`, `tcp`, `tls`,
    /// `https`, `quic`).
    ///
    /// Returns `None` for unrecognised input. The caller is expected to
    /// fail-closed (refuse to construct the cell) rather than silently
    /// fall back to UDP, so we don't pretend to know what the operator meant.
    pub fn parse(s: &str) -> Option<Self> {
        match s.trim().to_ascii_lowercase().as_str() {
            "do53-udp" | "do53udp" | "udp" | "" => Some(Self::Do53Udp),
            "do53-tcp" | "do53tcp" | "tcp" => Some(Self::Do53Tcp),
            "dot" | "tls" | "do-tls" => Some(Self::Dot),
            "doh" | "https" | "do-https" => Some(Self::Doh),
            "doq" | "quic" | "do-quic" => Some(Self::Doq),
            _ => None,
        }
    }

    /// Read the upstream transport selector from the process environment and
    /// parse it. Two env var names are accepted (in priority order):
    ///
    /// 1. `CELLOS_DNS_UPSTREAM_PROTOCOL` — the contract-aligned name
    ///    (matches `DnsAuthority::resolvers[].protocol`). Preferred.
    /// 2. `CELLOS_DNS_UPSTREAM_TRANSPORT` — the original legacy name,
    ///    preserved so existing operator scripts and integration tests
    ///    keep working byte-for-byte.
    ///
    /// Unset returns `Some(Do53Udp)` (the default). Set-but-unparseable
    /// returns `None` so the caller can refuse to start.
    pub fn from_env() -> Option<Self> {
        for var in [
            "CELLOS_DNS_UPSTREAM_PROTOCOL",
            "CELLOS_DNS_UPSTREAM_TRANSPORT",
        ] {
            match std::env::var(var) {
                Ok(s) if !s.trim().is_empty() => return Self::parse(&s),
                _ => continue,
            }
        }
        Some(Self::Do53Udp)
    }

    /// Stable string label suitable for stamping into events / logs. Used in
    /// [`UpstreamError::TransportNotEnabled`]'s `Display`.
    pub fn label(self) -> &'static str {
        match self {
            Self::Do53Udp => "do53-udp",
            Self::Do53Tcp => "do53-tcp",
            Self::Dot => "dot",
            Self::Doh => "doh",
            Self::Doq => "doq",
        }
    }
}

/// Optional per-transport extras the caller may supply alongside
/// [`UpstreamTransport`]. Populated by [`super::DnsProxyConfig::upstream_extras`]
/// in production; tests construct via `Default::default()` and then set only
/// the fields they care about.
///
/// All fields are optional so the `do53-udp` hot path doesn't require
/// any extras to function — unused fields incur zero cost.
#[derive(Debug, Clone, Default)]
pub struct UpstreamExtras {
    /// SNI hostname presented during the DoT TLS handshake. When `None`,
    /// rustls is given the resolver's IP literal — most public DoT resolvers
    /// (1.1.1.1 / 8.8.8.8) ship a cert that covers both the hostname and
    /// the IP, but operators with a private resolver will need to set this
    /// to the cert's CN/SAN.
    pub dot_sni: Option<String>,
    /// Operator-supplied DoT server host. When set (typically populated from
    /// the `CELLOS_DNS_UPSTREAM_DOT_SERVER` env var), the supervisor pre-resolves
    /// this to a `SocketAddr` and substitutes the proxy's `upstream_addr`,
    /// so the DoT roundtrip targets the operator's choice rather than the
    /// spec's do53 resolver. `None` falls back to the spec resolver's IP.
    ///
    /// Plain string here (not pre-resolved) so the config struct stays
    /// transport-agnostic and the resolution step lives in the composition
    /// root where DNS bootstrap is allowed.
    pub dot_server: Option<String>,
    /// Operator-supplied DoT port (RFC 7858 default = 853). Paired with
    /// [`Self::dot_server`]; when `None` the supervisor defaults to 853.
    pub dot_port: Option<u16>,
    /// DNS-DOH-2 — operator-supplied DoH endpoint URL. When `None` the
    /// DoH forward path defaults to [`DEFAULT_DOH_URL`]. Sourced from
    /// `CELLOS_DNS_UPSTREAM_DOH_URL` in production. Must be a full
    /// `https://…` URL including the `/dns-query` path; reqwest validates
    /// the scheme on first call and surfaces a typed
    /// [`UpstreamError::Io(InvalidInput)`] if it's malformed.
    pub doh_url: Option<String>,
    /// DNS-DOQ-2 — operator-supplied DoQ server (IP literal or hostname).
    /// `None` → default `1.1.1.1`. Sourced from
    /// `CELLOS_DNS_UPSTREAM_DOQ_SERVER`. Hostnames are passed through to
    /// `tokio::net::lookup_host` (which uses the OS resolver, NOT the
    /// supervisor's bootstrap path) — operators should prefer IP literals
    /// for the same reason DoT does (see [`parse_dot_target`]).
    pub doq_server: Option<String>,
    /// DNS-DOQ-2 — operator-supplied DoQ port. `None` → default `853`
    /// (RFC 9250). Sourced from `CELLOS_DNS_UPSTREAM_DOQ_PORT`.
    pub doq_port: Option<u16>,
}

impl UpstreamExtras {
    /// Read the DoT-specific operator overrides from the process environment.
    ///
    /// Recognised env vars:
    ///   - `CELLOS_DNS_UPSTREAM_DOT_SERVER` — host (IP literal or hostname)
    ///     for the DoT upstream. Default unset (caller falls back to the
    ///     spec resolver's IP, or to `1.1.1.1` if the caller has no spec
    ///     to fall back on).
    ///   - `CELLOS_DNS_UPSTREAM_DOT_PORT` — TCP port. Default unset
    ///     (caller falls back to `853` per RFC 7858).
    ///   - `CELLOS_DNS_UPSTREAM_DOT_SNI` — explicit SNI hostname. Default
    ///     unset (rustls receives the resolver's IP literal as ServerName).
    ///
    /// Unparseable port values are silently ignored (the field stays `None`)
    /// rather than failing — operators get the default behaviour instead of
    /// a refused cell. Strict parsing is the supervisor's job at the composition
    /// site if it wants to gate startup on a typo.
    pub fn from_env() -> Self {
        let dot_server = std::env::var("CELLOS_DNS_UPSTREAM_DOT_SERVER")
            .ok()
            .filter(|s| !s.trim().is_empty());
        let dot_port = std::env::var("CELLOS_DNS_UPSTREAM_DOT_PORT")
            .ok()
            .and_then(|s| s.trim().parse::<u16>().ok());
        let dot_sni = std::env::var("CELLOS_DNS_UPSTREAM_DOT_SNI")
            .ok()
            .filter(|s| !s.trim().is_empty());
        // DNS-DOH-2 / DNS-DOQ-2 — DoH/DoQ env wiring. Mirrors the DoT pattern:
        // unparseable port is silently treated as unset rather than failing,
        // so a typo lands the operator on the default (typed config error
        // surfaces later via reqwest URL validation in the DoH case).
        let doh_url = std::env::var("CELLOS_DNS_UPSTREAM_DOH_URL")
            .ok()
            .filter(|s| !s.trim().is_empty());
        let doq_server = std::env::var("CELLOS_DNS_UPSTREAM_DOQ_SERVER")
            .ok()
            .filter(|s| !s.trim().is_empty());
        let doq_port = std::env::var("CELLOS_DNS_UPSTREAM_DOQ_PORT")
            .ok()
            .and_then(|s| s.trim().parse::<u16>().ok());
        Self {
            dot_sni,
            dot_server,
            dot_port,
            doh_url,
            doq_server,
            doq_port,
        }
    }
}

/// Typed upstream-forward error.
///
/// Mapped at the [`super::run_one_shot`] call site to a SERVFAIL response
/// for the workload plus a `dns_query` CloudEvent stamped
/// `reasonCode: upstream_failure`. The variant discriminator is stable so
/// follow-up slots can refine the event's reason text without re-shaping
/// the matrix.
#[derive(Debug)]
pub enum UpstreamError {
    /// Upstream did not answer within the configured budget.
    Timeout,
    /// Lower-level I/O failure (refused, unreachable, transport reset).
    Io(io::Error),
    /// Operator selected a transport this build does not implement —
    /// typed SERVFAIL discriminator (D9 in the dispatch ledger).
    TransportNotEnabled(UpstreamTransport),
    /// rustls handshake failed (cert chain rejected, SNI mismatch, protocol
    /// error). Distinct from [`Self::Io`] so triage can tell "the resolver
    /// answered but with a bad cert" from "the resolver didn't answer".
    TlsHandshake(String),
    /// Async transport invoked outside a tokio runtime — programming error
    /// in production (the supervisor always runs inside one) but explicit
    /// rather than panicking.
    NoTokioRuntime,
}

impl std::fmt::Display for UpstreamError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Timeout => write!(f, "upstream timeout"),
            Self::Io(e) => write!(f, "upstream io: {e}"),
            Self::TransportNotEnabled(t) => {
                write!(
                    f,
                    "upstream transport '{}' not enabled in this build",
                    t.label()
                )
            }
            Self::TlsHandshake(msg) => write!(f, "tls handshake: {msg}"),
            Self::NoTokioRuntime => write!(f, "no tokio runtime in scope for async upstream"),
        }
    }
}

impl std::error::Error for UpstreamError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Io(e) => Some(e),
            _ => None,
        }
    }
}

impl From<io::Error> for UpstreamError {
    fn from(e: io::Error) -> Self {
        if matches!(
            e.kind(),
            io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock
        ) {
            Self::Timeout
        } else {
            Self::Io(e)
        }
    }
}

/// Synchronous upstream forward — the single entry point the proxy hot path
/// calls. Dispatches on [`UpstreamTransport`].
///
/// `udp_socket` is the pre-bound UDP socket used for the `Do53Udp` path; for
/// every other transport it is ignored (we open a fresh transport-specific
/// connection per query — no connection pooling in this slot). `out_buf` is
/// filled with the upstream's response and the byte count returned; on error
/// the buffer is left untouched.
pub fn forward(
    transport: UpstreamTransport,
    udp_socket: &UdpSocket,
    upstream: SocketAddr,
    query: &[u8],
    out_buf: &mut [u8],
    timeout: Duration,
    extras: &UpstreamExtras,
) -> Result<usize, UpstreamError> {
    if query.len() > MAX_PAYLOAD {
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidInput,
            "query exceeds MAX_PAYLOAD",
        )));
    }
    match transport {
        UpstreamTransport::Do53Udp => forward_udp(udp_socket, upstream, query, out_buf, timeout),
        UpstreamTransport::Do53Tcp => forward_tcp(upstream, query, out_buf, timeout),
        UpstreamTransport::Dot => forward_dot(upstream, query, out_buf, timeout, extras),
        UpstreamTransport::Doh => forward_doh(query, out_buf, timeout, extras),
        UpstreamTransport::Doq => forward_doq(query, out_buf, timeout, extras),
    }
}

fn forward_udp(
    upstream: &UdpSocket,
    addr: SocketAddr,
    query: &[u8],
    buf: &mut [u8],
    timeout: Duration,
) -> Result<usize, UpstreamError> {
    upstream.send_to(query, addr)?;
    upstream.set_read_timeout(Some(timeout))?;
    let deadline = std::time::Instant::now() + timeout;
    loop {
        match upstream.recv_from(buf) {
            Ok((n, _peer)) => return Ok(n),
            Err(e)
                if matches!(
                    e.kind(),
                    io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
                ) =>
            {
                if std::time::Instant::now() >= deadline {
                    return Err(UpstreamError::Timeout);
                }
            }
            Err(e) if matches!(e.kind(), io::ErrorKind::Interrupted) => continue,
            Err(e) => return Err(UpstreamError::Io(e)),
        }
    }
}

fn forward_tcp(
    upstream: SocketAddr,
    query: &[u8],
    buf: &mut [u8],
    timeout: Duration,
) -> Result<usize, UpstreamError> {
    use std::io::{Read, Write};
    use std::net::TcpStream;

    let mut stream = TcpStream::connect_timeout(&upstream, timeout)?;
    stream.set_read_timeout(Some(timeout))?;
    stream.set_write_timeout(Some(timeout))?;

    let len = query.len() as u16;
    stream.write_all(&len.to_be_bytes())?;
    stream.write_all(query)?;

    let mut len_buf = [0u8; 2];
    stream.read_exact(&mut len_buf)?;
    let resp_len = u16::from_be_bytes(len_buf) as usize;
    if resp_len > buf.len() {
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            "tcp response exceeds out buffer",
        )));
    }
    stream.read_exact(&mut buf[..resp_len])?;
    Ok(resp_len)
}

fn forward_dot(
    upstream: SocketAddr,
    query: &[u8],
    buf: &mut [u8],
    timeout: Duration,
    extras: &UpstreamExtras,
) -> Result<usize, UpstreamError> {
    let handle =
        tokio::runtime::Handle::try_current().map_err(|_| UpstreamError::NoTokioRuntime)?;

    // Resolve target SocketAddr. Precedence:
    //   1. `extras.dot_server` (operator override, typically from
    //      `CELLOS_DNS_UPSTREAM_DOT_SERVER`) — interpreted as either a
    //      bare IP literal or a `host:port` pair. If the parse fails, we
    //      surface a typed `Io(InvalidInput)` so the caller fail-safes to
    //      SERVFAIL rather than silently routing to the wrong resolver.
    //   2. The `upstream` SocketAddr the caller supplied (spec resolver).
    //   3. Port defaults to 853 (RFC 7858) when zero.
    let mut target = match extras.dot_server.as_deref() {
        Some(host) => parse_dot_target(host, extras.dot_port).map_err(|e| {
            UpstreamError::Io(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("CELLOS_DNS_UPSTREAM_DOT_SERVER='{host}' did not parse: {e}"),
            ))
        })?,
        None => {
            let mut t = upstream;
            if let Some(p) = extras.dot_port {
                t.set_port(p);
            }
            t
        }
    };
    if target.port() == 0 {
        target.set_port(DEFAULT_DOT_PORT);
    }

    let sni = extras.dot_sni.clone();
    let query = query.to_vec();
    let buf_len = buf.len();
    // `forward` is a sync entry point so it can sit on the proxy's
    // blocking I/O thread. The DoT path needs an async runtime for the
    // TLS handshake. `block_in_place` is tokio's official escape hatch:
    // it tells the multi-thread scheduler "this worker is about to
    // block; redistribute the other tasks", which keeps `block_on`
    // from panicking with `Cannot start a runtime from within a runtime`
    // when callers happen to dispatch us from inside an async context
    // (notably integration tests like supervisor_dns_proxy_dot.rs).
    let result: Result<Vec<u8>, UpstreamError> = tokio::task::block_in_place(|| {
        handle.block_on(async move {
            match tokio::time::timeout(timeout, dot_roundtrip(target, &query, &sni, buf_len)).await
            {
                Ok(inner) => inner,
                Err(_) => Err(UpstreamError::Timeout),
            }
        })
    });
    let resp = result?;
    if resp.len() > buf.len() {
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            "dot response exceeds out buffer",
        )));
    }
    buf[..resp.len()].copy_from_slice(&resp);
    Ok(resp.len())
}

/// Parse the operator-supplied DoT server string into a `SocketAddr`.
///
/// Accepted shapes:
///   - bare IPv4 literal (`1.1.1.1`)
///   - bare IPv6 literal (`2606:4700:4700::1111`)
///   - `IPv4:port` (`1.1.1.1:853`)
///   - `[IPv6]:port` (`[2606:4700:4700::1111]:853`)
///
/// Hostnames are explicitly rejected: the in-netns dataplane has no DNS
/// bootstrap path, and accepting a hostname here would either require a
/// circular DNS lookup or a silent fallback to the system resolver outside
/// the cell's authority. The composition root must pre-resolve any hostname
/// to an IP literal before populating `UpstreamExtras::dot_server`.
fn parse_dot_target(host: &str, port_override: Option<u16>) -> Result<SocketAddr, String> {
    let trimmed = host.trim();
    if trimmed.is_empty() {
        return Err("empty string".to_string());
    }

    // If the operator wrote `host:port`, prefer that port over the override.
    // SocketAddr::from_str handles both `IPv4:port` and `[IPv6]:port`.
    if let Ok(sa) = trimmed.parse::<SocketAddr>() {
        return Ok(sa);
    }

    // Otherwise treat the string as a bare IP literal and bolt on the port.
    let port = port_override.unwrap_or(DEFAULT_DOT_PORT);
    let ip: std::net::IpAddr = trimmed.parse().map_err(|_| {
        format!(
            "'{trimmed}' is not an IP literal (hostnames must be pre-resolved by the supervisor)"
        )
    })?;
    Ok(SocketAddr::new(ip, port))
}

/// Async DoT roundtrip — TCP connect, TLS handshake, length-prefixed query,
/// length-prefixed response.
async fn dot_roundtrip(
    target: SocketAddr,
    query: &[u8],
    sni: &Option<String>,
    out_cap: usize,
) -> Result<Vec<u8>, UpstreamError> {
    let config = build_dot_client_config();
    let connector = TlsConnector::from(Arc::new(config));

    let tcp = tokio::net::TcpStream::connect(target).await?;

    // Build the SNI ServerName. Operator-supplied SNI wins; otherwise we
    // fall back to the resolver's IP literal — rustls 0.23's `pki_types`
    // accepts `IpAddress` server names so the handshake will succeed
    // against resolvers whose certs cover the IP (1.1.1.1, 9.9.9.9, etc.).
    let server_name: ServerName<'static> = match sni {
        Some(host) if !host.is_empty() => ServerName::try_from(host.clone())
            .map_err(|e| UpstreamError::TlsHandshake(format!("invalid sni '{host}': {e}")))?,
        _ => ServerName::IpAddress(target.ip().into()),
    };

    let mut tls = connector
        .connect(server_name, tcp)
        .await
        .map_err(|e| UpstreamError::TlsHandshake(format!("{e}")))?;

    // RFC 7858 framing — 2-byte big-endian length prefix on both sides.
    let len = query.len() as u16;
    tls.write_all(&len.to_be_bytes()).await?;
    tls.write_all(query).await?;
    tls.flush().await?;

    let mut len_buf = [0u8; 2];
    tls.read_exact(&mut len_buf).await?;
    let resp_len = u16::from_be_bytes(len_buf) as usize;
    if resp_len > out_cap {
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            "dot response exceeds out buffer",
        )));
    }
    let mut resp = vec![0u8; resp_len];
    tls.read_exact(&mut resp).await?;
    Ok(resp)
}

/// DNS-DOH-1 — DoH (RFC 8484) forward path.
///
/// Synchronous entry point following the same `block_in_place` +
/// `Handle::block_on` pattern as [`forward_dot`]: the proxy hot path is a
/// blocking thread inside a tokio runtime, so we redistribute scheduler
/// work and drive an async reqwest POST. Returns the response wire bytes
/// or a typed [`UpstreamError`].
///
/// Endpoint resolution precedence:
///   1. `extras.doh_url` (operator override, typically from
///      `CELLOS_DNS_UPSTREAM_DOH_URL`).
///   2. [`DEFAULT_DOH_URL`] = Cloudflare's `https://1.1.1.1/dns-query`.
fn forward_doh(
    query: &[u8],
    buf: &mut [u8],
    timeout: Duration,
    extras: &UpstreamExtras,
) -> Result<usize, UpstreamError> {
    let handle =
        tokio::runtime::Handle::try_current().map_err(|_| UpstreamError::NoTokioRuntime)?;

    let url = extras
        .doh_url
        .clone()
        .unwrap_or_else(|| DEFAULT_DOH_URL.to_string());
    let query = query.to_vec();
    let buf_len = buf.len();

    let result: Result<Vec<u8>, UpstreamError> = tokio::task::block_in_place(|| {
        handle.block_on(async move {
            match tokio::time::timeout(timeout, doh_roundtrip(&url, &query, timeout, buf_len)).await
            {
                Ok(inner) => inner,
                Err(_) => Err(UpstreamError::Timeout),
            }
        })
    });
    let resp = result?;
    if resp.len() > buf.len() {
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            "doh response exceeds out buffer",
        )));
    }
    buf[..resp.len()].copy_from_slice(&resp);
    Ok(resp.len())
}

/// Async DoH roundtrip — single POST with the RFC 8484 content-type contract.
async fn doh_roundtrip(
    url: &str,
    query: &[u8],
    timeout: Duration,
    out_cap: usize,
) -> Result<Vec<u8>, UpstreamError> {
    // reqwest client constructed per call. No pool reuse in this slot — the
    // upstream forward path is already on the slow path (every cache miss
    // pays the TLS handshake), and connection reuse would require a static
    // client + careful cross-cell isolation. We keep it stateless to match
    // DoT.
    let client = reqwest::Client::builder()
        .timeout(timeout)
        .build()
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doh client build: {e}"))))?;

    let resp = client
        .post(url)
        .header(reqwest::header::CONTENT_TYPE, "application/dns-message")
        .header(reqwest::header::ACCEPT, "application/dns-message")
        .body(query.to_vec())
        .send()
        .await
        .map_err(|e| {
            // reqwest's timeout-error variant maps to our `Timeout`; everything
            // else (DNS failure on the URL host, TLS handshake, connection
            // refused) lands on `Io` so triage can tell "the resolver answered"
            // from "the resolver didn't answer" via the discriminator.
            if e.is_timeout() {
                UpstreamError::Timeout
            } else {
                UpstreamError::Io(io::Error::other(format!("doh request: {e}")))
            }
        })?;

    if !resp.status().is_success() {
        return Err(UpstreamError::Io(io::Error::other(format!(
            "doh upstream returned HTTP {}",
            resp.status()
        ))));
    }

    let bytes = resp
        .bytes()
        .await
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doh body: {e}"))))?;
    if bytes.len() > out_cap {
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            "doh response exceeds out buffer",
        )));
    }
    Ok(bytes.to_vec())
}

/// DNS-DOQ-1 — DoQ (RFC 9250) forward path.
///
/// Synchronous entry point. QUIC connection + bidi stream + 2-byte
/// length-prefixed query/response. Defaults to `1.1.1.1:853` if no operator
/// override is present.
fn forward_doq(
    query: &[u8],
    buf: &mut [u8],
    timeout: Duration,
    extras: &UpstreamExtras,
) -> Result<usize, UpstreamError> {
    let handle =
        tokio::runtime::Handle::try_current().map_err(|_| UpstreamError::NoTokioRuntime)?;

    let server = extras
        .doq_server
        .clone()
        .unwrap_or_else(|| DEFAULT_DOQ_SERVER.to_string());
    let port = extras.doq_port.unwrap_or(DEFAULT_DOQ_PORT);
    let query = query.to_vec();
    let buf_len = buf.len();

    let result: Result<Vec<u8>, UpstreamError> = tokio::task::block_in_place(|| {
        handle.block_on(async move {
            match tokio::time::timeout(timeout, doq_roundtrip(&server, port, &query, buf_len)).await
            {
                Ok(inner) => inner,
                Err(_) => Err(UpstreamError::Timeout),
            }
        })
    });
    let resp = result?;
    if resp.len() > buf.len() {
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            "doq response exceeds out buffer",
        )));
    }
    buf[..resp.len()].copy_from_slice(&resp);
    Ok(resp.len())
}

/// Async DoQ roundtrip — quinn 0.11 client endpoint, bidi stream, RFC 9250
/// length-prefixed framing.
async fn doq_roundtrip(
    server: &str,
    port: u16,
    query: &[u8],
    out_cap: usize,
) -> Result<Vec<u8>, UpstreamError> {
    use std::net::IpAddr;

    // Resolve `server` to an `IpAddr` + keep the original string for SNI. If
    // the operator supplied an IP literal, we use that for both the QUIC
    // target and as a fallback SNI; if it's a hostname, we rely on
    // `tokio::net::lookup_host` for the address and pass the hostname through
    // as the rustls SNI (matching standard browser-style DoQ resolvers).
    let (target_addr, sni): (SocketAddr, ServerName<'static>) =
        if let Ok(ip) = server.parse::<IpAddr>() {
            let sa = SocketAddr::new(ip, port);
            let sni = ServerName::IpAddress(match ip {
                IpAddr::V4(v4) => rustls::pki_types::IpAddr::V4(v4.into()),
                IpAddr::V6(v6) => rustls::pki_types::IpAddr::V6(v6.into()),
            });
            (sa, sni)
        } else {
            // Hostname path — first address wins. This uses the OS resolver
            // (NOT the supervisor's bootstrap path), so operators in a
            // sealed netns should prefer IP literals.
            let mut iter = tokio::net::lookup_host((server, port)).await.map_err(|e| {
                UpstreamError::Io(io::Error::new(
                    e.kind(),
                    format!("doq lookup '{server}': {e}"),
                ))
            })?;
            let sa = iter.next().ok_or_else(|| {
                UpstreamError::Io(io::Error::new(
                    io::ErrorKind::AddrNotAvailable,
                    format!("doq lookup '{server}' returned no addresses"),
                ))
            })?;
            let sni = ServerName::try_from(server.to_string())
                .map_err(|e| UpstreamError::TlsHandshake(format!("invalid sni '{server}': {e}")))?;
            (sa, sni)
        };

    // Bind a local UDP socket for the QUIC client. `0.0.0.0:0` lets the OS
    // pick an ephemeral port; the QUIC endpoint owns this socket for the
    // life of the connection.
    let bind_addr: SocketAddr = match target_addr {
        SocketAddr::V4(_) => "0.0.0.0:0".parse().unwrap(),
        SocketAddr::V6(_) => "[::]:0".parse().unwrap(),
    };
    let mut endpoint = quinn::Endpoint::client(bind_addr)
        .map_err(|e| UpstreamError::Io(io::Error::new(e.kind(), format!("doq endpoint: {e}"))))?;

    // rustls config with ALPN = "doq" (RFC 9250 §4.1.1).
    let mut roots = RootCertStore::empty();
    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
    let provider = Arc::new(rustls::crypto::ring::default_provider());
    let mut crypto = ClientConfig::builder_with_provider(provider)
        .with_safe_default_protocol_versions()
        .map_err(|e| UpstreamError::TlsHandshake(format!("doq rustls protocols: {e}")))?
        .with_root_certificates(roots)
        .with_no_client_auth();
    crypto.alpn_protocols = vec![b"doq".to_vec()];

    let quic_crypto = quinn::crypto::rustls::QuicClientConfig::try_from(crypto)
        .map_err(|e| UpstreamError::TlsHandshake(format!("doq quic crypto: {e}")))?;
    let client_config = quinn::ClientConfig::new(Arc::new(quic_crypto));
    endpoint.set_default_client_config(client_config);

    // SNI server-name string — rustls wants `&str`, but quinn's `connect`
    // takes `&str` separately from the rustls config. We extract it back
    // from `sni` for the API call.
    let sni_str: String = match &sni {
        ServerName::DnsName(d) => d.as_ref().to_string(),
        ServerName::IpAddress(_) => server.to_string(),
        _ => server.to_string(),
    };

    let connecting = endpoint
        .connect(target_addr, &sni_str)
        .map_err(|e| UpstreamError::TlsHandshake(format!("doq connect: {e}")))?;
    let connection = connecting
        .await
        .map_err(|e| UpstreamError::TlsHandshake(format!("doq handshake: {e}")))?;

    let (mut send, mut recv) = connection
        .open_bi()
        .await
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doq open_bi: {e}"))))?;

    // RFC 9250 §4.2.1 — 2-byte big-endian length prefix on both query and
    // response. The query carries the same wire format as DNS-over-TCP.
    let len = query.len() as u16;
    send.write_all(&len.to_be_bytes())
        .await
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doq send len: {e}"))))?;
    send.write_all(query)
        .await
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doq send body: {e}"))))?;
    send.finish()
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doq finish: {e}"))))?;

    let mut len_buf = [0u8; 2];
    recv.read_exact(&mut len_buf)
        .await
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doq recv len: {e}"))))?;
    let resp_len = u16::from_be_bytes(len_buf) as usize;
    if resp_len > out_cap {
        // Close the connection gracefully before bailing — keeps the QUIC
        // peer's state machine clean.
        connection.close(0u32.into(), b"oversized");
        endpoint.wait_idle().await;
        return Err(UpstreamError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            "doq response exceeds out buffer",
        )));
    }
    let mut resp = vec![0u8; resp_len];
    recv.read_exact(&mut resp)
        .await
        .map_err(|e| UpstreamError::Io(io::Error::other(format!("doq recv body: {e}"))))?;

    connection.close(0u32.into(), b"done");
    endpoint.wait_idle().await;
    Ok(resp)
}

/// Build a default rustls 0.23 client config with Mozilla CA roots from
/// `webpki-roots = "1"`.
///
/// Explicitly threads the `ring` crypto provider so we don't depend on a
/// process-wide default (rustls 0.23's `ClientConfig::builder()` panics if
/// neither provider has been installed via `install_default()` and the
/// process picks neither feature unambiguously). `cellos-supervisor` pins
/// the `ring` feature in its Cargo.toml so this is the intended provider.
fn build_dot_client_config() -> ClientConfig {
    let mut roots = RootCertStore::empty();
    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
    let provider = Arc::new(rustls::crypto::ring::default_provider());
    ClientConfig::builder_with_provider(provider)
        .with_safe_default_protocol_versions()
        .expect("ring provider supports default rustls protocol versions")
        .with_root_certificates(roots)
        .with_no_client_auth()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;

    /// Tests in this module that mutate `CELLOS_DNS_UPSTREAM_*` env vars must
    /// hold this mutex for their entire duration — cargo runs tests in
    /// parallel by default and the env namespace is process-global.
    static ENV_LOCK: Mutex<()> = Mutex::new(());

    #[test]
    fn parse_canonical_names() {
        assert_eq!(
            UpstreamTransport::parse("do53-udp"),
            Some(UpstreamTransport::Do53Udp)
        );
        assert_eq!(
            UpstreamTransport::parse("do53-tcp"),
            Some(UpstreamTransport::Do53Tcp)
        );
        assert_eq!(
            UpstreamTransport::parse("dot"),
            Some(UpstreamTransport::Dot)
        );
        assert_eq!(
            UpstreamTransport::parse("doh"),
            Some(UpstreamTransport::Doh)
        );
        assert_eq!(
            UpstreamTransport::parse("doq"),
            Some(UpstreamTransport::Doq)
        );
    }

    #[test]
    fn parse_aliases_case_insensitive() {
        assert_eq!(
            UpstreamTransport::parse("UDP"),
            Some(UpstreamTransport::Do53Udp)
        );
        assert_eq!(
            UpstreamTransport::parse("TCP"),
            Some(UpstreamTransport::Do53Tcp)
        );
        assert_eq!(
            UpstreamTransport::parse("Tls"),
            Some(UpstreamTransport::Dot)
        );
        assert_eq!(
            UpstreamTransport::parse("HTTPS"),
            Some(UpstreamTransport::Doh)
        );
        assert_eq!(
            UpstreamTransport::parse("quic"),
            Some(UpstreamTransport::Doq)
        );
    }

    #[test]
    fn parse_rejects_unknown() {
        assert_eq!(UpstreamTransport::parse("dnscrypt"), None);
        assert_eq!(UpstreamTransport::parse("xxx"), None);
    }

    #[test]
    fn default_is_udp() {
        assert_eq!(UpstreamTransport::default(), UpstreamTransport::Do53Udp);
    }

    #[test]
    fn label_round_trips() {
        for t in [
            UpstreamTransport::Do53Udp,
            UpstreamTransport::Do53Tcp,
            UpstreamTransport::Dot,
            UpstreamTransport::Doh,
            UpstreamTransport::Doq,
        ] {
            assert_eq!(UpstreamTransport::parse(t.label()), Some(t));
        }
    }

    #[test]
    fn extras_from_env_reads_doh_url() {
        // DNS-DOH-2 — `CELLOS_DNS_UPSTREAM_DOH_URL` lands on extras.doh_url.
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let saved = std::env::var("CELLOS_DNS_UPSTREAM_DOH_URL").ok();
        unsafe {
            std::env::set_var(
                "CELLOS_DNS_UPSTREAM_DOH_URL",
                "https://cloudflare-dns.com/dns-query",
            );
        }
        let extras = UpstreamExtras::from_env();
        assert_eq!(
            extras.doh_url.as_deref(),
            Some("https://cloudflare-dns.com/dns-query")
        );
        unsafe {
            match saved {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_DOH_URL", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_DOH_URL"),
            }
        }
    }

    #[test]
    fn extras_from_env_reads_doq_server_and_port() {
        // DNS-DOQ-2 — server + port env wiring lands on extras.
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let saved = (
            std::env::var("CELLOS_DNS_UPSTREAM_DOQ_SERVER").ok(),
            std::env::var("CELLOS_DNS_UPSTREAM_DOQ_PORT").ok(),
        );
        unsafe {
            std::env::set_var("CELLOS_DNS_UPSTREAM_DOQ_SERVER", "9.9.9.9");
            std::env::set_var("CELLOS_DNS_UPSTREAM_DOQ_PORT", "8853");
        }
        let extras = UpstreamExtras::from_env();
        assert_eq!(extras.doq_server.as_deref(), Some("9.9.9.9"));
        assert_eq!(extras.doq_port, Some(8853));
        unsafe {
            match saved.0 {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_DOQ_SERVER", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_DOQ_SERVER"),
            }
            match saved.1 {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_DOQ_PORT", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_DOQ_PORT"),
            }
        }
    }

    #[test]
    fn parse_dot_target_accepts_bare_ipv4() {
        let sa = parse_dot_target("1.1.1.1", None).expect("bare ipv4 parses");
        assert_eq!(sa, "1.1.1.1:853".parse::<SocketAddr>().unwrap());
    }

    #[test]
    fn parse_dot_target_accepts_bare_ipv4_with_port_override() {
        let sa = parse_dot_target("9.9.9.9", Some(8853)).expect("ipv4 + override parses");
        assert_eq!(sa, "9.9.9.9:8853".parse::<SocketAddr>().unwrap());
    }

    #[test]
    fn parse_dot_target_accepts_ipv4_with_inline_port() {
        // Inline port wins over override (operator wrote it explicitly).
        let sa = parse_dot_target("1.1.1.1:9999", Some(853)).expect("inline port parses");
        assert_eq!(sa, "1.1.1.1:9999".parse::<SocketAddr>().unwrap());
    }

    #[test]
    fn parse_dot_target_accepts_bracketed_ipv6() {
        let sa =
            parse_dot_target("[2606:4700:4700::1111]:853", None).expect("bracketed ipv6 parses");
        assert_eq!(
            sa,
            "[2606:4700:4700::1111]:853".parse::<SocketAddr>().unwrap()
        );
    }

    #[test]
    fn parse_dot_target_rejects_hostname() {
        let err = parse_dot_target("dns.example.com", None)
            .expect_err("hostname must be rejected (no DNS bootstrap in netns)");
        assert!(err.contains("hostnames must be pre-resolved"));
    }

    #[test]
    fn parse_dot_target_rejects_empty() {
        assert!(parse_dot_target("", None).is_err());
        assert!(parse_dot_target("   ", None).is_err());
    }

    #[test]
    fn extras_from_env_reads_dot_server_port_sni() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        // Save and restore env state so concurrent tests can't see our
        // overrides. (cargo test runs tests in parallel by default.)
        let saved = (
            std::env::var("CELLOS_DNS_UPSTREAM_DOT_SERVER").ok(),
            std::env::var("CELLOS_DNS_UPSTREAM_DOT_PORT").ok(),
            std::env::var("CELLOS_DNS_UPSTREAM_DOT_SNI").ok(),
        );

        // Use a process-local mutex via a static OnceLock would be cleaner,
        // but the supervisor's existing tests rely on the same env-var
        // pattern (see upstream `CELLOS_DNS_UPSTREAM_TRANSPORT` tests in
        // supervisor.rs); we follow the convention and accept the
        // serial-test caveat for this unit test.
        // SAFETY: these env vars are only read by this test and by
        // `UpstreamExtras::from_env`; no other thread mutates them
        // concurrently within the cellos-supervisor test binary.
        unsafe {
            std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_SERVER", "8.8.8.8");
            std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_PORT", "8853");
            std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_SNI", "dns.google");
        }

        let extras = UpstreamExtras::from_env();
        assert_eq!(extras.dot_server.as_deref(), Some("8.8.8.8"));
        assert_eq!(extras.dot_port, Some(8853));
        assert_eq!(extras.dot_sni.as_deref(), Some("dns.google"));

        // Restore prior env state.
        unsafe {
            match saved.0 {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_SERVER", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_DOT_SERVER"),
            }
            match saved.1 {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_PORT", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_DOT_PORT"),
            }
            match saved.2 {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_SNI", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_DOT_SNI"),
            }
        }
    }

    #[test]
    fn extras_from_env_ignores_unparseable_port() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let saved = std::env::var("CELLOS_DNS_UPSTREAM_DOT_PORT").ok();
        // Clear sibling vars so a parallel test's leftover state can't leak
        // into our observation. The ENV_LOCK serialises us, but tests prior
        // to the lock's introduction may have left state from a panic.
        unsafe {
            std::env::remove_var("CELLOS_DNS_UPSTREAM_DOT_SERVER");
            std::env::remove_var("CELLOS_DNS_UPSTREAM_DOT_SNI");
        }
        unsafe {
            std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_PORT", "not-a-number");
        }
        let extras = UpstreamExtras::from_env();
        assert_eq!(extras.dot_port, None);
        unsafe {
            match saved {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_DOT_PORT", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_DOT_PORT"),
            }
        }
    }

    #[test]
    fn from_env_prefers_protocol_over_transport() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let saved = (
            std::env::var("CELLOS_DNS_UPSTREAM_PROTOCOL").ok(),
            std::env::var("CELLOS_DNS_UPSTREAM_TRANSPORT").ok(),
        );
        unsafe {
            std::env::set_var("CELLOS_DNS_UPSTREAM_PROTOCOL", "dot");
            std::env::set_var("CELLOS_DNS_UPSTREAM_TRANSPORT", "do53-udp");
        }
        // PROTOCOL takes priority — operator's contract-aligned name wins.
        assert_eq!(UpstreamTransport::from_env(), Some(UpstreamTransport::Dot));
        unsafe {
            std::env::remove_var("CELLOS_DNS_UPSTREAM_PROTOCOL");
        }
        // PROTOCOL unset, TRANSPORT still honoured (back-compat).
        assert_eq!(
            UpstreamTransport::from_env(),
            Some(UpstreamTransport::Do53Udp)
        );
        unsafe {
            match saved.0 {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_PROTOCOL", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_PROTOCOL"),
            }
            match saved.1 {
                Some(v) => std::env::set_var("CELLOS_DNS_UPSTREAM_TRANSPORT", v),
                None => std::env::remove_var("CELLOS_DNS_UPSTREAM_TRANSPORT"),
            }
        }
    }

    #[test]
    fn udp_path_round_trips_against_synthetic_upstream() {
        // Spawn a synthetic UDP echo-style upstream that replies with a
        // 13-byte canned packet for any incoming query. Confirms the
        // Do53Udp dispatch goes through `forward_udp` and the result is
        // surfaced byte-for-byte.
        let echo = UdpSocket::bind("127.0.0.1:0").unwrap();
        echo.set_read_timeout(Some(Duration::from_millis(500)))
            .unwrap();
        let echo_addr = echo.local_addr().unwrap();
        std::thread::spawn(move || {
            let mut b = [0u8; 1500];
            if let Ok((_n, peer)) = echo.recv_from(&mut b) {
                let _ = echo.send_to(b"\x00\x00ABCDEFGHIJK", peer);
            }
        });
        let upstream_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
        let mut out = [0u8; 1500];
        let n = forward(
            UpstreamTransport::Do53Udp,
            &upstream_sock,
            echo_addr,
            b"\x00\x00\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00",
            &mut out,
            Duration::from_millis(500),
            &UpstreamExtras::default(),
        )
        .expect("udp round-trip");
        assert_eq!(n, 13);
        assert_eq!(&out[..13], b"\x00\x00ABCDEFGHIJK");
    }
}