umbral_socket/stream/

server.rs

use std::io::{self, Read, Result, Write};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::UnixListener as StdUnixListener;
use std::path::Path;
use std::slice;

use mio::net::{UnixListener, UnixStream};
use mio::{Events, Interest, Poll, Token};

use super::protocol::{
    MAX_PAYLOAD_LEN, MethodId, REQUEST_HEADER_LEN, SOCKET_PERMISSIONS, UmbralStatus,
};

const SERVER: Token = Token(0);
const CONNECTION_START: usize = 1;

type Handler<S> = fn(&S, &[u8]) -> io::Result<&'static [u8]>;

pub struct UmbralServer<S> {
    state: S,
    handlers: [Option<Handler<S>>; 256],
}

struct Connection {
    stream: UnixStream,
    state: ConnectionState,
    header: [u8; REQUEST_HEADER_LEN],
    header_offset: usize,
    request: Vec<u8>,
    request_len: usize,
    write_header: [u8; REQUEST_HEADER_LEN],
    write_header_offset: usize,
    write_payload: &'static [u8],
    write_payload_offset: usize,
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum ConnectionState {
    ReadingHeader,
    ReadingBody,
    Writing,
}

impl<S> UmbralServer<S> {
    pub fn new(state: S) -> Self {
        Self {
            state,
            handlers: std::array::from_fn(|_| None),
        }
    }

    pub fn route(mut self, method: MethodId, handler: Handler<S>) -> Self {
        self.handlers[method as usize] = Some(handler);
        self
    }

    pub fn run(self, socket: &str) -> Result<()> {
        let path = Path::new(socket);
        if path.exists() {
            std::fs::remove_file(path)?;
        }

        let std_listener = StdUnixListener::bind(path)?;
        std_listener.set_nonblocking(true)?;
        let permissions = std::fs::Permissions::from_mode(SOCKET_PERMISSIONS);
        std::fs::set_permissions(path, permissions)?;

        let mut listener = UnixListener::from_std(std_listener);
        let mut poll = Poll::new()?;
        poll.registry()
            .register(&mut listener, SERVER, Interest::READABLE)?;

        let mut events = Events::with_capacity(1024);
        let mut connections = Vec::new();
        let mut next_token = CONNECTION_START;

        loop {
            poll.poll(&mut events, None)?;

            for event in events.iter() {
                let token = event.token();
                if token == SERVER {
                    accept_connections(
                        poll.registry(),
                        &mut listener,
                        &mut connections,
                        &mut next_token,
                    )?;
                    continue;
                }

                if event.is_readable() {
                    self.read_connection(poll.registry(), token, &mut connections)?;
                }
                if event.is_writable() {
                    flush_connection(poll.registry(), token, &mut connections)?;
                }
            }
        }
    }

    fn read_connection(
        &self,
        registry: &mio::Registry,
        token: Token,
        connections: &mut Vec<Option<Connection>>,
    ) -> Result<()> {
        loop {
            let Some(connection) = connection_mut(connections, token) else {
                return Ok(());
            };

            if connection.state == ConnectionState::Writing {
                return update_interest(registry, token, connections);
            }

            match read_request(connection, MAX_PAYLOAD_LEN) {
                Ok(Some(method)) => {
                    self.handle_request(connection, method);
                    flush_connection(registry, token, connections)?;
                }
                Ok(None) => return update_interest(registry, token, connections),
                Err(error)
                    if matches!(
                        error.kind(),
                        io::ErrorKind::InvalidData | io::ErrorKind::UnexpectedEof
                    ) =>
                {
                    remove_connection(registry, token, connections)?;
                    return Ok(());
                }
                Err(error) => return Err(error),
            };
        }
    }

    fn handle_request(&self, connection: &mut Connection, method: MethodId) {
        let Some(handler) = self.handlers[method as usize] else {
            prepare_status(connection, UmbralStatus::MethodNotFound);
            return;
        };

        match handler(&self.state, &connection.request) {
            Ok(payload) => {
                if payload.len() > MAX_PAYLOAD_LEN {
                    prepare_status(connection, UmbralStatus::HandlerError);
                    return;
                }

                prepare_response(connection, UmbralStatus::Ok, payload);
            }
            Err(_) => {
                prepare_status(connection, UmbralStatus::HandlerError);
            }
        }
    }
}

impl Connection {
    fn new(stream: UnixStream) -> Self {
        Self {
            stream,
            state: ConnectionState::ReadingHeader,
            header: [0; REQUEST_HEADER_LEN],
            header_offset: 0,
            request: Vec::new(),
            request_len: 0,
            write_header: [0; REQUEST_HEADER_LEN],
            write_header_offset: 0,
            write_payload: b"",
            write_payload_offset: 0,
        }
    }

    fn reset_read(&mut self) {
        self.state = ConnectionState::ReadingHeader;
        self.header = [0; REQUEST_HEADER_LEN];
        self.header_offset = 0;
        self.request.clear();
        self.request_len = 0;
        self.write_header_offset = 0;
        self.write_payload = b"";
        self.write_payload_offset = 0;
    }
}

fn accept_connections(
    registry: &mio::Registry,
    listener: &mut UnixListener,
    connections: &mut Vec<Option<Connection>>,
    next_token: &mut usize,
) -> Result<()> {
    loop {
        let (mut stream, _) = match listener.accept() {
            Ok(accepted) => accepted,
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => return Ok(()),
            Err(error) if error.raw_os_error() == Some(24) => return Ok(()),
            Err(error) => return Err(error),
        };

        let token = Token(*next_token);
        *next_token += 1;
        registry.register(&mut stream, token, Interest::READABLE)?;
        let index = connection_index(token);
        if connections.len() <= index {
            connections.resize_with(index + 1, || None);
        }
        connections[index] = Some(Connection::new(stream));
    }
}

fn read_request(connection: &mut Connection, max_payload_len: usize) -> Result<Option<MethodId>> {
    loop {
        match connection.state {
            ConnectionState::ReadingHeader => {
                while connection.header_offset < REQUEST_HEADER_LEN {
                    match connection
                        .stream
                        .read(&mut connection.header[connection.header_offset..])
                    {
                        Ok(0) => {
                            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "eof"));
                        }
                        Ok(read) => connection.header_offset += read,
                        Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
                            return Ok(None);
                        }
                        Err(error) => return Err(error),
                    }
                }

                connection.request_len = u32::from_be_bytes([
                    connection.header[1],
                    connection.header[2],
                    connection.header[3],
                    connection.header[4],
                ]) as usize;

                if connection.request_len > max_payload_len {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        "payload too large",
                    ));
                }

                connection.request.clear();
                connection.state = ConnectionState::ReadingBody;

                if connection.request_len == 0 {
                    return Ok(Some(connection.header[0]));
                }
            }
            ConnectionState::ReadingBody => {
                while connection.request.len() < connection.request_len {
                    match read_into_spare(
                        &mut connection.stream,
                        &mut connection.request,
                        connection.request_len,
                    ) {
                        Ok(0) => {
                            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "eof"));
                        }
                        Ok(_) => {}
                        Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
                            return Ok(None);
                        }
                        Err(error) => return Err(error),
                    }
                }

                return Ok(Some(connection.header[0]));
            }
            ConnectionState::Writing => return Ok(None),
        }
    }
}

fn read_into_spare(
    stream: &mut UnixStream,
    payload: &mut Vec<u8>,
    target_len: usize,
) -> Result<usize> {
    let remaining = target_len - payload.len();
    payload.reserve(remaining);

    let before = payload.len();
    let buffer = unsafe { slice::from_raw_parts_mut(payload.as_mut_ptr().add(before), remaining) };
    let read = stream.read(buffer)?;

    unsafe {
        payload.set_len(before + read);
    }

    Ok(read)
}

fn prepare_status(connection: &mut Connection, status: UmbralStatus) {
    prepare_response(connection, status, b"");
}

fn prepare_response(connection: &mut Connection, status: UmbralStatus, payload: &'static [u8]) {
    let len = payload.len() as u32;
    connection.write_header[0] = status as u8;
    connection.write_header[1..].copy_from_slice(&len.to_be_bytes());
    connection.write_header_offset = 0;
    connection.write_payload = payload;
    connection.write_payload_offset = 0;
    connection.state = ConnectionState::Writing;
}

fn flush_connection(
    registry: &mio::Registry,
    token: Token,
    connections: &mut Vec<Option<Connection>>,
) -> Result<()> {
    let Some(connection) = connection_mut(connections, token) else {
        return Ok(());
    };

    if flush_write(connection)? {
        connection.reset_read();
    }

    update_interest(registry, token, connections)
}

fn flush_write(connection: &mut Connection) -> Result<bool> {
    while connection.write_header_offset < connection.write_header.len() {
        match connection
            .stream
            .write(&connection.write_header[connection.write_header_offset..])
        {
            Ok(0) => return write_zero(),
            Ok(written) => connection.write_header_offset += written,
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => return Ok(false),
            Err(error) => return Err(error),
        }
    }

    while connection.write_payload_offset < connection.write_payload.len() {
        match connection
            .stream
            .write(&connection.write_payload[connection.write_payload_offset..])
        {
            Ok(0) => return write_zero(),
            Ok(written) => connection.write_payload_offset += written,
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => return Ok(false),
            Err(error) => return Err(error),
        }
    }

    Ok(true)
}

fn write_zero<T>() -> Result<T> {
    Err(io::Error::new(
        io::ErrorKind::WriteZero,
        "failed to write umbral frame",
    ))
}

fn update_interest(
    registry: &mio::Registry,
    token: Token,
    connections: &mut Vec<Option<Connection>>,
) -> Result<()> {
    let Some(connection) = connection_mut(connections, token) else {
        return Ok(());
    };
    let interest = if connection.state == ConnectionState::Writing {
        Interest::WRITABLE
    } else {
        Interest::READABLE
    };
    registry.reregister(&mut connection.stream, token, interest)
}

fn remove_connection(
    registry: &mio::Registry,
    token: Token,
    connections: &mut Vec<Option<Connection>>,
) -> Result<()> {
    let index = connection_index(token);
    if let Some(mut connection) = connections.get_mut(index).and_then(Option::take) {
        registry.deregister(&mut connection.stream)?;
    }
    Ok(())
}

fn connection_index(token: Token) -> usize {
    token.0 - CONNECTION_START
}

fn connection_mut(connections: &mut [Option<Connection>], token: Token) -> Option<&mut Connection> {
    connections.get_mut(connection_index(token))?.as_mut()
}