mise-server 0.1.11

use crate::{
    mise::{Request, Response},
    routes::{ProcessorResponse, RouteMethod},
};
use http_by_chunks::{ConverterTo, HttpByChunks};
use mio::Registry;
use serde_json::Value;
use std::{
    collections::VecDeque,
    io::{ErrorKind, Read, Write},
    net::{Shutdown, SocketAddr},
};
use tracing::{error, warn};

pub(crate) struct Connection {
    socket: mio::net::TcpStream,
    addr: SocketAddr,
    closed: bool,
    // the read buffer but dynamically growing
    read_buffer: Vec<u8>,
    // partial filling of fixed sized 4k
    // essentially the read buffer fixed block
    static_read_buf: [u8; 4096],
    // The current request being built
    cur_request: HttpByChunks<'static, Value>,
    queued_requests: VecDeque<Request>,
    // All the remaining write buffer
    write_buffer: Vec<u8>,
}

struct JsonConverter;
impl ConverterTo<Value> for JsonConverter {
    fn convert(&self, buf: &[u8]) -> Value {
        serde_json::from_slice(buf).unwrap_or(Value::Null)
    }
}

static JSON_CONVERTER: JsonConverter = JsonConverter;

impl From<(mio::net::TcpStream, SocketAddr)> for Connection {
    fn from(value: (mio::net::TcpStream, SocketAddr)) -> Self {
        Self {
            socket: value.0,
            addr: value.1,
            closed: false,
            read_buffer: vec![],
            static_read_buf: [0; 4096],
            cur_request: HttpByChunks::<Value>::new(&JSON_CONVERTER),
            queued_requests: VecDeque::new(),
            write_buffer: vec![],
        }
    }
}

impl Connection {
    pub(crate) fn write_buffer(&mut self) -> bool {
        if self.write_buffer.is_empty() {
            return false;
        }
        let count = match self.socket.write(&self.write_buffer) {
            Ok(0) => {
                return false;
            }
            Ok(x) => x,
            Err(e) if e.kind() == ErrorKind::WouldBlock => {
                return false;
            }
            Err(e) if e.kind() == ErrorKind::Interrupted => {
                return false;
            }
            Err(e) if is_connection_closed(e.kind()) => {
                let _ = self.socket.shutdown(Shutdown::Both);
                self.closed = true;
                return false;
            }
            Err(e) => {
                error!("{e}");
                return false;
            }
        };
        self.write_buffer = self.write_buffer[count..].to_vec();
        true
    }

    pub(crate) fn read_buffer(&mut self) -> bool {
        let count = match self.socket.read(&mut self.static_read_buf) {
            Ok(0) => {
                return false;
            }
            Ok(x) => x,
            Err(e) if e.kind() == ErrorKind::WouldBlock => {
                return false;
            }
            Err(e) if e.kind() == ErrorKind::Interrupted => {
                return false;
            }
            Err(e) if is_connection_closed(e.kind()) => {
                self.closed = true;
                return false;
            }
            Err(e) => {
                error!("{e}");
                return false;
            }
        };
        self.read_buffer
            .extend_from_slice(&self.static_read_buf[..count]);
        self.append_buffer_and_push_in_queue();
        true
    }

    fn append_buffer_and_push_in_queue(&mut self) {
        if self.read_buffer.is_empty() {
            return;
        }
        self.append_buffer_once();
        while !self.closed
            && let Some(res) = self.cur_request.build_request()
        {
            match res {
                Ok(request) => {
                    self.queued_requests.push_back(Request(request));
                }
                Err(e) => warn!("Unparseable request ({e}), skipping"),
            }
            self.cur_request = HttpByChunks::<Value>::new(&JSON_CONVERTER);

            // Only check this after attempting to build a request, giving it
            // a chance to queue a complete request if the remainder ends up
            // exactly at size 0
            if self.read_buffer.is_empty() {
                return;
            }

            // There may be another one queued in the buffer so reappend more
            // before checking again.
            // This is important otherwise a buffered multi request in the
            // same channel will end up blocking, as mio will sleep while
            // self.read_buffer is still containing data.
            self.append_buffer_once();
        }
    }

    fn append_buffer_once(&mut self) {
        self.read_buffer = self.cur_request.append(&self.read_buffer).to_vec();
        if self.is_pipe_broken() {
            warn!(
                "The request is broken beyond repair (lost sequentiality), closing connection ({}).",
                self.addr
            );
            let _ = self.socket.shutdown(Shutdown::Both);
            self.closed = true;
        }
    }

    fn is_pipe_broken(&self) -> bool {
        let Some(m) = self.cur_request.request_method() else {
            return false;
        };
        let Some(u) = self.cur_request.request_uri() else {
            return false;
        };

        if RouteMethod::try_from(m).is_err() {
            warn!("Invalid/unsupported method: {m}");
            return true;
        }

        // checking the uri is valid here by using the request validator
        if let Err(e) = http::Request::builder().method(m).uri(u).body(()) {
            warn!("Invalid request built: {e}");
            return true;
        }

        false
    }

    // will keep returning None as long as the buffer does not contain a
    // complete request.
    // There may be queued up multiple requests at this point, so keep calling
    // this method until it returns None.
    pub(crate) fn pop_request(&mut self) -> Option<Request> {
        self.queued_requests.pop_front()
    }

    pub(crate) fn send_response(&mut self, resp: ProcessorResponse) {
        match resp {
            ProcessorResponse::Json(resp) => self.send_response_json(&resp),
            ProcessorResponse::Text(txt) => {
                self.write_buffer.extend_from_slice(
                    format!(
                        "HTTP/1.1 200 OK\r\nContent-length:{}\r\n\r\n{}",
                        txt.len(),
                        txt
                    )
                    .as_bytes(),
                );
            }
        }
    }

    fn send_response_json(&mut self, resp: &Response) {
        let mut output: Vec<u8> = vec![];
        output.extend_from_slice(
            format!(
                "HTTP/1.1 {} {}\r\n",
                resp.0.status().as_str(),
                resp.0.status().canonical_reason().unwrap_or("")
            )
            .as_bytes(),
        );
        for h in resp.0.headers() {
            output.extend_from_slice(
                format!(
                    "{}: {}\r\n",
                    h.0,
                    h.1.to_str().expect("Could not translate header value")
                )
                .as_bytes(),
            );
        }
        let body = resp.0.body();
        let bin = serde_json::to_vec(body).expect("Could not serialize json");
        output.extend_from_slice(format!("Content-Length: {}\r\n\r\n", bin.len()).as_bytes());
        output.extend_from_slice(&bin);
        self.write_buffer.extend_from_slice(&output);
    }

    pub(crate) fn is_closed(&self) -> bool {
        self.closed
    }

    pub(crate) fn unregister(&mut self, reg: &Registry) {
        let _ = reg.deregister(&mut self.socket);
    }
}

fn is_connection_closed(kind: ErrorKind) -> bool {
    kind == ErrorKind::BrokenPipe
        || kind == ErrorKind::ConnectionRefused
        || kind == ErrorKind::ConnectionReset
        || kind == ErrorKind::HostUnreachable
        || kind == ErrorKind::NetworkUnreachable
        || kind == ErrorKind::ConnectionAborted
        || kind == ErrorKind::NotConnected
        || kind == ErrorKind::NetworkDown
        || kind == ErrorKind::BrokenPipe
        || kind == ErrorKind::UnexpectedEof
}