DS4Server/src/server/http.rs

use super::*;

pub(super) fn send_sse_headers(stream: &mut impl Write) -> Result<(), String> {
    stream
        .write_all(
            b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n",
        )
        .map_err(|error| error.to_string())
}

pub(super) fn send_sse(stream: &mut impl Write, value: &Value) -> Result<(), String> {
    let body = serde_json::to_vec(value).map_err(|error| error.to_string())?;
    stream
        .write_all(b"data: ")
        .map_err(|error| error.to_string())?;
    stream.write_all(&body).map_err(|error| error.to_string())?;
    stream.write_all(b"\n\n").map_err(|error| error.to_string())
}

pub(super) fn send_sse_error(stream: &mut impl Write, message: &str) -> Result<(), String> {
    let body = serde_json::to_vec(&json!({"error": {"message": message, "type": "server_error"}}))
        .map_err(|error| error.to_string())?;
    stream
        .write_all(b"event: error\ndata: ")
        .map_err(|error| error.to_string())?;
    stream.write_all(&body).map_err(|error| error.to_string())?;
    stream.write_all(b"\n\n").map_err(|error| error.to_string())
}

pub(super) fn send_json(stream: &mut TcpStream, code: u16, value: &Value) -> Result<(), String> {
    let mut body = serde_json::to_vec(value).map_err(|error| error.to_string())?;
    body.push(b'\n');
    send_response(stream, code, Some("application/json"), &body)
}

pub(super) fn send_error(stream: &mut TcpStream, code: u16, message: &str) -> Result<(), String> {
    send_json(
        stream,
        code,
        &json!({"error": {"message": message, "type": "invalid_request_error"}}),
    )
}

pub(super) fn send_response(
    stream: &mut TcpStream,
    code: u16,
    content_type: Option<&str>,
    body: &[u8],
) -> Result<(), String> {
    let reason = match code {
        200 => "OK",
        204 => "No Content",
        400 => "Bad Request",
        404 => "Not Found",
        409 => "Conflict",
        500 => "Internal Server Error",
        _ => "Error",
    };
    let mut header = format!(
        "HTTP/1.1 {code} {reason}\r\nContent-Length: {}\r\n",
        body.len()
    );
    if let Some(content_type) = content_type {
        header.push_str("Content-Type: ");
        header.push_str(content_type);
        header.push_str("\r\n");
    }
    header.push_str("Connection: close\r\n\r\n");
    stream
        .write_all(header.as_bytes())
        .and_then(|()| stream.write_all(body))
        .map_err(|error| error.to_string())
}

pub(super) fn read_request(stream: &mut TcpStream) -> Result<HttpRequest, String> {
    let mut bytes = Vec::new();
    let header_end = loop {
        if bytes.len() >= MAX_HEADER_BYTES {
            return Err("HTTP headers are too large".into());
        }
        let mut chunk = [0_u8; 4096];
        let read = stream.read(&mut chunk).map_err(|error| error.to_string())?;
        if read == 0 {
            return Err("bad HTTP request".into());
        }
        bytes.extend_from_slice(&chunk[..read]);
        if let Some(end) = find_header_end(&bytes) {
            break end;
        }
    };
    let header = std::str::from_utf8(&bytes[..header_end])
        .map_err(|_| "HTTP headers are not UTF-8".to_owned())?;
    let mut lines = header.lines();
    let request_line = lines.next().ok_or_else(|| "bad HTTP request".to_owned())?;
    let mut parts = request_line.split_whitespace();
    let method = parts
        .next()
        .ok_or_else(|| "bad HTTP request".to_owned())?
        .to_owned();
    let path = parts
        .next()
        .ok_or_else(|| "bad HTTP request".to_owned())?
        .to_owned();
    let length = lines
        .find_map(|line| {
            line.split_once(':').and_then(|(name, value)| {
                name.eq_ignore_ascii_case("content-length")
                    .then(|| value.trim().parse::<usize>().ok())
                    .flatten()
            })
        })
        .unwrap_or(0);
    if length > MAX_BODY_BYTES {
        return Err("HTTP body is too large".into());
    }
    while bytes.len() < header_end + length {
        let mut chunk = [0_u8; 8192];
        let read = stream.read(&mut chunk).map_err(|error| error.to_string())?;
        if read == 0 {
            return Err("incomplete HTTP body".into());
        }
        bytes.extend_from_slice(&chunk[..read]);
    }
    Ok(HttpRequest {
        method,
        path: path.split('?').next().unwrap_or(&path).to_owned(),
        body: bytes[header_end..header_end + length].to_vec(),
    })
}

fn find_header_end(bytes: &[u8]) -> Option<usize> {
    bytes
        .windows(4)
        .position(|window| window == b"\r\n\r\n")
        .map(|position| position + 4)
        .or_else(|| {
            bytes
                .windows(2)
                .position(|window| window == b"\n\n")
                .map(|position| position + 2)
        })
}

pub(super) fn unix_time() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

pub(super) fn random_id(prefix: &str) -> String {
    random_hex_id::<12>(prefix)
}

pub(super) fn random_tool_id(prefix: &str) -> String {
    random_hex_id::<16>(prefix)
}

fn random_hex_id<const N: usize>(prefix: &str) -> String {
    let mut bytes = [0_u8; N];
    if File::open("/dev/urandom")
        .and_then(|mut file| file.read_exact(&mut bytes))
        .is_err()
    {
        let fallback = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos();
        bytes.copy_from_slice(&fallback.to_le_bytes()[..N]);
    }
    let mut id = String::with_capacity(prefix.len() + N * 2);
    id.push_str(prefix);
    for byte in bytes {
        use std::fmt::Write as _;
        let _ = write!(id, "{byte:02x}");
    }
    id
}