cellos-supervisor 0.5.1

//! SEAM-1 / L2-04 DNS proxy — forward-only UDP proxy that enforces
//! `dnsAuthority.hostnameAllowlist` at the DNS protocol layer and emits one
//! per-query CloudEvent for every observed query.
//!
//! ## Phase 1 scope
//!
//! - **UDP only.** The single-listener proxy reads a workload query, parses
//!   the question section (see [`parser`]), evaluates the allowlist, then
//!   either forwards the bytes verbatim to the upstream `do53-udp` resolver
//!   or builds a REFUSED response on the spot. TCP fallback is documented as
//!   a follow-up.
//! - **Forward-only.** No recursion, no caching. The proxy talks to a single
//!   declared upstream resolver per cell and relays the response (rcode,
//!   answers, additional records) verbatim to the workload, preserving the
//!   transaction id binding.
//! - **Allowlist-strict.** A query whose name does not match an entry in
//!   `hostname_allowlist` (literal or single-leading-`*.` wildcard) is
//!   short-circuited with REFUSED — the upstream resolver never sees it.
//! - **Malformed-drop.** Packets that fail [`parser::parse_query`] are
//!   dropped without a response (matching common resolver behaviour against
//!   adversarial input). The corresponding `dns_query` event still fires
//!   with `reasonCode: malformed_query`.
//! - **SERVFAIL on upstream timeout.** When the upstream resolver fails to
//!   answer within [`DnsProxyConfig::upstream_timeout`] the proxy synthesizes
//!   a SERVFAIL response (RCODE=2) so the workload sees a deterministic
//!   error rather than hanging.
//!
//! ## What this module does NOT do
//!
//! - **TCP fallback** — workloads that get a `TC` (truncated) flag from
//!   upstream are expected to retry on TCP; Phase 1 forwards UDP-only and
//!   relays the truncated answer if the upstream sets `TC`. A real TCP
//!   listener path is a Phase 2 follow-up.
//! - **EDNS / DoT / DoH / DoQ upstream** — Phase 1's upstream is `do53-udp`
//!   only. The contract layer (`DnsAuthority::resolvers[]`) admits these
//!   protocols but the dataplane forwards over plain UDP.
//! - **No netns plumbing** — the caller pre-binds the listening + upstream
//!   sockets in the cell's network namespace and hands them to
//!   [`run_one_shot`]. This module is platform-neutral and contains no
//!   `nsenter` or `setns` calls.
//!
//! ## Caller responsibility
//!
//! - Pre-bind the listener socket on the address the workload's
//!   `/etc/resolv.conf` points at (typically the same IP/port declared as a
//!   `do53-udp` resolver in `dnsAuthority.resolvers[]` so the existing nft
//!   accept rule lets the traffic through).
//! - Pre-bind the upstream socket in the same netns; the proxy uses it for
//!   forwarding.
//! - Drive the loop until `shutdown.store(true)` to terminate it; the
//!   listener's `read_timeout` should be set to a short interval (e.g.
//!   200ms) so the loop checks the shutdown flag promptly.

pub mod dnssec;
pub mod parser;
pub mod spawn;
pub mod upstream;

use std::io;
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream, UdpSocket};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use cellos_core::{
    cloud_event_v1_dns_authority_dnssec_failed, cloud_event_v1_dns_query_permitted,
    cloud_event_v1_dns_query_refused, qtype_to_dns_query_type, CloudEventV1,
    DnsAuthorityDnssecFailed, DnsAuthorityDnssecFailureReason, DnsQueryDecision, DnsQueryEvent,
    DnsQueryReasonCode, DnsQueryType,
};

use dnssec::{DataplaneDnssecOutcome, DataplaneDnssecValidator};
use parser::{parse_query, DnsParseError, DnsQueryView, DNS_HEADER_LEN};
use upstream::{UpstreamExtras, UpstreamTransport};

/// Default query types when the operator did not narrow `allowed_query_types`.
const DEFAULT_QUERY_TYPES: &[DnsQueryType] = &[
    DnsQueryType::A,
    DnsQueryType::AAAA,
    DnsQueryType::CNAME,
    DnsQueryType::HTTPS,
];

/// Maximum DNS UDP payload size the proxy buffers. EDNS answers can be larger
/// but Phase 1 forwards over plain UDP; if upstream returns more than this
/// the answer is truncated at the wire and the workload should retry on TCP.
const MAX_UDP_PAYLOAD: usize = 1500;

/// SEAM-1 Phase 2 — default per-connection idle timeout for the workload-side
/// TCP listener path when [`DnsProxyConfig::tcp_idle_timeout`] is at its
/// type-default zero value. Applied as the per-read/write deadline on
/// accepted TCP streams so a stuck workload cannot pin a worker thread
/// forever; a timed-out read drops the connection. Slot **A5** promoted this
/// from a hardcoded constant to a [`DnsProxyConfig`] knob; the constant
/// survives as the fallback for callers that build the config without
/// setting the new field explicitly (e.g. pre-A5 integration tests).
const DEFAULT_TCP_IDLE_TIMEOUT: Duration = Duration::from_secs(30);

/// Configuration for [`run_one_shot`].
///
/// All fields are owned (no borrows / lifetimes) so the supervisor can build
/// this once and pass it into a `spawn_blocking` thread without lifetime
/// gymnastics.
#[derive(Debug, Clone)]
pub struct DnsProxyConfig {
    /// Address the listener socket is bound to. Used only for diagnostic logs;
    /// the caller passes the actual pre-bound socket in.
    pub bind_addr: SocketAddr,
    /// Upstream resolver address (must be a `do53-udp` resolver from
    /// `dnsAuthority.resolvers[]`).
    pub upstream_addr: SocketAddr,
    /// Lowercased hostname allowlist. Entries may have a single leading
    /// `*.` for subdomain wildcards (e.g. `*.cdn.example.com` matches
    /// `foo.cdn.example.com` but not `cdn.example.com` itself).
    pub hostname_allowlist: Vec<String>,
    /// Permitted query types. Empty = default `[A, AAAA, CNAME, HTTPS]`.
    pub allowed_query_types: Vec<DnsQueryType>,
    /// Cell identifier (mirrors `lifecycle.started.cellId`).
    pub cell_id: String,
    /// Run identifier (mirrors `lifecycle.started.runId`).
    pub run_id: String,
    /// Optional `policyDigest` to bind into emitted events.
    pub policy_digest: Option<String>,
    /// Optional `keysetId` to bind into emitted events.
    pub keyset_id: Option<String>,
    /// Optional `issuerKid` to bind into emitted events.
    pub issuer_kid: Option<String>,
    /// Optional `correlationId` to bind into emitted events.
    pub correlation_id: Option<String>,
    /// Resolver identifier — mirrors a `dnsAuthority.resolvers[].resolverId`.
    /// Stamped into events on the allow path so audit can attribute the
    /// upstream answer to a declared resolver.
    pub upstream_resolver_id: String,
    /// Round-trip timeout for the upstream forward. On timeout the proxy
    /// returns SERVFAIL to the workload and emits `reasonCode: upstream_failure`.
    pub upstream_timeout: Duration,
    /// SEAM-1 / L2-04 Slot **A5** — per-connection idle timeout for the
    /// workload-facing TCP/53 listener. Applied as `set_read_timeout` /
    /// `set_write_timeout` on every accepted TCP stream so a stuck workload
    /// cannot pin a worker thread forever. A zero `Duration` is treated as
    /// "unset" and falls back to [`DEFAULT_TCP_IDLE_TIMEOUT`] (30s, matching
    /// the UDP SERVFAIL ceiling). Only consulted by [`run_tcp_one_shot`];
    /// ignored by the UDP-only [`run_one_shot`] entry point.
    pub tcp_idle_timeout: Duration,
    /// SEC-21 Phase 3h.1 — optional dataplane DNSSEC validator. `None`
    /// means the cell's `DnsAuthority` did not request DNSSEC validation
    /// on the dataplane (mode = off); the proxy hot path is unchanged.
    /// `Some(_)` means every allowlisted query is post-validated through
    /// the validator before the upstream answer is relayed to the
    /// workload — see [`dnssec::DataplaneDnssecValidator`] for the full
    /// behaviour matrix.
    pub dnssec_validator: Option<Arc<DataplaneDnssecValidator>>,
    /// SEC-21 Phase 3h.2 / T2.B Slot **A6** — upstream transport selector.
    /// `Default = Do53Udp` so the existing UDP-only behaviour is the no-op
    /// upgrade path. See [`upstream::UpstreamTransport`] for the full
    /// dispatch matrix and [`upstream::forward`] for the hot-path entry.
    pub transport: UpstreamTransport,
    /// A6 — per-transport extras (DoT SNI, etc.). Unused when
    /// `transport == Do53Udp`; populated by callers selecting DoT/DoH/DoQ.
    pub upstream_extras: UpstreamExtras,
}

/// Sink the proxy uses to publish per-query events. Trait-erased so unit tests
/// can plug in an in-memory collector and the supervisor can adapt the
/// production `EventSink` to a synchronous-fire-and-forget shape.
pub trait DnsQueryEmitter: Send + Sync {
    /// Publish a single CloudEvent. Implementations should not block — the
    /// proxy's recv loop runs synchronously and any blocking emit will
    /// directly increase per-query latency for the workload.
    fn emit(&self, event: CloudEventV1);
}

/// Aggregate counters returned by [`run_one_shot`] when the loop terminates.
/// Mirrors the cumulative state observed across a single cell run; the
/// supervisor logs these at teardown for sanity-checking ratios.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct DnsProxyStats {
    pub queries_total: u64,
    pub queries_allowed: u64,
    pub queries_denied: u64,
    pub queries_malformed: u64,
    pub upstream_failures: u64,
}

/// Run the proxy loop synchronously on the calling thread.
///
/// The loop terminates when `shutdown.load(Ordering::SeqCst) == true` between
/// iterations. The `socket` should have a non-zero `read_timeout` so the loop
/// can check the flag promptly between recvs.
///
/// Returns the cumulative [`DnsProxyStats`] when the loop exits cleanly. Errors
/// from the listening socket itself (other than timeouts) bubble up and abort
/// the loop — the supervisor treats those as fatal and tears the proxy down.
pub fn run_one_shot(
    cfg: &DnsProxyConfig,
    socket: &UdpSocket,
    upstream: &UdpSocket,
    emitter: &dyn DnsQueryEmitter,
    shutdown: &AtomicBool,
) -> io::Result<DnsProxyStats> {
    let mut stats = DnsProxyStats::default();
    let mut recv_buf = [0u8; MAX_UDP_PAYLOAD];
    let mut up_buf = [0u8; MAX_UDP_PAYLOAD];

    while !shutdown.load(Ordering::SeqCst) {
        let (n, peer) = match socket.recv_from(&mut recv_buf) {
            Ok(t) => t,
            Err(e) if is_timeout(&e) => continue,
            Err(e) if matches!(e.kind(), io::ErrorKind::Interrupted) => continue,
            Err(e) => return Err(e),
        };
        stats.queries_total = stats.queries_total.saturating_add(1);

        let pkt = &recv_buf[..n];
        match parse_query(pkt) {
            Err(parse_err) => {
                stats.queries_malformed = stats.queries_malformed.saturating_add(1);
                let event = build_event(
                    cfg,
                    EventInputs {
                        view: None,
                        decision: DnsQueryDecision::Deny,
                        reason_code: malformed_reason(parse_err),
                        response_rcode: None,
                        upstream_resolver_id: None,
                        upstream_latency_ms: None,
                        response_target_count: None,
                    },
                );
                emit_event(emitter, event);
                // Malformed packets are dropped — no response to the workload.
                continue;
            }
            Ok(view) => {
                let qtype_known = qtype_to_dns_query_type(view.qtype);
                let allowed_types = if cfg.allowed_query_types.is_empty() {
                    DEFAULT_QUERY_TYPES
                } else {
                    cfg.allowed_query_types.as_slice()
                };
                let qtype_in_set = qtype_known.is_some_and(|t| allowed_types.contains(&t));
                if !qtype_in_set {
                    stats.queries_denied = stats.queries_denied.saturating_add(1);
                    let resp = build_refused_response(pkt, &view);
                    let _ = socket.send_to(&resp, peer);
                    let event = build_event(
                        cfg,
                        EventInputs {
                            view: Some(&view),
                            decision: DnsQueryDecision::Deny,
                            reason_code: DnsQueryReasonCode::DeniedQueryType,
                            response_rcode: Some(5),
                            upstream_resolver_id: None,
                            upstream_latency_ms: None,
                            response_target_count: Some(0),
                        },
                    );
                    emit_event(emitter, event);
                    // SEAM-1 Phase 3 — short-form per-query refusal event.
                    emit_query_refused(cfg, emitter, &view, "denied_query_type");
                    continue;
                }

                if !hostname_in_allowlist(&view.qname, &cfg.hostname_allowlist) {
                    stats.queries_denied = stats.queries_denied.saturating_add(1);
                    let resp = build_refused_response(pkt, &view);
                    let _ = socket.send_to(&resp, peer);
                    let event = build_event(
                        cfg,
                        EventInputs {
                            view: Some(&view),
                            decision: DnsQueryDecision::Deny,
                            reason_code: DnsQueryReasonCode::DeniedNotInAllowlist,
                            response_rcode: Some(5),
                            upstream_resolver_id: None,
                            upstream_latency_ms: None,
                            response_target_count: Some(0),
                        },
                    );
                    emit_event(emitter, event);
                    // SEAM-1 Phase 3 — short-form per-query refusal event.
                    emit_query_refused(cfg, emitter, &view, "denied_not_in_allowlist");
                    continue;
                }

                // SEAM-1 Phase 3 — allowlist + qtype passed; emit the
                // short-form `query_permitted` event BEFORE forwarding so
                // operators can audit query intent independent of upstream
                // outcome.
                emit_query_permitted(cfg, emitter, &view);

                // Allow path — forward upstream.
                let started = Instant::now();
                let upstream_result = upstream::forward(
                    cfg.transport,
                    upstream,
                    cfg.upstream_addr,
                    pkt,
                    &mut up_buf,
                    cfg.upstream_timeout,
                    &cfg.upstream_extras,
                );
                let elapsed_ms = started.elapsed().as_millis() as u64;
                match upstream_result {
                    Ok(resp_len) => {
                        // SEC-21 Phase 3h.1 — interpose the dataplane
                        // DNSSEC validator BEFORE relaying the upstream
                        // answer to the workload. When the validator is
                        // None (mode = off), this whole block is
                        // skipped and the proxy behaviour is
                        // byte-identical to pre-P3h.1.
                        if let Some(validator) = cfg.dnssec_validator.as_ref() {
                            let outcome = validator.validate(pkt, &up_buf[..resp_len]);
                            let action =
                                decide_dnssec_action(validator.is_require_mode(), &outcome);
                            match action {
                                DnssecAction::Forward => {
                                    // Validator approved (or, in
                                    // best_effort, tolerated unsigned).
                                    // Fall through to relay the
                                    // upstream answer as today.
                                }
                                DnssecAction::Servfail { reason } => {
                                    let resp = build_servfail_response(pkt, &view);
                                    let _ = socket.send_to(&resp, peer);
                                    stats.queries_denied = stats.queries_denied.saturating_add(1);
                                    // Per-query observability event —
                                    // signals dnssec_failed via the
                                    // `denied_dnssec` reasonCode (added
                                    // alongside the existing reasons in
                                    // P3h.1).
                                    let q_event = build_event(
                                        cfg,
                                        EventInputs {
                                            view: Some(&view),
                                            decision: DnsQueryDecision::Deny,
                                            reason_code: DnsQueryReasonCode::DeniedDnssec,
                                            response_rcode: Some(2),
                                            upstream_resolver_id: Some(
                                                cfg.upstream_resolver_id.clone(),
                                            ),
                                            upstream_latency_ms: Some(elapsed_ms),
                                            response_target_count: Some(0),
                                        },
                                    );
                                    emit_event(emitter, q_event);
                                    // Domain-specific event — the
                                    // dataplane-stamped
                                    // `dns_authority_dnssec_failed`
                                    // CloudEvent.
                                    let dnssec_event = build_dataplane_dnssec_failed_event(
                                        cfg, &view, validator, reason,
                                    );
                                    emit_event(emitter, dnssec_event);
                                    continue;
                                }
                                DnssecAction::ForwardUnsignedBestEffort => {
                                    // best_effort + Unsigned — explicit
                                    // "tolerate unsigned zones" branch.
                                    // Fall through to relay; NO event
                                    // (the spec is explicit on this).
                                }
                            }
                        }
                        let resp = &up_buf[..resp_len];
                        let _ = socket.send_to(resp, peer);
                        stats.queries_allowed = stats.queries_allowed.saturating_add(1);
                        let answer_count = parse_response_target_count(resp, view.qtype);
                        let event = build_event(
                            cfg,
                            EventInputs {
                                view: Some(&view),
                                decision: DnsQueryDecision::Allow,
                                reason_code: DnsQueryReasonCode::AllowedByAllowlist,
                                response_rcode: Some(parse_response_rcode(resp)),
                                upstream_resolver_id: Some(cfg.upstream_resolver_id.clone()),
                                upstream_latency_ms: Some(elapsed_ms),
                                response_target_count: Some(answer_count),
                            },
                        );
                        emit_event(emitter, event);
                    }
                    Err(_e) => {
                        stats.upstream_failures = stats.upstream_failures.saturating_add(1);
                        let resp = build_servfail_response(pkt, &view);
                        let _ = socket.send_to(&resp, peer);
                        let event = build_event(
                            cfg,
                            EventInputs {
                                view: Some(&view),
                                decision: DnsQueryDecision::Deny,
                                reason_code: DnsQueryReasonCode::UpstreamFailure,
                                response_rcode: Some(2),
                                upstream_resolver_id: Some(cfg.upstream_resolver_id.clone()),
                                upstream_latency_ms: Some(elapsed_ms),
                                response_target_count: Some(0),
                            },
                        );
                        emit_event(emitter, event);
                    }
                }
            }
        }
    }

    Ok(stats)
}

/// SEAM-1 / L2-04 Phase 2 — TCP/53 listener variant of [`run_one_shot`].
///
/// Drives a workload-facing TCP listener while reusing the UDP upstream
/// socket for the actual forward (Phase 1/2 contract: `dnsAuthority.resolvers[]`
/// admits do53-udp at the upstream side; the TCP path at the workload edge
/// exists so responses larger than the classic 512-byte UDP limit can reach
/// the workload without the client having to retry on a separate connection).
///
/// Wire format on the workload side is RFC 1035 §4.2.2: each message is
/// preceded by a 2-byte big-endian length prefix. Connections are persistent
/// — multiple length-prefixed queries may arrive on a single accepted
/// connection until the peer closes or the proxy is shut down. The same
/// allowlist / query-type gates as [`run_one_shot`] apply per query; deny
/// paths build a framed REFUSED reply (rcode=5) and continue serving the
/// connection. Malformed messages (parser rejection) emit one
/// `dns_query{reasonCode: malformed_query}` event and the connection is
/// closed without a response — matching the UDP-side "drop on malformed"
/// behaviour.
///
/// The accept loop polls non-blockingly so the `shutdown` flag is observed
/// promptly; per-connection workers handle their own queries on dedicated
/// threads and contribute to the aggregated [`DnsProxyStats`] via an
/// internal [`Mutex`]. The function returns once `shutdown` is set AND all
/// outstanding worker threads have completed.
///
/// `upstream` is shared (`Arc<UdpSocket>`) because each per-connection worker
/// performs its own forward against the single upstream socket — DNS UDP
/// transactions are stateless per query and the upstream send/recv timeout
/// is enforced inside [`forward_upstream`] for each call.
pub fn run_tcp_one_shot(
    cfg: &DnsProxyConfig,
    listener: &TcpListener,
    upstream: Arc<UdpSocket>,
    emitter: Arc<dyn DnsQueryEmitter>,
    shutdown: &AtomicBool,
) -> io::Result<DnsProxyStats> {
    // Non-blocking accept so we can observe `shutdown` between polls.
    listener.set_nonblocking(true)?;

    let stats = Arc::new(Mutex::new(DnsProxyStats::default()));
    let mut workers: Vec<std::thread::JoinHandle<()>> = Vec::new();

    // A5 — resolve the per-connection idle timeout once per loop entry. A
    // zero Duration on the config is treated as "unset" and falls back to
    // [`DEFAULT_TCP_IDLE_TIMEOUT`] so pre-A5 callers that built the config
    // without setting this field keep their previous behaviour.
    let tcp_idle_timeout = if cfg.tcp_idle_timeout.is_zero() {
        DEFAULT_TCP_IDLE_TIMEOUT
    } else {
        cfg.tcp_idle_timeout
    };

    while !shutdown.load(Ordering::SeqCst) {
        match listener.accept() {
            Ok((stream, _peer)) => {
                // Per-connection blocking I/O — flip the inherited
                // non-blocking flag back off so the worker's read_exact
                // does not spin.
                if let Err(_e) = stream.set_nonblocking(false) {
                    continue;
                }
                // Per-read deadline: a stuck client cannot pin a worker
                // thread forever. The read loop treats a timeout as
                // "drop the connection". Sourced from
                // [`DnsProxyConfig::tcp_idle_timeout`]; see A5 notes on
                // that field for the zero-Duration fallback rule.
                let _ = stream.set_read_timeout(Some(tcp_idle_timeout));
                let _ = stream.set_write_timeout(Some(tcp_idle_timeout));

                let cfg_owned = cfg.clone();
                let upstream = upstream.clone();
                let emitter = emitter.clone();
                let stats = stats.clone();
                let handle = std::thread::spawn(move || {
                    handle_tcp_connection(&cfg_owned, stream, &upstream, &*emitter, &stats);
                });
                workers.push(handle);
            }
            Err(e) if matches!(e.kind(), io::ErrorKind::WouldBlock) => {
                std::thread::sleep(Duration::from_millis(50));
                continue;
            }
            Err(e) if matches!(e.kind(), io::ErrorKind::Interrupted) => continue,
            Err(e) => {
                let _ = listener.set_nonblocking(false);
                return Err(e);
            }
        }
    }

    // Drain workers so the caller sees a clean termination point.
    for h in workers {
        let _ = h.join();
    }

    let _ = listener.set_nonblocking(false);
    let final_stats = *stats.lock().expect("dns proxy stats mutex poisoned");
    Ok(final_stats)
}

/// Handle a single accepted TCP connection: read length-prefixed queries
/// in a loop, gate them against the allowlist, forward to UDP upstream,
/// reply length-prefixed. Returns when the peer closes, on malformed
/// input, or on any I/O error.
fn handle_tcp_connection(
    cfg: &DnsProxyConfig,
    mut stream: TcpStream,
    upstream: &UdpSocket,
    emitter: &dyn DnsQueryEmitter,
    stats: &Mutex<DnsProxyStats>,
) {
    let mut up_buf = [0u8; MAX_UDP_PAYLOAD];
    loop {
        // 2-byte big-endian length prefix.
        let mut len_buf = [0u8; 2];
        if stream.read_exact(&mut len_buf).is_err() {
            // EOF / timeout / reset — drop connection cleanly.
            return;
        }
        let msg_len = u16::from_be_bytes(len_buf) as usize;
        if msg_len == 0 {
            // Zero-length message is malformed; emit and close.
            bump_malformed(stats);
            let event = build_event(
                cfg,
                EventInputs {
                    view: None,
                    decision: DnsQueryDecision::Deny,
                    reason_code: DnsQueryReasonCode::MalformedQuery,
                    response_rcode: None,
                    upstream_resolver_id: None,
                    upstream_latency_ms: None,
                    response_target_count: None,
                },
            );
            emit_event(emitter, event);
            return;
        }
        let mut pkt = vec![0u8; msg_len];
        if stream.read_exact(&mut pkt).is_err() {
            return;
        }

        bump_total(stats);

        match parse_query(&pkt) {
            Err(parse_err) => {
                bump_malformed(stats);
                let event = build_event(
                    cfg,
                    EventInputs {
                        view: None,
                        decision: DnsQueryDecision::Deny,
                        reason_code: malformed_reason(parse_err),
                        response_rcode: None,
                        upstream_resolver_id: None,
                        upstream_latency_ms: None,
                        response_target_count: None,
                    },
                );
                emit_event(emitter, event);
                // Drop the connection — the framing context is suspect.
                return;
            }
            Ok(view) => {
                let qtype_known = qtype_to_dns_query_type(view.qtype);
                let allowed_types = if cfg.allowed_query_types.is_empty() {
                    DEFAULT_QUERY_TYPES
                } else {
                    cfg.allowed_query_types.as_slice()
                };
                let qtype_in_set = qtype_known.is_some_and(|t| allowed_types.contains(&t));
                if !qtype_in_set {
                    bump_denied(stats);
                    let resp = build_refused_response(&pkt, &view);
                    if write_framed(&mut stream, &resp).is_err() {
                        return;
                    }
                    let event = build_event(
                        cfg,
                        EventInputs {
                            view: Some(&view),
                            decision: DnsQueryDecision::Deny,
                            reason_code: DnsQueryReasonCode::DeniedQueryType,
                            response_rcode: Some(5),
                            upstream_resolver_id: None,
                            upstream_latency_ms: None,
                            response_target_count: Some(0),
                        },
                    );
                    emit_event(emitter, event);
                    // SEAM-1 Phase 3 — short-form per-query refusal event.
                    emit_query_refused(cfg, emitter, &view, "denied_query_type");
                    continue;
                }

                if !hostname_in_allowlist(&view.qname, &cfg.hostname_allowlist) {
                    bump_denied(stats);
                    let resp = build_refused_response(&pkt, &view);
                    if write_framed(&mut stream, &resp).is_err() {
                        return;
                    }
                    let event = build_event(
                        cfg,
                        EventInputs {
                            view: Some(&view),
                            decision: DnsQueryDecision::Deny,
                            reason_code: DnsQueryReasonCode::DeniedNotInAllowlist,
                            response_rcode: Some(5),
                            upstream_resolver_id: None,
                            upstream_latency_ms: None,
                            response_target_count: Some(0),
                        },
                    );
                    emit_event(emitter, event);
                    // SEAM-1 Phase 3 — short-form per-query refusal event.
                    emit_query_refused(cfg, emitter, &view, "denied_not_in_allowlist");
                    continue;
                }

                // SEAM-1 Phase 3 — short-form per-query permit event,
                // mirroring the UDP path.
                emit_query_permitted(cfg, emitter, &view);

                // Allow path — forward to UDP upstream (Phase 1/2
                // contract: do53-udp upstream, TCP edge to workload).
                // A6 replaced the inline `forward_upstream` helper with
                // a transport-dispatching `upstream::forward(...)`; the
                // TCP edge is independent of the upstream transport
                // dispatch — both paths share `cfg.transport` /
                // `cfg.upstream_extras`.
                let started = Instant::now();
                let upstream_result = upstream::forward(
                    cfg.transport,
                    upstream,
                    cfg.upstream_addr,
                    &pkt,
                    &mut up_buf,
                    cfg.upstream_timeout,
                    &cfg.upstream_extras,
                );
                let elapsed_ms = started.elapsed().as_millis() as u64;
                match upstream_result {
                    Ok(resp_len) => {
                        // SEC-21 Phase 3h.1 — DNSSEC validator interpose
                        // mirrors the UDP path; the workload-facing reply
                        // is framed length-prefixed in either branch.
                        if let Some(validator) = cfg.dnssec_validator.as_ref() {
                            let outcome = validator.validate(&pkt, &up_buf[..resp_len]);
                            let action =
                                decide_dnssec_action(validator.is_require_mode(), &outcome);
                            match action {
                                DnssecAction::Forward | DnssecAction::ForwardUnsignedBestEffort => {
                                    // Fall through to relay; the
                                    // best_effort+Unsigned branch
                                    // intentionally emits no event.
                                }
                                DnssecAction::Servfail { reason } => {
                                    let resp = build_servfail_response(&pkt, &view);
                                    if write_framed(&mut stream, &resp).is_err() {
                                        return;
                                    }
                                    bump_denied(stats);
                                    let q_event = build_event(
                                        cfg,
                                        EventInputs {
                                            view: Some(&view),
                                            decision: DnsQueryDecision::Deny,
                                            reason_code: DnsQueryReasonCode::DeniedDnssec,
                                            response_rcode: Some(2),
                                            upstream_resolver_id: Some(
                                                cfg.upstream_resolver_id.clone(),
                                            ),
                                            upstream_latency_ms: Some(elapsed_ms),
                                            response_target_count: Some(0),
                                        },
                                    );
                                    emit_event(emitter, q_event);
                                    let dnssec_event = build_dataplane_dnssec_failed_event(
                                        cfg, &view, validator, reason,
                                    );
                                    emit_event(emitter, dnssec_event);
                                    continue;
                                }
                            }
                        }
                        let resp = &up_buf[..resp_len];
                        if write_framed(&mut stream, resp).is_err() {
                            return;
                        }
                        bump_allowed(stats);
                        let answer_count = parse_response_target_count(resp, view.qtype);
                        let event = build_event(
                            cfg,
                            EventInputs {
                                view: Some(&view),
                                decision: DnsQueryDecision::Allow,
                                reason_code: DnsQueryReasonCode::AllowedByAllowlist,
                                response_rcode: Some(parse_response_rcode(resp)),
                                upstream_resolver_id: Some(cfg.upstream_resolver_id.clone()),
                                upstream_latency_ms: Some(elapsed_ms),
                                response_target_count: Some(answer_count),
                            },
                        );
                        emit_event(emitter, event);
                    }
                    Err(_e) => {
                        bump_upstream_failure(stats);
                        let resp = build_servfail_response(&pkt, &view);
                        if write_framed(&mut stream, &resp).is_err() {
                            return;
                        }
                        let event = build_event(
                            cfg,
                            EventInputs {
                                view: Some(&view),
                                decision: DnsQueryDecision::Deny,
                                reason_code: DnsQueryReasonCode::UpstreamFailure,
                                response_rcode: Some(2),
                                upstream_resolver_id: Some(cfg.upstream_resolver_id.clone()),
                                upstream_latency_ms: Some(elapsed_ms),
                                response_target_count: Some(0),
                            },
                        );
                        emit_event(emitter, event);
                    }
                }
            }
        }
    }
}

/// Send a DNS message to the workload framed with the 2-byte big-endian
/// length prefix per RFC 1035 §4.2.2.
fn write_framed(stream: &mut TcpStream, msg: &[u8]) -> io::Result<()> {
    let len = u16::try_from(msg.len()).map_err(|_| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            "DNS message exceeds 65535-byte TCP frame limit",
        )
    })?;
    stream.write_all(&len.to_be_bytes())?;
    stream.write_all(msg)?;
    stream.flush()?;
    Ok(())
}

fn bump_total(stats: &Mutex<DnsProxyStats>) {
    if let Ok(mut s) = stats.lock() {
        s.queries_total = s.queries_total.saturating_add(1);
    }
}

fn bump_allowed(stats: &Mutex<DnsProxyStats>) {
    if let Ok(mut s) = stats.lock() {
        s.queries_allowed = s.queries_allowed.saturating_add(1);
    }
}

fn bump_denied(stats: &Mutex<DnsProxyStats>) {
    if let Ok(mut s) = stats.lock() {
        s.queries_denied = s.queries_denied.saturating_add(1);
    }
}

fn bump_malformed(stats: &Mutex<DnsProxyStats>) {
    if let Ok(mut s) = stats.lock() {
        s.queries_malformed = s.queries_malformed.saturating_add(1);
    }
}

fn bump_upstream_failure(stats: &Mutex<DnsProxyStats>) {
    if let Ok(mut s) = stats.lock() {
        s.upstream_failures = s.upstream_failures.saturating_add(1);
    }
}

/// SEC-21 Phase 3h.1 — call-site policy for one
/// [`DataplaneDnssecOutcome`] given the validator's mode (`fail_closed`).
///
/// Pure function; no I/O. Trivially unit-testable. Encodes the spec's
/// behaviour matrix:
///
/// - `Validated` (any mode) → Forward.
/// - `Unsigned` + require → Servfail with `unsigned_in_require_mode`.
/// - `Unsigned` + best_effort → ForwardUnsignedBestEffort (no event).
/// - `Failed` (any mode) → Servfail with `validation_failed`. Bogus is
///   ALWAYS rejected, even in best_effort.
/// - `Skip` + require → Servfail with `unsupported_query_type_in_require_mode`.
/// - `Skip` + best_effort → Forward (the call site forwards the raw
///   upstream answer unvalidated; documented honestly as a known gap
///   for non-A/AAAA queries).
fn decide_dnssec_action(require_mode: bool, outcome: &DataplaneDnssecOutcome) -> DnssecAction {
    match (outcome, require_mode) {
        (DataplaneDnssecOutcome::Validated, _) => DnssecAction::Forward,
        (DataplaneDnssecOutcome::Failed { reason }, _) => DnssecAction::Servfail { reason },
        (DataplaneDnssecOutcome::Unsigned, true) => DnssecAction::Servfail {
            reason: "unsigned_in_require_mode",
        },
        (DataplaneDnssecOutcome::Unsigned, false) => DnssecAction::ForwardUnsignedBestEffort,
        (DataplaneDnssecOutcome::Skip, true) => DnssecAction::Servfail {
            reason: "unsupported_query_type_in_require_mode",
        },
        (DataplaneDnssecOutcome::Skip, false) => DnssecAction::Forward,
    }
}

/// Three call-site actions returned by [`decide_dnssec_action`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DnssecAction {
    /// Relay the upstream answer to the workload as today (validated, OR
    /// best_effort + Skip — no event).
    Forward,
    /// Build a SERVFAIL response, send to workload, emit
    /// `dns_authority_dnssec_failed{source:"dataplane"}` with `reason`.
    /// `reason` is the schema enum literal.
    Servfail { reason: &'static str },
    /// best_effort + Unsigned — relay the upstream answer; NO event.
    ForwardUnsignedBestEffort,
}

/// Build the dataplane-stamped `dns_authority_dnssec_failed` CloudEvent.
///
/// Mirrors [`crate::resolver_refresh`]'s emission of the same event one-for-one
/// EXCEPT for the `source` field which is stamped `"dataplane"` here vs
/// `"resolver_refresh"` there. The schema's `reason` enum was extended in
/// P3h.1 with `unsigned_in_require_mode` and
/// `unsupported_query_type_in_require_mode` — both are dataplane-only.
fn build_dataplane_dnssec_failed_event(
    cfg: &DnsProxyConfig,
    view: &DnsQueryView,
    validator: &DataplaneDnssecValidator,
    reason: &'static str,
) -> CloudEventV1 {
    let observed_at = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
    // Two dataplane-specific reason strings (added to the schema enum in
    // P3h.1) pass through as-is; everything else collapses to the P3h
    // `validation_failed` enum literal so SIEM rules that already
    // dispatch on resolver-refresh `validation_failed` continue to fire
    // for the dataplane equivalent. `trust_anchor_missing` cannot
    // surface from `validate()` today — `from_authority` rejects bad
    // anchor paths at activation, so the validator never constructs.
    let reason_str = match reason {
        "unsigned_in_require_mode" => "unsigned_in_require_mode",
        "unsupported_query_type_in_require_mode" => "unsupported_query_type_in_require_mode",
        _ => DnsAuthorityDnssecFailureReason::ValidationFailed.as_str(),
    };
    let payload = DnsAuthorityDnssecFailed {
        schema_version: "1.0.0".into(),
        cell_id: cfg.cell_id.clone(),
        run_id: cfg.run_id.clone(),
        resolver_id: cfg.upstream_resolver_id.clone(),
        hostname: view.qname.clone(),
        reason: reason_str.into(),
        // require mode → fail_closed=true; best_effort → false. The
        // event's failClosed field reflects the policy that produced
        // the SERVFAIL (best_effort + Failed still SERVFAILs, and
        // failClosed=false correctly says "policy was best_effort but
        // a hard validation failure was observed").
        fail_closed: validator.is_require_mode(),
        trust_anchor_source: validator.trust_anchor_source().to_string(),
        policy_digest: cfg.policy_digest.clone(),
        keyset_id: cfg.keyset_id.clone(),
        issuer_kid: cfg.issuer_kid.clone(),
        correlation_id: cfg.correlation_id.clone(),
        // SEC-21 Phase 3h.1 — additive `source` discriminator. The
        // dataplane stamps "dataplane"; the resolver-refresh path
        // stamps "resolver_refresh".
        source: Some("dataplane".into()),
        observed_at: observed_at.clone(),
    };
    cloud_event_v1_dns_authority_dnssec_failed("cellos-dns-proxy", &observed_at, &payload)
        .expect("DnsAuthorityDnssecFailed serializes to JSON")
}

struct EventInputs<'a> {
    view: Option<&'a DnsQueryView>,
    decision: DnsQueryDecision,
    reason_code: DnsQueryReasonCode,
    response_rcode: Option<u8>,
    upstream_resolver_id: Option<String>,
    upstream_latency_ms: Option<u64>,
    response_target_count: Option<u32>,
}

fn build_event(cfg: &DnsProxyConfig, inputs: EventInputs<'_>) -> CloudEventV1 {
    let observed_at = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
    let (qname, qtype) = match inputs.view {
        Some(v) => (
            v.qname.clone(),
            qtype_to_dns_query_type(v.qtype).unwrap_or(DnsQueryType::A),
        ),
        None => (String::new(), DnsQueryType::A),
    };
    let payload = DnsQueryEvent {
        schema_version: "1.0.0".into(),
        cell_id: cfg.cell_id.clone(),
        run_id: cfg.run_id.clone(),
        query_id: uuid::Uuid::new_v4().to_string(),
        // Schema requires queryName.minLength=1; for malformed/empty queries
        // we substitute the conventional placeholder so the event still
        // validates while preserving the malformed signal in `reasonCode`.
        query_name: if qname.is_empty() {
            "(unknown)".into()
        } else {
            qname
        },
        query_type: qtype,
        decision: inputs.decision,
        reason_code: inputs.reason_code,
        response_rcode: inputs.response_rcode,
        upstream_resolver_id: inputs.upstream_resolver_id,
        upstream_latency_ms: inputs.upstream_latency_ms,
        response_target_count: inputs.response_target_count,
        keyset_id: cfg.keyset_id.clone(),
        issuer_kid: cfg.issuer_kid.clone(),
        policy_digest: cfg.policy_digest.clone(),
        correlation_id: cfg.correlation_id.clone(),
        observed_at: observed_at.clone(),
    };
    cellos_core::cloud_event_v1_dns_query("cellos-dns-proxy", &observed_at, &payload)
        .expect("DnsQueryEvent serializes to JSON")
}

fn emit_event(emitter: &dyn DnsQueryEmitter, event: CloudEventV1) {
    emitter.emit(event);
}

/// Map a [`DnsQueryType`] to its on-wire IANA mnemonic string used in
/// per-query CloudEvent payloads.
fn dns_query_type_str(t: DnsQueryType) -> &'static str {
    match t {
        DnsQueryType::A => "A",
        DnsQueryType::AAAA => "AAAA",
        DnsQueryType::CNAME => "CNAME",
        DnsQueryType::TXT => "TXT",
        DnsQueryType::MX => "MX",
        DnsQueryType::SRV => "SRV",
        DnsQueryType::NS => "NS",
        DnsQueryType::PTR => "PTR",
        DnsQueryType::HTTPS => "HTTPS",
        DnsQueryType::SVCB => "SVCB",
    }
}

/// SEAM-1 Phase 3 — best-effort emission of
/// `dev.cellos.events.cell.dns.v1.query_permitted` immediately after the
/// allowlist + qtype gates pass and before the proxy forwards the query
/// upstream. Fires once per query on the allow path, regardless of the
/// upstream outcome that follows.
fn emit_query_permitted(cfg: &DnsProxyConfig, emitter: &dyn DnsQueryEmitter, view: &DnsQueryView) {
    let qtype = qtype_to_dns_query_type(view.qtype).unwrap_or(DnsQueryType::A);
    let observed_at = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
    let event = cloud_event_v1_dns_query_permitted(
        "cellos-dns-proxy",
        &observed_at,
        &view.qname,
        dns_query_type_str(qtype),
        &cfg.cell_id,
        &cfg.upstream_resolver_id,
    );
    emit_event(emitter, event);
}

/// SEAM-1 Phase 3 — best-effort emission of
/// `dev.cellos.events.cell.dns.v1.query_refused` whenever the proxy
/// short-circuits a workload query with REFUSED (allowlist miss or
/// disallowed query-type). `reason` is the upstream-aggregate event's
/// `reasonCode` string so SIEM rules can correlate.
fn emit_query_refused(
    cfg: &DnsProxyConfig,
    emitter: &dyn DnsQueryEmitter,
    view: &DnsQueryView,
    reason: &str,
) {
    let qtype = qtype_to_dns_query_type(view.qtype).unwrap_or(DnsQueryType::A);
    let observed_at = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
    let event = cloud_event_v1_dns_query_refused(
        "cellos-dns-proxy",
        &observed_at,
        &view.qname,
        dns_query_type_str(qtype),
        &cfg.cell_id,
        reason,
    );
    emit_event(emitter, event);
}

fn malformed_reason(_e: DnsParseError) -> DnsQueryReasonCode {
    // All parser errors collapse to `malformed_query` in the emitted event;
    // the schema's reasonCode enum doesn't carry the underlying parser-error
    // detail. The parser's structured error type is retained for log triage.
    DnsQueryReasonCode::MalformedQuery
}

/// Match `qname` against a leading-`*.` wildcard or literal allowlist.
///
/// Thin wrapper over the shared
/// [`cellos_core::hostname_allowlist::matches_allowlist`] helper — preserved as
/// a local symbol so existing tests in this module continue to call it. Both
/// the SEAM-1 DNS proxy and the SEC-22 Phase 2 SNI proxy share the same
/// matcher; see that module for the documented semantics.
fn hostname_in_allowlist(qname: &str, allowlist: &[String]) -> bool {
    cellos_core::hostname_allowlist::matches_allowlist(qname, allowlist)
}

fn is_timeout(e: &io::Error) -> bool {
    matches!(
        e.kind(),
        io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
    )
}

/// Build a REFUSED (RCODE=5) response from the workload's query bytes.
///
/// The response copies the question section verbatim, sets `QR=1` + `RA=0`
/// + `RCODE=5` in the flags word, and zeroes ANCOUNT/NSCOUNT/ARCOUNT.
fn build_refused_response(query: &[u8], view: &DnsQueryView) -> Vec<u8> {
    build_error_response(query, view, 5)
}

fn build_servfail_response(query: &[u8], view: &DnsQueryView) -> Vec<u8> {
    build_error_response(query, view, 2)
}

fn build_error_response(query: &[u8], view: &DnsQueryView, rcode: u8) -> Vec<u8> {
    // Find the end of the question section (header + qname + 4 bytes qtype/qclass).
    // Fall back to the full packet length if our parser-derived offset is past
    // the end (defensive — should never trip given parse_query already validated).
    let mut question_end = DNS_HEADER_LEN;
    let mut idx = DNS_HEADER_LEN;
    while idx < query.len() {
        let b = query[idx];
        if b == 0 {
            idx += 1;
            break;
        }
        // parse_query rejects compression so we won't see top bits set here.
        idx += 1 + b as usize;
    }
    if idx + 4 <= query.len() {
        question_end = idx + 4;
    }
    let mut resp = Vec::with_capacity(question_end);
    resp.extend_from_slice(&view.txn_id.to_be_bytes());
    // Flags: QR=1, OPCODE copied from query, AA=0, TC=0, RD copied, RA=0, Z=0, RCODE.
    let mut flags = view.flags;
    flags |= 0x8000; // QR
    flags &= !0x0080; // RA cleared
    flags = (flags & 0xfff0) | u16::from(rcode & 0x0f);
    resp.extend_from_slice(&flags.to_be_bytes());
    // QDCOUNT=1, ANCOUNT=0, NSCOUNT=0, ARCOUNT=0.
    resp.extend_from_slice(&[0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
    // Question section verbatim.
    resp.extend_from_slice(&query[DNS_HEADER_LEN..question_end]);
    resp
}

/// Read the RCODE nibble from a response packet's flags word.
fn parse_response_rcode(resp: &[u8]) -> u8 {
    if resp.len() < 4 {
        return 2; // SERVFAIL on truncated upstream response
    }
    resp[3] & 0x0f
}

/// Count A/AAAA records in the answer section of an upstream response.
///
/// scope: walks the answer section, skipping each RR's owner name
/// (handling pointer-compression on the response side because upstream
/// resolvers DO use it), and counts records whose TYPE matches the
/// original query type when the query was A or AAAA. Returns 0 for
/// non-A/AAAA queries.
///
/// Best-effort: any parse failure during the walk returns 0 rather than
/// surfacing an error — the answer count is observability only.
fn parse_response_target_count(resp: &[u8], qtype: u16) -> u32 {
    if !matches!(qtype, 1 | 28) {
        return 0;
    }
    if resp.len() < DNS_HEADER_LEN {
        return 0;
    }
    let qdcount = u16::from_be_bytes([resp[4], resp[5]]) as usize;
    let ancount = u16::from_be_bytes([resp[6], resp[7]]) as usize;
    let mut idx = DNS_HEADER_LEN;
    // Skip QDCOUNT questions.
    for _ in 0..qdcount {
        idx = match skip_name(resp, idx) {
            Some(n) => n,
            None => return 0,
        };
        idx += 4; // qtype + qclass
        if idx > resp.len() {
            return 0;
        }
    }
    let mut count: u32 = 0;
    for _ in 0..ancount {
        idx = match skip_name(resp, idx) {
            Some(n) => n,
            None => return count,
        };
        if idx + 10 > resp.len() {
            return count;
        }
        let rtype = u16::from_be_bytes([resp[idx], resp[idx + 1]]);
        let rdlen = u16::from_be_bytes([resp[idx + 8], resp[idx + 9]]) as usize;
        idx += 10;
        if rtype == qtype {
            count = count.saturating_add(1);
        }
        idx += rdlen;
        if idx > resp.len() {
            return count;
        }
    }
    count
}

/// Skip over an encoded DNS name (handling pointer-compression in responses).
/// Returns the offset just past the terminating root label or pointer.
fn skip_name(buf: &[u8], mut idx: usize) -> Option<usize> {
    loop {
        if idx >= buf.len() {
            return None;
        }
        let b = buf[idx];
        if b == 0 {
            return Some(idx + 1);
        }
        if b & 0xc0 == 0xc0 {
            // 2-byte pointer.
            if idx + 1 >= buf.len() {
                return None;
            }
            return Some(idx + 2);
        }
        idx += 1 + b as usize;
    }
}

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

    /// In-memory emitter for tests.
    #[derive(Default)]
    struct MemEmitter {
        events: Mutex<Vec<CloudEventV1>>,
    }
    impl DnsQueryEmitter for MemEmitter {
        fn emit(&self, event: CloudEventV1) {
            self.events.lock().unwrap().push(event);
        }
    }

    fn build_query_packet(qname: &str, qtype: u16) -> Vec<u8> {
        let mut p = Vec::new();
        p.extend_from_slice(&[
            0xab, 0xcd, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ]);
        for label in qname.split('.') {
            p.push(label.len() as u8);
            p.extend_from_slice(label.as_bytes());
        }
        p.push(0);
        p.extend_from_slice(&qtype.to_be_bytes());
        p.extend_from_slice(&[0x00, 0x01]);
        p
    }

    /// Build a synthetic upstream response with `ancount` A records (RDATA = 4 bytes 0).
    fn build_a_response(query: &[u8], ancount: u16) -> Vec<u8> {
        // Reuse query header; set QR + ANCOUNT.
        let mut resp = query.to_vec();
        resp[2] = 0x81;
        resp[3] = 0x80;
        resp[6] = (ancount >> 8) as u8;
        resp[7] = (ancount & 0xff) as u8;
        for _ in 0..ancount {
            resp.extend_from_slice(&[0xc0, 0x0c]); // pointer to QNAME
            resp.extend_from_slice(&[0x00, 0x01]); // type A
            resp.extend_from_slice(&[0x00, 0x01]); // class IN
            resp.extend_from_slice(&[0x00, 0x00, 0x01, 0x2c]); // TTL 300
            resp.extend_from_slice(&[0x00, 0x04]); // RDLENGTH 4
            resp.extend_from_slice(&[203, 0, 113, 1]);
        }
        resp
    }

    /// Spawn a tiny localhost upstream that replies to the next N queries with
    /// `build_a_response`, optionally swallowing them (timeout simulation).
    fn spawn_upstream(swallow: bool, ancount: u16) -> (SocketAddr, std::thread::JoinHandle<()>) {
        let sock = UdpSocket::bind("127.0.0.1:0").unwrap();
        let addr = sock.local_addr().unwrap();
        sock.set_read_timeout(Some(Duration::from_millis(2000)))
            .unwrap();
        let h = std::thread::spawn(move || {
            let mut buf = [0u8; 1500];
            while let Ok((n, peer)) = sock.recv_from(&mut buf) {
                if swallow {
                    // Drop the packet; client should time out.
                    continue;
                }
                let resp = build_a_response(&buf[..n], ancount);
                let _ = sock.send_to(&resp, peer);
            }
        });
        (addr, h)
    }

    fn proxy_cfg(allowlist: Vec<&str>, upstream: SocketAddr) -> DnsProxyConfig {
        DnsProxyConfig {
            bind_addr: "127.0.0.1:0".parse().unwrap(),
            upstream_addr: upstream,
            hostname_allowlist: allowlist.into_iter().map(String::from).collect(),
            allowed_query_types: vec![],
            cell_id: "test-cell".into(),
            run_id: "test-run".into(),
            policy_digest: None,
            keyset_id: Some("test-keyset".into()),
            issuer_kid: Some("test-kid-001".into()),
            correlation_id: None,
            upstream_resolver_id: "resolver-test-001".into(),
            upstream_timeout: Duration::from_millis(300),
            // A5 — in-module tests don't exercise the TCP path; leave
            // tcp_idle_timeout at zero so the run_tcp_one_shot fallback
            // rule kicks in if these helpers ever get reused there.
            tcp_idle_timeout: Duration::ZERO,
            // SEC-21 Phase 3h.1 — these in-module tests pre-date the
            // dataplane DNSSEC validator and exercise the proxy's
            // base-level allowlist + forward behaviour with mode=off
            // (validator absent). The dedicated DNSSEC integration
            // tests live in `tests/dataplane_dnssec_validation.rs` and
            // construct their own DataplaneDnssecValidator via
            // `with_backend`.
            dnssec_validator: None,
            // A6 — UDP default preserves byte-identical pre-A6 behaviour
            // for these in-module tests; the new DoT path is exercised
            // in `tests/supervisor_dns_proxy_dot.rs`.
            transport: UpstreamTransport::Do53Udp,
            upstream_extras: UpstreamExtras::default(),
        }
    }

    #[test]
    fn proxy_allows_query_in_allowlist() {
        let (upstream_addr, _h) = spawn_upstream(false, 2);
        let listener = UdpSocket::bind("127.0.0.1:0").unwrap();
        listener
            .set_read_timeout(Some(Duration::from_millis(150)))
            .unwrap();
        let listen_addr = listener.local_addr().unwrap();
        let upstream_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
        let cfg = proxy_cfg(vec!["api.example.com"], upstream_addr);
        let emitter = std::sync::Arc::new(MemEmitter::default());
        let shutdown = std::sync::Arc::new(AtomicBool::new(false));

        let proxy_handle = {
            let emitter = emitter.clone();
            let shutdown = shutdown.clone();
            let cfg = cfg.clone();
            std::thread::spawn(move || {
                let _ = run_one_shot(&cfg, &listener, &upstream_sock, &*emitter, &shutdown);
            })
        };

        let client = UdpSocket::bind("127.0.0.1:0").unwrap();
        client
            .set_read_timeout(Some(Duration::from_secs(1)))
            .unwrap();
        let q = build_query_packet("api.example.com", 1);
        client.send_to(&q, listen_addr).unwrap();
        let mut rb = [0u8; 1500];
        let (n, _) = client.recv_from(&mut rb).unwrap();
        assert!(n > DNS_HEADER_LEN);
        let rcode = rb[3] & 0x0f;
        assert_eq!(rcode, 0, "expected NOERROR on allow path");

        shutdown.store(true, Ordering::SeqCst);
        proxy_handle.join().unwrap();
        let evs = emitter.events.lock().unwrap();
        // SEAM-1 Phase 3 — the allow path now emits TWO events: the
        // short-form `query_permitted` BEFORE the upstream forward, then
        // the aggregate `dns_query{decision:allow}` AFTER the answer.
        assert_eq!(evs.len(), 2);
        assert_eq!(evs[0].ty, "dev.cellos.events.cell.dns.v1.query_permitted");
        let permitted_data = evs[0].data.as_ref().unwrap();
        assert_eq!(permitted_data["queryName"], "api.example.com");
        assert_eq!(permitted_data["queryType"], "A");
        assert_eq!(permitted_data["resolver"], "resolver-test-001");
        let data = evs[1].data.as_ref().unwrap();
        assert_eq!(data["decision"], "allow");
        assert_eq!(data["reasonCode"], "allowed_by_allowlist");
        assert_eq!(data["responseRcode"], 0);
        assert_eq!(data["upstreamResolverId"], "resolver-test-001");
        assert_eq!(data["responseTargetCount"], 2);
    }

    #[test]
    fn proxy_denies_query_not_in_allowlist() {
        let (upstream_addr, _h) = spawn_upstream(false, 0);
        let listener = UdpSocket::bind("127.0.0.1:0").unwrap();
        listener
            .set_read_timeout(Some(Duration::from_millis(150)))
            .unwrap();
        let listen_addr = listener.local_addr().unwrap();
        let upstream_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
        let cfg = proxy_cfg(vec!["api.example.com"], upstream_addr);
        let emitter = std::sync::Arc::new(MemEmitter::default());
        let shutdown = std::sync::Arc::new(AtomicBool::new(false));

        let proxy_handle = {
            let emitter = emitter.clone();
            let shutdown = shutdown.clone();
            let cfg = cfg.clone();
            std::thread::spawn(move || {
                let _ = run_one_shot(&cfg, &listener, &upstream_sock, &*emitter, &shutdown);
            })
        };

        let client = UdpSocket::bind("127.0.0.1:0").unwrap();
        client
            .set_read_timeout(Some(Duration::from_secs(1)))
            .unwrap();
        let q = build_query_packet("blocked.example.com", 1);
        client.send_to(&q, listen_addr).unwrap();
        let mut rb = [0u8; 1500];
        let (n, _) = client.recv_from(&mut rb).unwrap();
        let rcode = rb[3] & 0x0f;
        assert_eq!(
            rcode, 5,
            "expected REFUSED on deny path, got rcode={rcode} n={n}"
        );

        shutdown.store(true, Ordering::SeqCst);
        proxy_handle.join().unwrap();
        let evs = emitter.events.lock().unwrap();
        // SEAM-1 Phase 3 — deny path now emits BOTH the aggregate
        // `dns_query{decision:deny}` AND the short-form
        // `dns.v1.query_refused`.
        assert_eq!(evs.len(), 2, "aggregate + short-form refusal event");
        let data = evs[0].data.as_ref().unwrap();
        assert_eq!(data["decision"], "deny");
        assert_eq!(data["reasonCode"], "denied_not_in_allowlist");
        assert_eq!(data["responseRcode"], 5);
        assert_eq!(evs[1].ty, "dev.cellos.events.cell.dns.v1.query_refused");
        let refused = evs[1].data.as_ref().unwrap();
        assert_eq!(refused["reason"], "denied_not_in_allowlist");
        assert_eq!(refused["queryName"], "blocked.example.com");
    }

    #[test]
    fn proxy_wildcard_matches_subdomain_only() {
        // *.cdn.example.com should match foo.cdn.example.com but NOT cdn.example.com
        assert!(hostname_in_allowlist(
            "foo.cdn.example.com",
            &["*.cdn.example.com".into()]
        ));
        assert!(hostname_in_allowlist(
            "deep.foo.cdn.example.com",
            &["*.cdn.example.com".into()]
        ));
        assert!(!hostname_in_allowlist(
            "cdn.example.com",
            &["*.cdn.example.com".into()]
        ));
        assert!(!hostname_in_allowlist(
            "evil-cdn.example.com",
            &["*.cdn.example.com".into()]
        ));
        // Literal entry: exact match only.
        assert!(hostname_in_allowlist(
            "api.example.com",
            &["api.example.com".into()]
        ));
        assert!(!hostname_in_allowlist(
            "x.api.example.com",
            &["api.example.com".into()]
        ));
    }

    #[test]
    fn proxy_denies_disallowed_query_type() {
        let (upstream_addr, _h) = spawn_upstream(false, 0);
        let listener = UdpSocket::bind("127.0.0.1:0").unwrap();
        listener
            .set_read_timeout(Some(Duration::from_millis(150)))
            .unwrap();
        let listen_addr = listener.local_addr().unwrap();
        let upstream_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
        let mut cfg = proxy_cfg(vec!["api.example.com"], upstream_addr);
        // Restrict to A/AAAA only — TXT (16) is not permitted.
        cfg.allowed_query_types = vec![DnsQueryType::A, DnsQueryType::AAAA];
        let emitter = std::sync::Arc::new(MemEmitter::default());
        let shutdown = std::sync::Arc::new(AtomicBool::new(false));

        let proxy_handle = {
            let emitter = emitter.clone();
            let shutdown = shutdown.clone();
            let cfg = cfg.clone();
            std::thread::spawn(move || {
                let _ = run_one_shot(&cfg, &listener, &upstream_sock, &*emitter, &shutdown);
            })
        };

        let client = UdpSocket::bind("127.0.0.1:0").unwrap();
        client
            .set_read_timeout(Some(Duration::from_secs(1)))
            .unwrap();
        let q = build_query_packet("api.example.com", 16); // TXT
        client.send_to(&q, listen_addr).unwrap();
        let mut rb = [0u8; 1500];
        let (_n, _) = client.recv_from(&mut rb).unwrap();
        assert_eq!(rb[3] & 0x0f, 5);

        shutdown.store(true, Ordering::SeqCst);
        proxy_handle.join().unwrap();
        let evs = emitter.events.lock().unwrap();
        // SEAM-1 Phase 3 — aggregate + short-form `query_refused`.
        assert_eq!(evs.len(), 2);
        let data = evs[0].data.as_ref().unwrap();
        assert_eq!(data["reasonCode"], "denied_query_type");
        assert_eq!(evs[1].ty, "dev.cellos.events.cell.dns.v1.query_refused");
        assert_eq!(evs[1].data.as_ref().unwrap()["reason"], "denied_query_type");
    }

    #[test]
    fn proxy_emits_event_per_query() {
        let (upstream_addr, _h) = spawn_upstream(false, 1);
        let listener = UdpSocket::bind("127.0.0.1:0").unwrap();
        listener
            .set_read_timeout(Some(Duration::from_millis(150)))
            .unwrap();
        let listen_addr = listener.local_addr().unwrap();
        let upstream_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
        let cfg = proxy_cfg(vec!["api.example.com", "*.cdn.example.com"], upstream_addr);
        let emitter = std::sync::Arc::new(MemEmitter::default());
        let shutdown = std::sync::Arc::new(AtomicBool::new(false));

        let proxy_handle = {
            let emitter = emitter.clone();
            let shutdown = shutdown.clone();
            let cfg = cfg.clone();
            std::thread::spawn(move || {
                let _ = run_one_shot(&cfg, &listener, &upstream_sock, &*emitter, &shutdown);
            })
        };

        let client = UdpSocket::bind("127.0.0.1:0").unwrap();
        client
            .set_read_timeout(Some(Duration::from_secs(1)))
            .unwrap();
        for name in [
            "api.example.com",
            "img.cdn.example.com",
            "blocked.example.com",
        ] {
            let q = build_query_packet(name, 1);
            client.send_to(&q, listen_addr).unwrap();
            let mut rb = [0u8; 1500];
            let _ = client.recv_from(&mut rb).unwrap();
        }

        shutdown.store(true, Ordering::SeqCst);
        proxy_handle.join().unwrap();
        let evs = emitter.events.lock().unwrap();
        // SEAM-1 Phase 3 — 3 queries × (short-form + aggregate) = 6 events.
        // Allow paths emit `query_permitted` BEFORE the aggregate; deny
        // paths emit the aggregate BEFORE `query_refused`.
        assert_eq!(evs.len(), 6);
        // Query 1 — api.example.com allowed.
        assert_eq!(evs[0].ty, "dev.cellos.events.cell.dns.v1.query_permitted");
        assert_eq!(
            evs[0].data.as_ref().unwrap()["queryName"],
            "api.example.com"
        );
        let data1_agg = evs[1].data.as_ref().unwrap();
        assert_eq!(data1_agg["decision"], "allow");
        assert_eq!(data1_agg["queryName"], "api.example.com");
        assert_eq!(data1_agg["upstreamResolverId"], "resolver-test-001");
        // Query 2 — img.cdn.example.com allowed (wildcard match).
        assert_eq!(evs[2].ty, "dev.cellos.events.cell.dns.v1.query_permitted");
        let data3_agg = evs[3].data.as_ref().unwrap();
        assert_eq!(data3_agg["decision"], "allow");
        assert_eq!(data3_agg["queryName"], "img.cdn.example.com");
        // Query 3 — blocked.example.com denied. Aggregate first, refused after.
        let data4_agg = evs[4].data.as_ref().unwrap();
        assert_eq!(data4_agg["decision"], "deny");
        assert_eq!(data4_agg["queryName"], "blocked.example.com");
        assert_eq!(evs[5].ty, "dev.cellos.events.cell.dns.v1.query_refused");
        assert_eq!(
            evs[5].data.as_ref().unwrap()["reason"],
            "denied_not_in_allowlist"
        );
    }

    #[test]
    fn proxy_returns_servfail_on_upstream_timeout() {
        let (upstream_addr, _h) = spawn_upstream(true, 0); // swallow → forced timeout
        let listener = UdpSocket::bind("127.0.0.1:0").unwrap();
        listener
            .set_read_timeout(Some(Duration::from_millis(150)))
            .unwrap();
        let listen_addr = listener.local_addr().unwrap();
        let upstream_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
        let mut cfg = proxy_cfg(vec!["api.example.com"], upstream_addr);
        cfg.upstream_timeout = Duration::from_millis(120);
        let emitter = std::sync::Arc::new(MemEmitter::default());
        let shutdown = std::sync::Arc::new(AtomicBool::new(false));

        let proxy_handle = {
            let emitter = emitter.clone();
            let shutdown = shutdown.clone();
            let cfg = cfg.clone();
            std::thread::spawn(move || {
                let _ = run_one_shot(&cfg, &listener, &upstream_sock, &*emitter, &shutdown);
            })
        };

        let client = UdpSocket::bind("127.0.0.1:0").unwrap();
        client
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        let q = build_query_packet("api.example.com", 1);
        client.send_to(&q, listen_addr).unwrap();
        let mut rb = [0u8; 1500];
        let (_n, _) = client.recv_from(&mut rb).unwrap();
        assert_eq!(rb[3] & 0x0f, 2, "expected SERVFAIL on upstream timeout");

        shutdown.store(true, Ordering::SeqCst);
        proxy_handle.join().unwrap();
        let evs = emitter.events.lock().unwrap();
        // SEAM-1 Phase 3 — allowlist match still fires `query_permitted`
        // before the upstream forward; the upstream-timeout aggregate
        // event then follows. No `query_refused` event is emitted
        // because the proxy DECISION was allow — the SERVFAIL is an
        // upstream-side failure, not a policy refusal.
        assert_eq!(evs.len(), 2);
        assert_eq!(evs[0].ty, "dev.cellos.events.cell.dns.v1.query_permitted");
        let data = evs[1].data.as_ref().unwrap();
        assert_eq!(data["reasonCode"], "upstream_failure");
        assert_eq!(data["responseRcode"], 2);
    }

    #[test]
    fn parse_response_target_count_counts_a_records() {
        let q = build_query_packet("api.example.com", 1);
        let r = build_a_response(&q, 3);
        assert_eq!(parse_response_target_count(&r, 1), 3);
        // Non-A query type returns 0 even with answers present.
        assert_eq!(parse_response_target_count(&r, 16), 0);
    }
}