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

Latest change 3ad787ae157fc60f711c536c1811b3b199b0ac1a - init: AkurAI-Framework foundation (history reset after FS corruption) 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::{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 response. `Send + Sync` so it can be
/// shared across worker threads.
pub trait Handler: Send + Sync + 'static {
    fn handle(&self, req: &Request) -> Response;
}

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

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 response = match Request::parse_head(&head) {
        Ok(req) => handler.handle(&req),
        Err(_) => Response::bad_request(),
    };

    stream.write_all(&response.to_bytes())?;
    stream.flush()
}

/// 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);
    }
}