conn.rs

//! Connection state machine.
//!
//! Drives one UDS connection: read request → call service → write
//! response → loop (if keep-alive).

// Implementation notes (not user-facing):
//
// * The state machine is an `async fn` (`serve`) kept inside
//   `Connection` as a boxed future — the per-connection allocation
//   is negligible compared to the syscall costs we are dominating.
// * `Connection` is generic over the inner future type `F` rather
//   than erasing it behind `dyn Future`, so `Send`/`Sync` auto
//   traits flow through naturally (mirroring hyper's conditional
//   `impl<T, S> Send for Connection<T, S> where S: Send, …`).

use std::future::Future;
use std::io;
use std::os::fd::AsFd;
use std::os::unix::net::UnixStream as StdUnixStream;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Poll};

use bytes::{Buf, BytesMut};

use crate::codec::{self, BodyLength, DecodeStatus};
use crate::error::{Error, Result};
use crate::fd_stream::FdStream;
use crate::message::Request;
use crate::service::Service;

/// Tunable knobs shared between [`crate::Builder`] and the connection.
#[derive(Debug, Clone)]
pub(crate) struct Config {
    pub keep_alive: bool,
    pub max_head_size: usize,
    pub max_body_size: usize,
    pub max_fds_per_request: usize,
    pub initial_read_buf: usize,
    pub initial_write_buf: usize,
    pub recv_chunk: usize,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            keep_alive: true,
            max_head_size: 16 * 1024,
            max_body_size: 16 * 1024 * 1024,
            max_fds_per_request: 16,
            initial_read_buf: 8 * 1024,
            initial_write_buf: 8 * 1024,
            recv_chunk: 8 * 1024,
        }
    }
}

/// A future representing the lifetime of one served UDS connection.
///
/// Returned by [`crate::Builder::serve_connection`]. Polling it drives
/// HTTP I/O; resolution means the connection has been closed (cleanly
/// or via error).
///
/// `Connection` is `Send` whenever the user's [`Service`] and its
/// `Future` are `Send`, so it can be spawned on a multi-threaded
/// runtime.
// `F` is the anonymous future produced by `serve`; auto traits
// (`Send`/`Sync`/`Unpin`) propagate from it through `Pin<Box<F>>`.
pub struct Connection<F> {
    fut: Pin<Box<F>>,
    shutdown: Arc<AtomicBool>,
}

impl<F> Connection<F>
where
    F: Future<Output = Result<()>>,
{
    pub(crate) fn new(fut: F, shutdown: Arc<AtomicBool>) -> Self {
        Self {
            fut: Box::pin(fut),
            shutdown,
        }
    }

    /// Request the connection to stop accepting new requests after
    /// the in-flight one (if any) finishes. The connection future
    /// must continue to be polled until it resolves.
    pub fn graceful_shutdown(&self) {
        self.shutdown.store(true, Ordering::Release);
    }
}

impl<F> Future for Connection<F>
where
    F: Future<Output = Result<()>>,
{
    type Output = Result<()>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // `Pin<Box<F>>` and `Arc<AtomicBool>` are both `Unpin`, so
        // `Connection<F>: Unpin` for all `F` and projection is trivial.
        self.get_mut().fut.as_mut().poll(cx)
    }
}

// Async driver implementing the per-connection state machine.
pub(crate) async fn serve<S>(
    cfg: Config,
    stream: StdUnixStream,
    service: S,
    shutdown: Arc<AtomicBool>,
) -> Result<()>
where
    S: Service,
{
    let mut io = FdStream::from_std(stream)?;
    let mut read_buf = BytesMut::with_capacity(cfg.initial_read_buf);
    let mut write_buf: Vec<u8> = Vec::with_capacity(cfg.initial_write_buf);

    loop {
        if shutdown.load(Ordering::Acquire) {
            return Ok(());
        }

        // ---- Read head ----
        let (head, consumed) = match read_head(&mut io, &mut read_buf, &cfg).await? {
            ReadHeadOutcome::Head { head, consumed } => (head, consumed),
            ReadHeadOutcome::CleanEof => return Ok(()),
        };
        read_buf.advance(consumed);

        // ---- Read body ----
        let body_len = match head.body_len {
            BodyLength::Known(n) => {
                if n as usize > cfg.max_body_size {
                    return Err(Error::BodyTooLarge);
                }
                n as usize
            }
            BodyLength::None => 0,
        };
        while read_buf.len() < body_len {
            let need = body_len - read_buf.len();
            let want = need.min(cfg.recv_chunk);
            read_buf.reserve(want);
            // SAFETY: we read into a temporary stack/heap-backed slice
            // and only extend the BytesMut by the bytes actually read.
            let mut tmp = vec![0u8; want];
            let n = io.recv(&mut tmp).await?;
            if n == 0 {
                return Err(Error::UnexpectedEof);
            }
            read_buf.extend_from_slice(&tmp[..n]);
        }
        let body = codec::split_body(&mut read_buf, body_len);

        // ---- Take fds ----
        let fds = io.take_fds();
        if fds.len() > cfg.max_fds_per_request {
            return Err(Error::TooManyFds);
        }
        if head.fd_names.len() > fds.len() {
            return Err(Error::InvalidXFdHeader);
        }
        // If client attached more fds than X-FD names, we keep them
        // (caller-defined), but truncate fd_names alignment view.
        // For now, surface only as many fds as fd_names declares;
        // any surplus is dropped (closed) here.
        let (fds, _surplus) = if head.fd_names.len() < fds.len() {
            let mut fds = fds;
            let surplus = fds.split_off(head.fd_names.len());
            (fds, surplus)
        } else {
            (fds, Vec::new())
        };

        let keep_alive = head.keep_alive && cfg.keep_alive;
        let req = Request {
            method: head.method,
            uri: head.uri,
            version: head.version,
            headers: head.headers,
            fd_names: head.fd_names,
            fds,
            body,
        };

        // ---- Dispatch ----
        let resp = service
            .call(req)
            .await
            .map_err(|e| Error::Service(e.into()))?;

        // ---- Encode response ----
        write_buf.clear();
        codec::encode_response(&resp, keep_alive, &mut write_buf);

        // Borrow the response fds (after this point we still own them
        // via `resp` so the BorrowedFd is valid).
        let borrowed: Vec<_> = resp.fds.iter().map(|f| f.as_fd()).collect();

        // sendmsg loop: cmsg goes with the very first call.
        let mut sent = 0usize;
        let mut first = true;
        while sent < write_buf.len() {
            let chunk = &write_buf[sent..];
            let n = if first {
                first = false;
                io.send(chunk, &borrowed).await?
            } else {
                io.send(chunk, &[]).await?
            };
            if n == 0 {
                return Err(Error::Io(io::Error::new(
                    io::ErrorKind::WriteZero,
                    "sendmsg returned 0",
                )));
            }
            sent += n;
        }
        // `borrowed` (and the `resp.fds` it borrows from) are released
        // here when the function-scope `resp` is dropped at the end of
        // this iteration; no explicit drops needed.
        drop(borrowed);
        drop(resp);

        if !keep_alive {
            return Ok(());
        }
    }
}

// PERF: see the same rationale on `codec::DecodeStatus`. `Head` carries
// the inline `DecodedHead` to avoid a heap allocation per parsed request.
#[allow(clippy::large_enum_variant)]
enum ReadHeadOutcome {
    Head {
        head: codec::DecodedHead,
        consumed: usize,
    },
    /// Peer closed cleanly with no in-flight data.
    CleanEof,
}

async fn read_head(
    io: &mut FdStream,
    read_buf: &mut BytesMut,
    cfg: &Config,
) -> Result<ReadHeadOutcome> {
    loop {
        match codec::decode_head(read_buf.as_ref())? {
            DecodeStatus::Complete { head, consumed } => {
                return Ok(ReadHeadOutcome::Head { head, consumed });
            }
            DecodeStatus::Partial => {
                if read_buf.len() >= cfg.max_head_size {
                    return Err(Error::HeadTooLarge);
                }
                let want = cfg.recv_chunk.min(cfg.max_head_size - read_buf.len());
                let mut tmp = vec![0u8; want];
                let n = io.recv(&mut tmp).await?;
                if n == 0 {
                    return if read_buf.is_empty() {
                        Ok(ReadHeadOutcome::CleanEof)
                    } else {
                        Err(Error::UnexpectedEof)
                    };
                }
                read_buf.extend_from_slice(&tmp[..n]);
            }
        }
    }
}