AkurAI-Framework / crates/http/src/server.rs

Latest change 5458532ad0ba451030d0554d03be8ff383ccddb8 - feat(ws,http): pure-std WebSocket (RFC 6455) + Server-Sent Events by Ólafur Búi Ólafsson

//! The blocking, thread-pooled HTTP/1.1 server.
//!
//! [`Server::bind`] grabs the socket (so the caller can read the real local
//! address before serving — useful for `:0` ephemeral ports in tests), and
//! [`Server::run`] accepts connections, parses each request head, calls the
//! handler, and writes the response.

use crate::pool::ThreadPool;
use crate::{Reply, Request, Response};
use std::io::{self, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs};
use std::sync::Arc;
use std::thread;

/// Largest request head we will buffer before giving up (64 KiB). Bodies are
/// not read here — static GET serving doesn't need them yet.
const MAX_HEAD_BYTES: usize = 64 * 1024;

/// A handler turns a parsed request into a [`Reply`] — either a buffered
/// [`Response`] or a connection [`Upgrade`](crate::Upgrade). `Send + Sync` so it
/// can be shared across worker threads. A closure returning a plain `Response`
/// works directly: `Response` converts into `Reply` via `From`.
pub trait Handler: Send + Sync + 'static {
    fn handle(&self, req: &Request) -> Reply;
}

impl<F, R> Handler for F
where
    F: Fn(&Request) -> R + Send + Sync + 'static,
    R: Into<Reply>,
{
    fn handle(&self, req: &Request) -> Reply {
        self(req).into()
    }
}

pub struct Server {
    listener: TcpListener,
    workers: usize,
}

impl Server {
    /// Bind without serving yet. Use `"127.0.0.1:0"` to get an OS-assigned port.
    pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<Server> {
        Ok(Server {
            listener: TcpListener::bind(addr)?,
            workers: default_workers(),
        })
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.listener.local_addr()
    }

    /// Accept connections forever, dispatching each to the handler on a worker
    /// thread. Blocks the calling thread.
    pub fn run<H: Handler>(self, handler: H) -> io::Result<()> {
        let pool = ThreadPool::new(self.workers);
        let handler = Arc::new(handler);
        for stream in self.listener.incoming() {
            let stream = stream?;
            let handler = Arc::clone(&handler);
            pool.execute(move || {
                if let Err(e) = serve_connection(stream, handler.as_ref()) {
                    // A broken pipe / reset mid-write is normal; don't crash.
                    let _ = e;
                }
            });
        }
        Ok(())
    }
}

fn serve_connection<H: Handler>(mut stream: TcpStream, handler: &H) -> io::Result<()> {
    let head = match read_head(&mut stream)? {
        Some(bytes) => bytes,
        None => return Ok(()), // connection closed before a full head arrived
    };

    let reply = match Request::parse_head(&head) {
        Ok(req) => handler.handle(&req),
        Err(_) => Reply::Response(Response::bad_request()),
    };

    match reply {
        Reply::Response(response) => {
            stream.write_all(&response.to_bytes())?;
            stream.flush()
        }
        Reply::Upgrade(up) => {
            // Send the head (101 / event-stream …), then hand the socket to the
            // protocol. We move it to a *dedicated* thread so a long-lived
            // WebSocket or SSE stream doesn't pin a pool worker for its whole
            // life — the request pool stays free to accept new connections.
            stream.write_all(&up.head.to_head_bytes())?;
            stream.flush()?;
            let hijack = up.hijack;
            thread::Builder::new()
                .name("akurai-upgrade".into())
                .spawn(move || {
                    let _ = hijack(stream); // a dropped peer is normal; don't crash
                })?;
            Ok(())
        }
    }
}

/// Read until the blank line that ends the request head. Returns `None` if the
/// peer closed cleanly first, and errors if the head exceeds [`MAX_HEAD_BYTES`].
fn read_head(stream: &mut TcpStream) -> io::Result<Option<Vec<u8>>> {
    let mut buf = Vec::with_capacity(1024);
    let mut chunk = [0u8; 1024];
    loop {
        if find_head_end(&buf).is_some() {
            return Ok(Some(buf));
        }
        if buf.len() > MAX_HEAD_BYTES {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "request head too large",
            ));
        }
        let n = stream.read(&mut chunk)?;
        if n == 0 {
            return Ok(if buf.is_empty() { None } else { Some(buf) });
        }
        buf.extend_from_slice(&chunk[..n]);
    }
}

/// Index just past the header terminator (`\r\n\r\n`, or bare `\n\n`).
fn find_head_end(buf: &[u8]) -> Option<usize> {
    buf.windows(4)
        .position(|w| w == b"\r\n\r\n")
        .map(|i| i + 4)
        .or_else(|| buf.windows(2).position(|w| w == b"\n\n").map(|i| i + 2))
}

fn default_workers() -> usize {
    thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(4)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::TcpStream;

    #[test]
    fn serves_a_request_end_to_end() {
        let server = Server::bind("127.0.0.1:0").unwrap();
        let addr = server.local_addr().unwrap();

        std::thread::spawn(move || {
            server
                .run(|req: &Request| Response::ok().with_text(&format!("path={}", req.path)))
                .unwrap();
        });

        let mut conn = TcpStream::connect(addr).unwrap();
        conn.write_all(b"GET /hello HTTP/1.1\r\nHost: x\r\n\r\n")
            .unwrap();
        let mut resp = String::new();
        conn.read_to_string(&mut resp).unwrap();

        assert!(resp.starts_with("HTTP/1.1 200 OK\r\n"), "got: {resp:?}");
        assert!(resp.ends_with("path=/hello"));
    }

    #[test]
    fn finds_head_terminators() {
        assert_eq!(find_head_end(b"GET / HTTP/1.1\r\n\r\n"), Some(18));
        assert_eq!(find_head_end(b"GET / HTTP/1.1\n\n"), Some(16));
        assert_eq!(find_head_end(b"incomplete\r\n"), None);
    }
}