notescraft/crates/engine/tests/live_stub/mod.rs

//! A tiny HTTP server for the live session tests: it answers the live
//! endpoints of the sync server and streams the events the test pushes.
//!
//! The stub is deliberately small: it records the requests it answered, it
//! holds the text of one room and it can drop the event streams it opened.

#![allow(dead_code)] // every helper is used by one of the test cases

use std::io::{BufRead, BufReader, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::{Arc, Mutex};

/// The largest request body the stub reads, in bytes.
const MAX_BODY: usize = 1024 * 1024;

/// One request the stub answered.
#[derive(Clone, Debug)]
pub struct Call {
    /// The path of the request, without the query.
    pub path: String,
    /// The body of the request.
    pub body: String,
}

/// The state of the stub, shared by its threads.
struct State {
    /// The text of the room.
    text: Mutex<String>,
    /// The revision of the room.
    rev: Mutex<u64>,
    /// The joins that arrived.
    joins: Mutex<Vec<Call>>,
    /// The operations that arrived.
    ops: Mutex<Vec<Call>>,
    /// The status the next operation answers with.
    ops_status: Mutex<Option<u16>>,
    /// The number of event streams that were opened.
    streams: Mutex<usize>,
    /// The senders of the open event streams.
    senders: Mutex<Vec<Sender<String>>>,
}

/// A running stub server.
pub struct Stub {
    /// The base URL of the server.
    pub base: String,
    state: Arc<State>,
}

impl Stub {
    /// Starts a server whose room holds `text`.
    pub fn start(text: &str) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
        let port = listener.local_addr().expect("addr").port();
        let state = Arc::new(State {
            text: Mutex::new(text.to_owned()),
            rev: Mutex::new(0),
            joins: Mutex::new(Vec::new()),
            ops: Mutex::new(Vec::new()),
            ops_status: Mutex::new(None),
            streams: Mutex::new(0),
            senders: Mutex::new(Vec::new()),
        });
        let server = Self { base: format!("http://127.0.0.1:{port}"), state: Arc::clone(&state) };
        std::thread::spawn({
            let state = Arc::clone(&state);
            move || {
                for stream in listener.incoming() {
                    let Ok(stream) = stream else { continue };
                    let state = Arc::clone(&state);
                    std::thread::spawn(move || serve(&state, stream));
                }
            }
        });
        server
    }

    /// The joins the stub answered.
    pub fn joins(&self) -> Vec<Call> {
        self.state.joins.lock().expect("joins").clone()
    }

    /// The operations the stub answered.
    pub fn ops(&self) -> Vec<Call> {
        self.state.ops.lock().expect("ops").clone()
    }

    /// Makes the next operation answer with this status.
    pub fn set_ops_status(&self, status: u16) {
        *self.state.ops_status.lock().expect("status") = Some(status);
    }

    /// Changes the text of the room, as another peer would.
    pub fn set_room_text(&self, text: &str) {
        let mut room = self.state.text.lock().expect("text");
        text.clone_into(&mut room);
    }

    /// The number of event streams that were opened.
    pub fn stream_count(&self) -> usize {
        *self.state.streams.lock().expect("streams")
    }

    /// Sends raw Server-Sent Events text to the open streams.
    pub fn push_event(&self, raw: &str) {
        let mut senders = self.state.senders.lock().expect("senders");
        senders.retain(|sender| sender.send(raw.to_owned()).is_ok());
    }

    /// Drops the open event streams: the clients see the stream end.
    pub fn drop_streams(&self) {
        let mut senders = self.state.senders.lock().expect("senders");
        senders.clear();
    }
}

/// Answers one request of a client.
fn serve(state: &Arc<State>, mut stream: TcpStream) {
    let Some(request) = read_request(&mut stream) else {
        return;
    };
    if request.path.ends_with("/join") {
        state.joins.lock().expect("joins").push(request);
        let reply = format!("{{\"text\":{},\"rev\":{},\"peers\":[]}}", quote(&state.text()), state.rev());
        respond(&mut stream, 200, &reply);
    } else if request.path.ends_with("/ops") {
        state.ops.lock().expect("ops").push(request);
        let status = state.ops_status.lock().expect("status").take();
        if let Some(status) = status {
            respond(&mut stream, status, "{\"error\":\"stale\"}");
        } else {
            let rev = state.bump_rev();
            respond(&mut stream, 200, &format!("{{\"rev\":{rev}}}"));
        }
    } else if request.path.ends_with("/presence") {
        respond(&mut stream, 200, "{\"peers\":[]}");
    } else if request.path.ends_with("/leave") {
        respond(&mut stream, 200, "{}");
    } else if request.path.ends_with("/events") {
        stream_events(state, stream);
    } else {
        respond(&mut stream, 404, "{}");
    }
}

impl State {
    /// The text of the room.
    fn text(&self) -> String {
        self.text.lock().expect("text").clone()
    }

    /// The revision of the room.
    fn rev(&self) -> u64 {
        *self.rev.lock().expect("rev")
    }

    /// Moves the revision of the room on.
    fn bump_rev(&self) -> u64 {
        let mut rev = self.rev.lock().expect("rev");
        *rev = rev.saturating_add(1);
        *rev
    }
}

/// Reads the request of a client: the line, the headers and the body.
fn read_request(stream: &mut TcpStream) -> Option<Call> {
    let mut reader = BufReader::new(stream.try_clone().ok()?);
    let mut line = String::new();
    reader.read_line(&mut line).ok()?;
    let path = line.split_whitespace().nth(1)?.split('?').next()?.to_owned();
    let mut length = 0usize;
    loop {
        let mut header = String::new();
        if reader.read_line(&mut header).ok()? == 0 {
            break;
        }
        let header = header.trim_end();
        if header.is_empty() {
            break;
        }
        if let Some((name, value)) = header.split_once(':')
            && name.eq_ignore_ascii_case("content-length")
        {
            length = value.trim().parse().ok()?;
        }
    }
    let mut body = vec![0u8; length.min(MAX_BODY)];
    if !body.is_empty() {
        reader.read_exact(&mut body).ok()?;
    }
    Some(Call { path, body: String::from_utf8_lossy(&body).into_owned() })
}

/// Writes a complete answer and closes the connection.
fn respond(stream: &mut TcpStream, status: u16, body: &str) {
    let reason = match status {
        200 => "OK",
        204 => "No Content",
        404 => "Not Found",
        409 => "Conflict",
        _ => "Error",
    };
    let _ = write!(stream, "HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len());
    let _ = stream.flush();
}

/// Streams the events of the test until the stream is dropped.
fn stream_events(state: &Arc<State>, mut stream: TcpStream) {
    let (sender, receiver): (Sender<String>, Receiver<String>) = mpsc::channel();
    state.senders.lock().expect("senders").push(sender);
    *state.streams.lock().expect("streams") += 1;
    let _ = write!(
        stream,
        "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n"
    );
    let _ = stream.flush();
    loop {
        match receiver.recv_timeout(std::time::Duration::from_millis(50)) {
            Ok(chunk) => {
                if write!(stream, "{:x}\r\n{chunk}\r\n", chunk.len()).is_err() || stream.flush().is_err() {
                    break;
                }
            }
            Err(mpsc::RecvTimeoutError::Timeout) => {}
            Err(mpsc::RecvTimeoutError::Disconnected) => break,
        }
    }
}

/// The text of the room as a JSON string.
fn quote(text: &str) -> String {
    use std::fmt::Write;
    let mut out = String::with_capacity(text.len() + 2);
    out.push('"');
    for character in text.chars() {
        match character {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            other if (other as u32) < 0x20 => {
                let _ = write!(out, "\\u{:04x}", other as u32);
            }
            other => out.push(other),
        }
    }
    out.push('"');
    out
}