urus/src/conn.rs

//! `Conn` — the value that flows through the entire plug pipeline.
//!
//! Carries request data, response state being built, and arbitrary user
//! data. Owned and moved at each step. Plugs are unaware of the underlying
//! HTTP version (v1.1 today, /2 later).

use std::any::{Any, TypeId};
use std::collections::HashMap;

// ---------------------------------------------------------------------------
// Method
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Method {
    Get,
    Post,
    Put,
    Delete,
    Patch,
    Head,
    Options,
    Other(String),
}

impl Method {
    pub fn parse(s: &str) -> Self {
        match s {
            "GET"     => Method::Get,
            "POST"    => Method::Post,
            "PUT"     => Method::Put,
            "DELETE"  => Method::Delete,
            "PATCH"   => Method::Patch,
            "HEAD"    => Method::Head,
            "OPTIONS" => Method::Options,
            _         => Method::Other(s.to_string()),
        }
    }

    pub fn as_str(&self) -> &str {
        match self {
            Method::Get     => "GET",
            Method::Post    => "POST",
            Method::Put     => "PUT",
            Method::Delete  => "DELETE",
            Method::Patch   => "PATCH",
            Method::Head    => "HEAD",
            Method::Options => "OPTIONS",
            Method::Other(s) => s.as_str(),
        }
    }
}

// ---------------------------------------------------------------------------
// HttpVersion
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HttpVersion {
    Http10,
    Http11,
}

impl HttpVersion {
    pub fn as_str(self) -> &'static str {
        match self {
            HttpVersion::Http10 => "HTTP/1.0",
            HttpVersion::Http11 => "HTTP/1.1",
        }
    }
}

// ---------------------------------------------------------------------------
// HeaderMap — small-vec of (name, value) string pairs.
// ---------------------------------------------------------------------------
//
// Per spec: most requests have <20 headers, so a linear-scan Vec is cheaper
// than a hash map. Names are case-insensitive (HTTP/1.1), normalised to
// lowercase on insert.

#[derive(Debug, Clone, Default)]
pub struct HeaderMap {
    inner: Vec<(String, String)>,
}

impl HeaderMap {
    pub fn new() -> Self { Self::default() }

    pub fn with_capacity(n: usize) -> Self {
        Self { inner: Vec::with_capacity(n) }
    }

    /// Append a header. Does not deduplicate (HTTP allows repeats; e.g.
    /// `Set-Cookie`).
    pub fn append(&mut self, name: &str, value: impl Into<String>) {
        self.inner.push((name.to_ascii_lowercase(), value.into()));
    }

    /// Replace any existing values for `name` with a single value.
    pub fn set(&mut self, name: &str, value: impl Into<String>) {
        let lower = name.to_ascii_lowercase();
        self.inner.retain(|(n, _)| n != &lower);
        self.inner.push((lower, value.into()));
    }

    /// First value for `name`, if any.
    pub fn get(&self, name: &str) -> Option<&str> {
        let lower = name.to_ascii_lowercase();
        self.inner.iter()
            .find(|(n, _)| n == &lower)
            .map(|(_, v)| v.as_str())
    }

    /// Iterate over (name, value) pairs in insertion order.
    pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
        self.inner.iter().map(|(n, v)| (n.as_str(), v.as_str()))
    }

    pub fn len(&self) -> usize { self.inner.len() }
    pub fn is_empty(&self) -> bool { self.inner.is_empty() }
}

// ---------------------------------------------------------------------------
// Body — request body, exposed as owned bytes.
// ---------------------------------------------------------------------------
//
// v1: bodies are read in full by the connection actor before the pipeline
// runs, when `Content-Length` is set. Streaming and chunked encoding will
// come later; for now this is the simplest correct thing.

#[derive(Debug, Clone, Default)]
pub struct Body {
    bytes: Vec<u8>,
}

impl Body {
    pub fn empty() -> Self { Self::default() }
    pub fn from_bytes(b: Vec<u8>) -> Self { Self { bytes: b } }
    pub fn as_bytes(&self) -> &[u8] { &self.bytes }
    pub fn into_bytes(self) -> Vec<u8> { self.bytes }
    pub fn len(&self) -> usize { self.bytes.len() }
    pub fn is_empty(&self) -> bool { self.bytes.is_empty() }
}

// ---------------------------------------------------------------------------
// RespBody — what plugs put on the wire.
// ---------------------------------------------------------------------------
//
// Enum, not a trait, so the connection actor can pattern-match.
//
// `Stream` is the pull-shaped streaming body (v0.3 design decision): the
// handler hands back a `smarm::Receiver<Vec<u8>>` — typically the read end
// of a channel whose `Sender` lives in a producer actor the handler
// spawned — and the connection actor pumps it onto the wire. The conn
// actor keeps sole ownership of the socket and of write deadlines; the
// producer never touches the fd. End-of-stream is signalled by dropping
// every `Sender` (the conn actor then emits the terminating 0-chunk).
// Empty `Vec`s are skipped by the pump (a zero-length chunk would
// terminate chunked framing early), so they are safe to send but useless.

pub enum RespBody {
    Empty,
    Bytes(Vec<u8>),
    Stream(StreamBody),
}

/// A streaming response body. Construct with [`StreamBody::new`] (or
/// `RespBody::from(rx)`) and hand it to [`Conn::put_body`].
///
/// On HTTP/1.1 the connection actor sends it with
/// `Transfer-Encoding: chunked` and the connection stays reusable after
/// the stream completes. On HTTP/1.0 (no chunked framing) the bytes are
/// written raw and the connection closes at end of stream to delimit the
/// body.
pub struct StreamBody {
    pub rx: smarm::Receiver<Vec<u8>>,
    /// If set, the connection actor's pump waits for each chunk with
    /// `recv_timeout(interval)` instead of a bare `recv`, and on expiry
    /// writes `payload` as its own chunk and keeps waiting. This is how
    /// SSE keep-alive comments work: liveness comes from the ping (a dead
    /// client is detected when a heartbeat write stalls past
    /// `write_timeout`), not from any request clock.
    pub heartbeat: Option<(std::time::Duration, Vec<u8>)>,
}

impl StreamBody {
    pub fn new(rx: smarm::Receiver<Vec<u8>>) -> Self {
        Self { rx, heartbeat: None }
    }

    pub fn with_heartbeat(mut self, interval: std::time::Duration, payload: Vec<u8>) -> Self {
        self.heartbeat = Some((interval, payload));
        self
    }
}

impl std::fmt::Debug for RespBody {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RespBody::Empty     => f.write_str("Empty"),
            RespBody::Bytes(b)  => f.debug_tuple("Bytes").field(&b.len()).finish(),
            RespBody::Stream(_) => f.write_str("Stream(..)"),
        }
    }
}

impl RespBody {
    /// Body length for fixed bodies; 0 for `Stream` (a stream has no
    /// length up front — that is the point — and is never serialised with
    /// a `content-length`).
    pub fn len_hint(&self) -> usize {
        match self {
            RespBody::Empty     => 0,
            RespBody::Bytes(b)  => b.len(),
            RespBody::Stream(_) => 0,
        }
    }
}

impl Default for RespBody {
    fn default() -> Self { RespBody::Empty }
}

// ---------------------------------------------------------------------------
// Params — path parameters extracted by the router.
// ---------------------------------------------------------------------------
//
// Small-vec semantics: routes have a handful of params; a Vec is the right
// shape. Empty until the router plug populates it.

#[derive(Debug, Clone, Default)]
pub struct Params {
    inner: Vec<(String, String)>,
}

impl Params {
    pub fn new() -> Self { Self::default() }

    pub fn put(&mut self, name: impl Into<String>, value: impl Into<String>) {
        self.inner.push((name.into(), value.into()));
    }

    pub fn get(&self, name: &str) -> Option<&str> {
        self.inner.iter()
            .find(|(n, _)| n == name)
            .map(|(_, v)| v.as_str())
    }

    pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
        self.inner.iter().map(|(n, v)| (n.as_str(), v.as_str()))
    }
}

// ---------------------------------------------------------------------------
// Assigns — type-erased map for arbitrary user data.
// ---------------------------------------------------------------------------
//
// Same pattern as Phoenix's `conn.assigns`. The `Option<HashMap>` keeps the
// common case (no assigns) at zero allocation cost.

#[derive(Debug, Default)]
pub struct Assigns {
    inner: Option<HashMap<TypeId, Box<dyn Any + Send>>>,
}

impl Assigns {
    pub fn new() -> Self { Self::default() }

    pub fn put<T: Any + Send>(&mut self, value: T) {
        self.inner
            .get_or_insert_with(HashMap::new)
            .insert(TypeId::of::<T>(), Box::new(value));
    }

    pub fn get<T: Any + Send>(&self) -> Option<&T> {
        self.inner
            .as_ref()?
            .get(&TypeId::of::<T>())?
            .downcast_ref::<T>()
    }

    pub fn take<T: Any + Send>(&mut self) -> Option<T> {
        let boxed = self.inner.as_mut()?.remove(&TypeId::of::<T>())?;
        boxed.downcast::<T>().ok().map(|b| *b)
    }
}

// ---------------------------------------------------------------------------
// Conn — the pipeline value.
// ---------------------------------------------------------------------------

#[derive(Debug)]
pub struct Conn {
    // Request (populated before the pipeline runs).
    pub method:  Method,
    pub path:    String,
    pub query:   Option<String>,
    pub version: HttpVersion,
    pub headers: HeaderMap,
    pub body:    Body,

    // Response (built up by plugs, written after the pipeline returns).
    pub status:       Option<u16>,
    pub resp_headers: HeaderMap,
    pub resp_body:    RespBody,

    // Pipeline state.
    pub params:  Params,
    pub assigns: Assigns,
    pub halted:  bool,

    /// Set by [`Conn::upgrade`] on an accepted WebSocket handshake; the
    /// connection actor sees it (with status 101) and leaves the HTTP
    /// loop after writing the response. `None` for plain HTTP.
    pub upgrade: Option<crate::ws::WsUpgrade>,
}

impl Conn {
    /// Construct an empty `Conn`. Tests use this; the connection actor uses
    /// `From<RequestParts>` (see parser.rs).
    pub fn new() -> Self {
        Self {
            method:  Method::Get,
            path:    String::new(),
            query:   None,
            version: HttpVersion::Http11,
            headers: HeaderMap::new(),
            body:    Body::empty(),

            status:       None,
            resp_headers: HeaderMap::new(),
            resp_body:    RespBody::Empty,

            params:  Params::new(),
            assigns: Assigns::new(),
            halted:  false,

            upgrade: None,
        }
    }

    // ----- Fluent builders. Most handler code uses these. -----

    pub fn put_status(mut self, status: u16) -> Self {
        self.status = Some(status);
        self
    }

    pub fn put_header(mut self, name: &str, value: impl Into<String>) -> Self {
        self.resp_headers.set(name, value);
        self
    }

    pub fn put_body(mut self, body: impl Into<RespBody>) -> Self {
        self.resp_body = body.into();
        self
    }

    pub fn put_params(mut self, params: Params) -> Self {
        self.params = params;
        self
    }

    pub fn halt(mut self) -> Self {
        self.halted = true;
        self
    }

    /// Accept a WebSocket upgrade on this request (RFC 6455 §4.2),
    /// handing `handler` to the duplex loop that takes the socket over.
    ///
    /// On a valid handshake: status `101`, the `upgrade`/`connection`/
    /// `sec-websocket-accept` response headers, the [`WsUpgrade`] payload
    /// the connection actor acts on, and the pipeline is halted. After
    /// the 101 the connection actor leaves HTTP and drives `handler` —
    /// see [`WsHandler`](crate::ws::WsHandler) for the callback contract.
    ///
    /// On an invalid handshake: `handler` is dropped untouched and the
    /// appropriate rejection goes out (`400`, or `426` +
    /// `sec-websocket-version: 13` for a version mismatch), empty body,
    /// pipeline halted; the connection stays plain HTTP with normal
    /// keep-alive semantics. Pre-check with
    /// [`ws::handshake::is_upgrade_request`](crate::ws::handshake::is_upgrade_request)
    /// to branch before committing.
    ///
    /// [`WsUpgrade`]: crate::ws::WsUpgrade
    pub fn upgrade(mut self, handler: impl crate::ws::WsHandler) -> Self {
        use crate::ws::{handshake, Rejection};
        match handshake::validate(&self) {
            Ok(accept) => {
                self.upgrade = Some(crate::ws::WsUpgrade {
                    handler: Box::new(handler),
                });
                self.put_status(101)
                    .put_header("upgrade", "websocket")
                    .put_header("connection", "Upgrade")
                    .put_header("sec-websocket-accept", accept)
                    .put_body(RespBody::Empty)
                    .halt()
            }
            Err(rej) => {
                let c = self
                    .put_status(rej.status())
                    .put_body(RespBody::Empty)
                    .halt();
                if rej == Rejection::UnsupportedVersion {
                    c.put_header("sec-websocket-version", "13")
                } else {
                    c
                }
            }
        }
    }
}

impl Default for Conn {
    fn default() -> Self { Self::new() }
}

// ----- Ergonomic conversions into RespBody. -----

impl From<Vec<u8>> for RespBody {
    fn from(v: Vec<u8>) -> Self { RespBody::Bytes(v) }
}
impl From<String> for RespBody {
    fn from(s: String) -> Self { RespBody::Bytes(s.into_bytes()) }
}
impl From<&'static str> for RespBody {
    fn from(s: &'static str) -> Self { RespBody::Bytes(s.as_bytes().to_vec()) }
}
impl From<&[u8]> for RespBody {
    fn from(s: &[u8]) -> Self { RespBody::Bytes(s.to_vec()) }
}
impl From<StreamBody> for RespBody {
    fn from(s: StreamBody) -> Self { RespBody::Stream(s) }
}
impl From<smarm::Receiver<Vec<u8>>> for RespBody {
    fn from(rx: smarm::Receiver<Vec<u8>>) -> Self {
        RespBody::Stream(StreamBody::new(rx))
    }
}