cellos_supervisor/dns_proxy/

spawn.rs

//! SEAM-1 Phase 2b — supervisor → cell-netns DNS proxy spawn.
//!
//! This module owns the *placement* of [`super::run_one_shot`] inside the
//! cell's network namespace. The proxy module itself ([`super::mod`]) is
//! platform-neutral and contains no `setns` / `nsenter` calls; the supervisor
//! pre-binds sockets and hands them to `run_one_shot`. Phase 2b's job is to
//! make those sockets exist in `/proc/<child_pid>/ns/net` rather than the host
//! netns.
//!
//! ## Why a dedicated OS thread
//!
//! `setns(2)` changes the *calling thread*'s namespace association. A
//! `tokio::task::spawn_blocking` task does not give us a stable thread —
//! tokio's blocking pool may reuse the worker for other tasks once we
//! return. Even if the proxy ran sync inside `spawn_blocking`, polluting
//! that worker's netns with the cell's namespace would corrupt every
//! subsequent task that landed on the same worker. We therefore allocate a
//! fresh `std::thread`, do the `setns` there, run the loop to completion,
//! and let the OS reclaim the thread on exit. The thread never returns to
//! a tokio worker.
//!
//! ## Bridging the sync emitter to async [`EventSink`]
//!
//! The proxy's [`super::DnsQueryEmitter`] is synchronous (it must not block
//! the recv loop). The supervisor's [`cellos_core::ports::EventSink`] is
//! async. [`EventSinkEmitter`] captures a [`tokio::runtime::Handle`] **before**
//! the OS thread starts and uses `Handle::spawn` to fire-and-forget every
//! event onto the runtime. This preserves the per-query event audit trail
//! without blocking the recv loop on emit completion.
//!
//! ## Shutdown coordination
//!
//! `run_one_shot` already checks an `AtomicBool` between iterations and
//! the listener has a short read_timeout. The teardown helper
//! [`signal_proxy_shutdown`] sets the flag *and* sends a 0-byte UDP
//! packet to the listen address — the recv loop wakes immediately,
//! observes the flag, and returns. Without the wake packet, the loop
//! would sit in `recv_from` for up to `read_timeout` after teardown
//! requested shutdown.

use std::net::SocketAddr;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;

use cellos_core::CloudEventV1;

#[cfg(target_os = "linux")]
use super::DnsProxyConfig;
use super::{DnsProxyStats, DnsQueryEmitter};

/// Read timeout on the proxy's listener socket. Bounds the worst-case
/// shutdown latency when the wake packet is somehow lost (e.g. the test
/// teardown didn't reach `signal_proxy_shutdown`). Kept short so the
/// `AtomicBool` check happens at a deterministic cadence.
#[cfg(target_os = "linux")]
const LISTENER_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(200);

/// Adapter from the proxy's synchronous [`DnsQueryEmitter`] trait to the
/// supervisor's async [`cellos_core::ports::EventSink`] pipeline.
///
/// **Tokio-context constraint:** `runtime_handle` MUST be captured on a
/// thread that already has a tokio runtime context (i.e. inside an `async`
/// block on the multi-thread runtime, OR inside `spawn_blocking`). The
/// proxy thread itself is a bare `std::thread` and has no tokio context;
/// the captured handle is what lets it `spawn` work back onto the runtime.
pub struct EventSinkEmitter {
    pub runtime_handle: tokio::runtime::Handle,
    pub sink: Arc<dyn cellos_core::ports::EventSink>,
    pub jsonl_sink: Option<Arc<dyn cellos_core::ports::EventSink>>,
}

impl EventSinkEmitter {
    /// Build an emitter capturing the **current** tokio runtime handle.
    ///
    /// Call from inside an async block (or `spawn_blocking`) — never from
    /// a bare std thread. Panics with a clear message if no tokio runtime
    /// is reachable, which is a programming error: this struct exists
    /// specifically to bridge from sync threads, but the bridge itself
    /// must be constructed where a runtime is available.
    pub fn capture_current(
        sink: Arc<dyn cellos_core::ports::EventSink>,
        jsonl_sink: Option<Arc<dyn cellos_core::ports::EventSink>>,
    ) -> Self {
        let handle = tokio::runtime::Handle::try_current()
            .expect("EventSinkEmitter::capture_current called outside a tokio runtime context");
        Self {
            runtime_handle: handle,
            sink,
            jsonl_sink,
        }
    }
}

impl DnsQueryEmitter for EventSinkEmitter {
    fn emit(&self, event: CloudEventV1) {
        // Fire-and-forget: the proxy recv loop runs synchronously and any
        // backpressure here turns directly into per-query latency for the
        // workload. The supervisor's emit() helper logs sink failures —
        // mirror that here without blocking the proxy.
        let sink = self.sink.clone();
        let jsonl = self.jsonl_sink.clone();
        let event_for_jsonl = event.clone();
        self.runtime_handle.spawn(async move {
            if let Err(e) = sink.emit(&event).await {
                tracing::warn!(
                    target: "cellos.supervisor.dns_proxy",
                    error = %e,
                    "primary sink emit failed for dns_query event"
                );
            }
        });
        if let Some(j) = jsonl {
            self.runtime_handle.spawn(async move {
                if let Err(e) = j.emit(&event_for_jsonl).await {
                    tracing::warn!(
                        target: "cellos.supervisor.dns_proxy",
                        error = %e,
                        "jsonl sink emit failed for dns_query event"
                    );
                }
            });
        }
    }
}

/// Shutdown coordination handle returned by [`spawn_proxy_in_netns`] (Linux)
/// — held by the supervisor and signalled at cell destroy time.
///
/// On non-Linux platforms the supervisor never constructs one of these; the
/// activation predicate evaluator simply returns `None`.
pub struct DnsProxyHandle {
    /// Shared shutdown flag. Setting `true` signals the proxy loop to exit
    /// at its next iteration (within `LISTENER_READ_TIMEOUT` worst case).
    pub shutdown: Arc<AtomicBool>,
    /// The address the proxy listener is bound to inside the cell's netns.
    /// Used by [`signal_proxy_shutdown`] to wake the recv loop.
    pub listen_addr: SocketAddr,
    /// Join handle for the proxy OS thread. `Some` until `join()` consumes it.
    #[cfg(target_os = "linux")]
    pub thread: Option<std::thread::JoinHandle<DnsProxyStats>>,
    /// Resolver id stamped into events; retained for diagnostics.
    pub upstream_resolver_id: String,
}

impl DnsProxyHandle {
    /// Signal the proxy to stop and join the OS thread. Returns the
    /// cumulative [`DnsProxyStats`] when the thread exited cleanly, or
    /// `None` if the thread had already been joined or panicked.
    ///
    /// Caller must already have set `self.shutdown` to `true` and called
    /// [`signal_proxy_shutdown`] (or otherwise woken the recv loop).
    /// Calling this without waking the loop will block for up to
    /// `LISTENER_READ_TIMEOUT` — bounded but not instant.
    #[cfg(target_os = "linux")]
    pub fn join(&mut self) -> Option<DnsProxyStats> {
        let handle = self.thread.take()?;
        match handle.join() {
            Ok(stats) => Some(stats),
            Err(_panic) => {
                tracing::warn!(
                    target: "cellos.supervisor.dns_proxy",
                    "DNS proxy thread panicked on join"
                );
                None
            }
        }
    }

    /// Stub join on non-Linux platforms — no thread was ever spawned.
    #[cfg(not(target_os = "linux"))]
    pub fn join(&mut self) -> Option<DnsProxyStats> {
        None
    }
}

/// Send a 0-byte UDP packet to `listen_addr` to wake a blocked `recv_from`.
///
/// The proxy's recv loop checks `shutdown` between iterations; without a
/// wake-up the loop would sit in `recv_from` for up to `LISTENER_READ_TIMEOUT`
/// after the supervisor flipped the flag. The wake packet collapses that
/// window to ~milliseconds.
///
/// 0-byte UDP datagrams are valid and parse-rejected by `parse_query` (the
/// header alone is 12 bytes), producing one `dns_query` event with
/// `reasonCode: malformed_query`. This is intentional: the wake packet is
/// indistinguishable from any other malformed input the workload could
/// have sent and the audit trail records it the same way. **The proxy
/// thread will NOT emit this event after the loop has already terminated**
/// — practically, the wake packet arrives, recv returns, the loop iterates
/// once with the shutdown flag set to true, and exits without re-emitting.
/// (The loop checks the flag at the top of each iteration; the recv that
/// produced the wake packet completes, and on the next iteration the flag
/// breaks the loop.)
///
/// Best-effort: if the wake socket cannot be bound or sent (extremely
/// unlikely on localhost), the supervisor logs and falls back to
/// timeout-based shutdown.
pub fn signal_proxy_shutdown(listen_addr: SocketAddr) {
    // Bind on a system-chosen port on the same family as the listener.
    let bind_str = if listen_addr.is_ipv6() {
        "[::1]:0"
    } else {
        "127.0.0.1:0"
    };
    let waker = match std::net::UdpSocket::bind(bind_str) {
        Ok(s) => s,
        Err(e) => {
            tracing::debug!(
                target: "cellos.supervisor.dns_proxy",
                error = %e,
                "wake-up socket bind failed; falling back to timeout-based shutdown"
            );
            return;
        }
    };
    if let Err(e) = waker.send_to(&[], listen_addr) {
        tracing::debug!(
            target: "cellos.supervisor.dns_proxy",
            error = %e,
            addr = %listen_addr,
            "wake-up packet send failed; falling back to timeout-based shutdown"
        );
    }
}

/// Linux-only: spawn the proxy on a dedicated OS thread, `setns(2)` into
/// `/proc/<child_pid>/ns/net`, bind UDP listener + upstream sockets in that
/// netns, and run [`super::run_one_shot`] until `shutdown` is set.
///
/// Returns a [`DnsProxyHandle`] carrying the join handle + the address the
/// listener is bound to. The caller is responsible for joining the handle
/// during teardown ([`DnsProxyHandle::join`]) and for waking the recv loop
/// via [`signal_proxy_shutdown`] before joining.
///
/// **Errors:** any failure before the loop starts (open `/proc` fd, `setns`,
/// bind listener / upstream) is returned to the caller as
/// `std::io::Error`. The supervisor turns this into a `dns_query` event
/// with `reasonCode: upstream_failure` so the audit trail surfaces the
/// spawn failure.
///
/// **Activation contract:** `child_pid` MUST be a live process whose
/// `/proc/<pid>/ns/net` is the cell's network namespace. The supervisor
/// calls this helper inside `linux_run_cell_command_isolated` immediately
/// after `cmd.spawn()` returns and before `child.wait()` — the PID is
/// guaranteed alive at that point (the child is on a `pre_exec` path that
/// has not yet completed `execve`).
#[cfg(target_os = "linux")]
pub fn spawn_proxy_in_netns(
    child_pid: u32,
    cfg: DnsProxyConfig,
    listen_addr: SocketAddr,
    upstream_addr: SocketAddr,
    emitter: Arc<dyn DnsQueryEmitter>,
    shutdown: Arc<AtomicBool>,
) -> std::io::Result<DnsProxyHandle> {
    use std::fs::File;
    use std::os::unix::io::AsRawFd;

    // Open the netns FD on the calling thread. We pass the FD into the proxy
    // thread; setns(2) will be called there. Opening here means errors like
    // "child died before we could grab the namespace" surface synchronously
    // to the supervisor rather than from a thread it would have to join to
    // discover the failure.
    let netns_path = format!("/proc/{child_pid}/ns/net");
    let netns_file = File::open(&netns_path)
        .map_err(|e| std::io::Error::new(e.kind(), format!("open netns at {netns_path}: {e}")))?;

    let upstream_resolver_id = cfg.upstream_resolver_id.clone();
    let shutdown_for_thread = shutdown.clone();

    let (ready_tx, ready_rx) = std::sync::mpsc::channel::<std::io::Result<SocketAddr>>();

    let thread = std::thread::Builder::new()
        .name(format!("cellos-dns-proxy-{child_pid}"))
        .spawn(move || {
            // SAFETY: setns is the documented Linux syscall for moving the
            // calling thread into the namespace referenced by `fd`. We hold
            // `netns_file` for the lifetime of the thread so the fd remains
            // valid until the kernel's reference is taken; the thread never
            // returns to a tokio worker so polluting "the" thread's netns
            // is intentional and isolated.
            let setns_rc = unsafe { libc::setns(netns_file.as_raw_fd(), libc::CLONE_NEWNET) };
            if setns_rc != 0 {
                let err = std::io::Error::last_os_error();
                let _ = ready_tx.send(Err(std::io::Error::new(
                    err.kind(),
                    format!("setns(CLONE_NEWNET) for pid={child_pid}: {err}"),
                )));
                return DnsProxyStats::default();
            }
            // From this point on the thread is in the cell's netns.
            // Sockets bound here will live in that netns and the workload
            // can reach them via the IP it sees on its loopback / declared
            // bridge interface.
            let listener = match std::net::UdpSocket::bind(listen_addr) {
                Ok(s) => s,
                Err(e) => {
                    let _ = ready_tx.send(Err(std::io::Error::new(
                        e.kind(),
                        format!("bind listener at {listen_addr} in cell netns: {e}"),
                    )));
                    return DnsProxyStats::default();
                }
            };
            if let Err(e) = listener.set_read_timeout(Some(LISTENER_READ_TIMEOUT)) {
                let _ = ready_tx.send(Err(e));
                return DnsProxyStats::default();
            }
            // Upstream socket: ephemeral port in the same netns. The upstream
            // resolver MUST be reachable from inside this netns; SEC-22's
            // nft path ensures the declared resolver IP/port is allowed.
            let upstream_sock = match std::net::UdpSocket::bind(if upstream_addr.is_ipv6() {
                "[::]:0"
            } else {
                "0.0.0.0:0"
            }) {
                Ok(s) => s,
                Err(e) => {
                    let _ = ready_tx.send(Err(std::io::Error::new(
                        e.kind(),
                        format!("bind upstream socket in cell netns: {e}"),
                    )));
                    return DnsProxyStats::default();
                }
            };
            // Report bound address back so the supervisor can confirm
            // before it lets the cell command run.
            let actual_listen = match listener.local_addr() {
                Ok(a) => a,
                Err(e) => {
                    let _ = ready_tx.send(Err(e));
                    return DnsProxyStats::default();
                }
            };

            // Slot **A5** — TCP/53 listener path in the same netns.
            //
            // The TCP listener binds on the SAME `listen_addr` as UDP (per
            // RFC 1035 §4.2.2 a resolver MUST accept both transports on
            // the declared address). We bind it from the proxy thread,
            // which is already setns'd into the cell netns; sub-threads
            // spawned from here inherit the parent thread's net namespace
            // (CLONE_NEWNET is per-task and Linux clone() without it
            // copies the nsproxy reference), so the sibling worker
            // running `run_tcp_one_shot` sees the same netns as the UDP
            // loop on the main thread.
            //
            // Failure to bind TCP is non-fatal for the UDP path — many
            // workloads only resolve over UDP and degrading silently to
            // UDP-only is preferable to refusing activation. The error
            // is logged and we proceed UDP-only.
            let tcp_listener = match std::net::TcpListener::bind(listen_addr) {
                Ok(l) => Some(l),
                Err(e) => {
                    tracing::warn!(
                        target: "cellos.supervisor.dns_proxy",
                        error = %e,
                        addr = %listen_addr,
                        "TCP listener bind FAILED in cell netns — continuing UDP-only"
                    );
                    None
                }
            };

            if ready_tx.send(Ok(actual_listen)).is_err() {
                // Supervisor side dropped the receiver — abandon.
                return DnsProxyStats::default();
            }

            // A5 — spawn the TCP loop on a sibling thread before entering
            // the UDP loop on the current thread. The sibling inherits
            // the netns we setns'd into above. We split this into a
            // small helper closure so the multi-step "clone the upstream
            // socket, build the worker config, spawn the thread" chain
            // does not nest Options inside the OK arm of the map.
            let tcp_worker: Option<std::thread::JoinHandle<()>> = match tcp_listener {
                None => None,
                Some(listener_tcp) => {
                    match upstream_sock.try_clone() {
                        Ok(upstream_clone) => {
                            let cfg_tcp = cfg.clone();
                            let emitter_tcp: Arc<dyn DnsQueryEmitter> = emitter.clone();
                            let shutdown_tcp = shutdown_for_thread.clone();
                            // TCP per-connection workers each forward
                            // against the upstream socket. UDP is
                            // stateless per query so sharing is safe.
                            let upstream_tcp = Arc::new(upstream_clone);
                            std::thread::Builder::new()
                                .name(format!("cellos-dns-proxy-tcp-{child_pid}"))
                                .spawn(move || {
                                    if let Err(e) = super::run_tcp_one_shot(
                                        &cfg_tcp,
                                        &listener_tcp,
                                        upstream_tcp,
                                        emitter_tcp,
                                        &shutdown_tcp,
                                    ) {
                                        tracing::warn!(
                                            target: "cellos.supervisor.dns_proxy",
                                            error = %e,
                                            "TCP proxy accept loop returned with I/O error"
                                        );
                                    }
                                })
                                .ok()
                        }
                        Err(e) => {
                            tracing::warn!(
                                target: "cellos.supervisor.dns_proxy",
                                error = %e,
                                "TCP upstream socket clone FAILED — TCP path will not start"
                            );
                            None
                        }
                    }
                }
            };

            let udp_stats = match super::run_one_shot(
                &cfg,
                &listener,
                &upstream_sock,
                emitter.as_ref(),
                &shutdown_for_thread,
            ) {
                Ok(stats) => stats,
                Err(e) => {
                    tracing::warn!(
                        target: "cellos.supervisor.dns_proxy",
                        error = %e,
                        "proxy recv loop returned with I/O error"
                    );
                    DnsProxyStats::default()
                }
            };

            // A5 — UDP loop has exited (shutdown observed). Join the TCP
            // sibling so the OS thread is reaped before we return; the
            // TCP accept loop also polls `shutdown_for_thread` and will
            // be exiting at this point. Stats from the TCP path are
            // intentionally NOT merged into the returned UDP stats — the
            // existing teardown logging consumes UDP stats only, and
            // changing that aggregation is out of scope for A5.
            if let Some(h) = tcp_worker {
                let _ = h.join();
            }
            udp_stats
        })?;

    // Wait for the thread to confirm setns + bind succeeded before we let
    // the supervisor proceed. Bounded — the thread either reports within
    // ~tens of ms or something is wrong and we surface the error.
    let bound_addr = ready_rx
        .recv_timeout(std::time::Duration::from_secs(2))
        .map_err(|e| std::io::Error::other(format!("proxy thread ready timeout: {e}")))??;

    Ok(DnsProxyHandle {
        shutdown,
        listen_addr: bound_addr,
        thread: Some(thread),
        upstream_resolver_id,
    })
}

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

    /// Minimal in-memory sink for the EventSinkEmitter test.
    struct CountingSink {
        count: Mutex<u64>,
    }
    #[async_trait::async_trait]
    impl cellos_core::ports::EventSink for CountingSink {
        async fn emit(&self, _event: &CloudEventV1) -> Result<(), cellos_core::error::CellosError> {
            *self.count.lock().unwrap() += 1;
            Ok(())
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn event_sink_emitter_dispatches_to_runtime() {
        let sink = Arc::new(CountingSink {
            count: Mutex::new(0),
        });
        let emitter = EventSinkEmitter::capture_current(sink.clone(), None);
        let event = CloudEventV1 {
            specversion: "1.0".into(),
            id: "evt-1".into(),
            source: "test".into(),
            ty: "test.event".into(),
            datacontenttype: Some("application/json".into()),
            data: Some(serde_json::json!({"k": "v"})),
            time: Some(chrono::Utc::now().to_rfc3339()),
            traceparent: None,
        };
        emitter.emit(event);
        // Yield until the spawned task runs. fire-and-forget so we poll.
        for _ in 0..50 {
            if *sink.count.lock().unwrap() >= 1 {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
        assert_eq!(
            *sink.count.lock().unwrap(),
            1,
            "sink should have received one event via the runtime handle"
        );
    }

    #[test]
    fn signal_proxy_shutdown_does_not_panic_on_unbound_addr() {
        // Sending a wake packet to a port nothing is listening on must not
        // panic — UDP is fire-and-forget.
        let addr: SocketAddr = "127.0.0.1:1".parse().unwrap();
        signal_proxy_shutdown(addr);
    }

    #[test]
    fn handle_join_returns_none_when_no_thread() {
        let mut h = DnsProxyHandle {
            shutdown: Arc::new(AtomicBool::new(false)),
            listen_addr: "127.0.0.1:0".parse().unwrap(),
            #[cfg(target_os = "linux")]
            thread: None,
            upstream_resolver_id: "test".into(),
        };
        assert!(h.join().is_none());
        // shutdown is still readable/usable post-join.
        h.shutdown.store(true, Ordering::SeqCst);
        assert!(h.shutdown.load(Ordering::SeqCst));
    }
}