ehttpd 0.14.0

//! Implements a simple threadpool-based server
#![cfg(feature = "server")]

mod pool;
mod worker;

use crate::bytes::{Sink, Source};
use crate::error::Error;
use crate::http::{Request, Response};
use crate::server::pool::{Executable, Threadpool};
use std::convert::Infallible;
use std::io::{BufReader, BufWriter, Write};
use std::net::{TcpListener, TcpStream, ToSocketAddrs};
use std::sync::Arc;
use std::time::Duration;

/// A connection handler
#[derive(Clone)]
enum Handler {
    /// A `source,sink`-handler
    #[allow(clippy::type_complexity, reason = "the type is trivially human readable")]
    SourceSink(Arc<dyn Fn(&mut Source, &mut Sink) -> bool + Send + Sync + 'static>),
    /// A `request->response`-handler
    RequestResponse(Arc<dyn Fn(Request) -> Response + Send + Sync + 'static>),
}
impl Handler {
    /// Handles a given connection and returns whether the handler wants to be rescheduled (e.g. keep-alive)
    pub fn exec(&self, source: &mut Source, sink: &mut Sink) -> bool {
        match self {
            Handler::SourceSink(handler) => handler(source, sink),
            Handler::RequestResponse(handler) => Self::bridge_request_response(source, sink, handler.as_ref()),
        }
    }

    /// Bridges a `request->response`-handler to a source-sink pattern
    fn bridge_request_response<F>(source: &mut Source, sink: &mut Sink, handler: &F) -> bool
    where
        F: Fn(Request) -> Response + ?Sized,
    {
        // Read request
        let Ok(Some(request)) = Request::from_stream(source) else {
            return false;
        };

        // Copy header info we need later on
        let is_head = request.method.eq_ignore_ascii_case(b"HEAD");
        let is_http10 = request.version.eq_ignore_ascii_case(b"HTTP/1.0");
        let has_connection_close = request.has_connection_close();
        let has_transfer_encoding = request.field("Transfer-Encoding").is_some();
        let content_length = request.content_length();
        let body_start = request.stream.bytes_read();

        // Handle the request accordingly
        let mut response = handler(request);
        if is_head {
            // Drop body for HEAD requests
            response.make_head();
        }

        // Write response
        let Ok(_) = response.to_stream(sink) else {
            return false;
        };

        // See if the request body has been fully consumed
        #[allow(clippy::arithmetic_side_effects, reason = "should never underflow")]
        let body_read_len = source.bytes_read() - body_start;
        let has_body_read_err = match content_length {
            Ok(Some(expected)) => body_read_len != expected,
            Ok(None) => body_read_len != 0,
            Err(_) => true,
        };

        // Reschedule the connection if it is not closed **and** we can be sure it has been fully read
        let do_connection_close = is_http10
            || has_connection_close
            || has_transfer_encoding
            || has_body_read_err
            || response.has_connection_close();
        !do_connection_close
    }
}

/// An encapsulated connection to pass to the thread pool
struct Connection<const STACK_SIZE: usize> {
    /// The connection handler
    pub handler: Handler,
    /// The receiving half of the stream
    pub source: Source,
    /// The writing half of the stream
    pub sink: Sink,
    /// The thread-pool so that keep-alive connections can requeue themselves
    pub threadpool: Threadpool<Self, STACK_SIZE>,
}
impl<const STACK_SIZE: usize> Executable for Connection<STACK_SIZE> {
    fn exec(mut self) {
        // Call the connection handler and flush the output
        let reschedule = self.handler.exec(&mut self.source, &mut self.sink);
        if self.sink.flush().is_ok() && reschedule {
            // Reschedule the connection
            let threadpool = self.threadpool.clone();
            let _ = threadpool.dispatch(self);
        }
    }
}

/// A threadpool-based HTTP server
pub struct Server<const STACK_SIZE: usize> {
    /// The thread pool to handle the incoming connections
    threadpool: Threadpool<Connection<STACK_SIZE>, STACK_SIZE>,
    /// The connection handler
    handler: Handler,
}
impl<const STACK_SIZE: usize> Server<STACK_SIZE> {
    /// Define a strict read timeout to close idle connections on time
    const CONNECTION_TIMEOUT: Duration = Duration::from_secs(20);

    /// Creates a new server with the given connection handler
    pub fn with_source_sink<F>(workers_max: usize, source_sink_handler: F) -> Self
    where
        F: Fn(&mut Source, &mut Sink) -> bool + Send + Sync + 'static,
    {
        // Create threadpool and init self
        let threadpool: Threadpool<_, STACK_SIZE> = Threadpool::new(workers_max);
        let handler = Handler::SourceSink(Arc::new(source_sink_handler));
        Self { threadpool, handler }
    }
    /// Creates a new server with the given connection handler
    pub fn with_request_response<F>(workers_max: usize, request_response_handler: F) -> Self
    where
        F: Fn(Request) -> Response + Send + Sync + 'static,
    {
        // Create threadpool and init self
        let threadpool: Threadpool<_, STACK_SIZE> = Threadpool::new(workers_max);
        let handler = Handler::RequestResponse(Arc::new(request_response_handler));
        Self { threadpool, handler }
    }

    /// Manually dispatches a connection
    pub fn dispatch(&self, source: Source, sink: Sink) -> Result<(), Error> {
        // Create and dispatch the job
        self.threadpool.dispatch(Connection {
            handler: self.handler.clone(),
            source,
            sink,
            threadpool: self.threadpool.clone(),
        })
    }

    /// Listens on the given address and accepts forever
    pub fn accept<A>(self, address: A) -> Result<Infallible, Error>
    where
        A: ToSocketAddrs,
    {
        // Bind and listen
        let socket = TcpListener::bind(address)?;
        'try_accept: loop {
            // Configure the socket or immediately drop the connection on error
            let Ok((source, sink)) = self.accept_connection(&socket) else {
                continue 'try_accept;
            };

            // Best-effort attempt to dispatch the connection; drop it otherwise
            let sink = BufWriter::new(sink);
            let source = BufReader::new(source);
            let _ = self.dispatch(source.into(), sink.into());
        }
    }

    /// Accepts a connection and prepares it for dispatch
    fn accept_connection(&self, socket: &TcpListener) -> Result<(TcpStream, TcpStream), Error> {
        // Accept the connection and setup timeouts
        let (connection, _) = socket.accept()?;
        connection.set_read_timeout(Some(Self::CONNECTION_TIMEOUT))?;
        connection.set_write_timeout(Some(Self::CONNECTION_TIMEOUT))?;

        // Duplicate the connection for independent source/sink handling
        let sink = connection.try_clone()?;
        Ok((connection, sink))
    }
}