haematite-cli 0.21.7

use std::io::{self, BufReader, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::mpsc;
use std::thread::JoinHandle;
use std::time::Duration;

use ureq::http::HeaderMap;

use super::{Config, listener};

pub(super) type TestResult = Result<(), Box<dyn std::error::Error + Send + Sync>>;
type TestTask = JoinHandle<TestResult>;

pub(super) struct Server {
    pub(super) address: SocketAddr,
    stop: Option<mpsc::Sender<()>>,
    task: Option<TestTask>,
}

impl Server {
    pub(super) fn mcp(
        config: Config,
        service: std::path::PathBuf,
        cap: usize,
        clients: usize,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let server = listener::Server::new(config, service, Duration::from_secs(5), cap, clients)?;
        Self::from_mcp(server)
    }

    pub(super) fn mcp_signalled(
        config: Config,
        service: std::path::PathBuf,
        cap: usize,
        clients: usize,
        admitted: mpsc::Sender<()>,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let server = listener::Server::new(config, service, Duration::from_secs(5), cap, clients)?
            .with_admission_signal(admitted)?;
        Self::from_mcp(server)
    }

    pub(super) fn from_mcp(
        server: listener::Server,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let socket = TcpListener::bind("127.0.0.1:0")?;
        let address = socket.local_addr()?;
        let (stop, stopped) = mpsc::channel();
        let task = std::thread::spawn(move || {
            server.until(socket, stopped);
            Ok(())
        });
        Ok(Self {
            address,
            stop: Some(stop),
            task: Some(task),
        })
    }

    pub(super) fn finish(mut self) -> TestResult {
        self.close()
    }

    fn close(&mut self) -> TestResult {
        if let Some(stop) = self.stop.take() {
            let signal = stop.send(());
            let wake = TcpStream::connect(self.address);
            let joined = self
                .task
                .take()
                .ok_or_else(|| io::Error::other("fixture task missing"))?
                .join()
                .map_err(|_| io::Error::other("fixture listener panicked"))?;
            joined?;
            signal
                .map_err(|_| io::Error::other("fixture listener ended before its finish signal"))?;
            drop(wake?);
        }
        Ok(())
    }
}

impl Drop for Server {
    fn drop(&mut self) {
        if let Err(error) = self.close() {
            eprintln!("http_fixture_cleanup_failed: {error}");
        }
    }
}

/// The HTTP MCP's settings every fixture starts from: loopback, one
/// permitted origin, and short header and body deadlines.
pub(super) fn config() -> Result<Config, serde_json::Error> {
    serde_json::from_value(serde_json::json!({
        "listen":"127.0.0.1:0",
        "allowed_origins":["https://agent.example"],
        "max_connections":8,
        "header_timeout_ms":1000,
        "body_timeout_ms":1000
    }))
}

pub(super) struct Reply {
    pub(super) status: u16,
    pub(super) headers: HeaderMap,
    pub(super) body: Vec<u8>,
}

impl Reply {
    pub(super) fn read(
        mut socket: TcpStream,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        Self::read_open(&mut socket)
    }

    pub(super) fn read_open(
        socket: &mut TcpStream,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let mut bytes = Vec::new();
        socket.read_to_end(&mut bytes)?;
        Self::parse(&bytes)
    }

    pub(super) fn read_frame(
        socket: TcpStream,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let mut socket = BufReader::new(socket);
        let mut bytes = Vec::new();
        let mut byte = [0];
        while !bytes.ends_with(b"\r\n\r\n") {
            if bytes.len() == 16_384 {
                return Err(io::Error::other("fixture response header exceeds its bound").into());
            }
            socket.read_exact(&mut byte)?;
            bytes.push(byte[0]);
        }
        let text = std::str::from_utf8(&bytes)?;
        let lengths = text
            .split("\r\n")
            .filter_map(|line| line.split_once(':'))
            .filter(|(name, _)| name.eq_ignore_ascii_case("content-length"))
            .collect::<Vec<_>>();
        if lengths.len() != 1 {
            return Err(io::Error::other("fixture response length must be unique").into());
        }
        let length: usize = lengths[0].1.trim().parse()?;
        if length > 131_072 {
            return Err(io::Error::other("fixture response body exceeds its bound").into());
        }
        let start = bytes.len();
        bytes.resize(start + length, 0);
        socket.read_exact(&mut bytes[start..])?;
        Self::parse(&bytes)
    }

    fn parse(bytes: &[u8]) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let split = bytes
            .windows(4)
            .position(|part| part == b"\r\n\r\n")
            .ok_or_else(|| io::Error::other("fixture response header missing"))?;
        let header = std::str::from_utf8(&bytes[..split])?;
        let mut lines = header.split("\r\n");
        let status = lines
            .next()
            .and_then(|line| line.split_ascii_whitespace().nth(1))
            .ok_or_else(|| io::Error::other("fixture response status missing"))?
            .parse()?;
        let mut headers = HeaderMap::new();
        for line in lines {
            let (name, value) = line
                .split_once(':')
                .ok_or_else(|| io::Error::other("fixture response header invalid"))?;
            headers.append(
                ureq::http::HeaderName::from_bytes(name.as_bytes())?,
                ureq::http::HeaderValue::from_str(value.trim())?,
            );
        }
        let body = bytes[split + 4..].to_vec();
        let expected: usize = headers
            .get("content-length")
            .ok_or_else(|| io::Error::other("fixture response length missing"))?
            .to_str()?
            .parse()?;
        assert_eq!(body.len(), expected);
        Ok(Self {
            status,
            headers,
            body,
        })
    }
}

pub(super) fn request(
    server: &Server,
    method: &str,
    body: &str,
    token: bool,
    extra: &[(&str, &str)],
) -> Result<TcpStream, io::Error> {
    let mut headers = vec![
        ("Content-Type", "application/json"),
        ("Accept", "application/json, text/event-stream"),
    ];
    if token {
        headers.push(("Authorization", "Bearer fixture-token"));
    }
    headers.extend_from_slice(extra);
    request_with_headers(server, method, body, &headers)
}

pub(super) fn request_with_headers(
    server: &Server,
    method: &str,
    body: &str,
    headers: &[(&str, &str)],
) -> Result<TcpStream, io::Error> {
    let mut stream = TcpStream::connect(server.address)?;
    let mut head = format!(
        "{method} /mcp HTTP/1.1\r\nHost: {}\r\nContent-Length: {}\r\n",
        server.address,
        body.len()
    );
    for (name, value) in headers {
        head.push_str(name);
        head.push_str(": ");
        head.push_str(value);
        head.push_str("\r\n");
    }
    head.push_str("\r\n");
    stream.write_all(head.as_bytes())?;
    stream.write_all(body.as_bytes())?;
    Ok(stream)
}

pub(super) fn refusal(response: &Reply, status: u16, code: &str) -> TestResult {
    assert_eq!(response.status, status);
    assert_eq!(
        response
            .headers
            .get("cache-control")
            .and_then(|value| value.to_str().ok()),
        Some("no-store")
    );
    let body: serde_json::Value = serde_json::from_slice(&response.body)?;
    assert_eq!(
        body.get("code").and_then(serde_json::Value::as_str),
        Some(code)
    );
    assert!(!std::str::from_utf8(&response.body)?.contains("fixture-token"));
    Ok(())
}

/// The caller's bearer pass: its own bytes, which the store must see unchanged.
pub(super) const PASS: &str = "fixture-token";

pub(super) const LIST: &str = r#"{"jsonrpc":"2.0","id":"list","method":"tools/list"}"#;