rightkit-http 0.3.8

//! In-process mock HTTP/1.1 server for end-to-end tests. Binds 127.0.0.1, records
//! every request exactly as received, and answers from a handler closure.
//! Never contacts the network.

use std::io::{BufRead, BufReader, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;

use crate::{AddressPolicy, ClientConfig};

/// Address policy that lets the [`MockServer`] be reached: the hosts `127.0.0.1`,
/// `localhost` and `[::1]`, loopback addresses only, through `AddressPolicy::allow_if`.
/// Every other destination keeps the default public-web rules.
pub fn loopback_policy() -> AddressPolicy {
    AddressPolicy::default().allow_if(
        vec!["127.0.0.1".into(), "localhost".into(), "::1".into()],
        |ip| ip.is_loopback(),
    )
}

/// `ClientConfig::default()` with [`loopback_policy`]; the starting point for tests that
/// talk to a [`MockServer`].
pub fn loopback_config() -> ClientConfig {
    ClientConfig::default().address_policy(loopback_policy())
}

/// Make clients built afterwards with [`crate::TrustRoots::Bundled`] trust `der` (an X.509
/// root certificate) as well. Process-wide and for tests only: it exists so a test can run a
/// local TLS server under a throwaway CA. The async client trusts it in addition to the
/// Mozilla roots; the blocking client, whose bundled set ureq cannot extend, trusts the
/// registered roots instead of them. [`crate::TrustRoots::Platform`] ignores it.
pub fn add_bundled_test_root(der: Vec<u8>) {
    crate::trust::add_test_root(der);
}

#[derive(Debug, Clone)]
pub struct RecordedRequest {
    pub method: String,
    /// Path plus query, as sent.
    pub target: String,
    /// Lower-cased names.
    pub headers: Vec<(String, String)>,
    pub body: Vec<u8>,
}

impl RecordedRequest {
    pub fn header(&self, name: &str) -> Option<&str> {
        self.headers
            .iter()
            .find(|(k, _)| k.eq_ignore_ascii_case(name))
            .map(|(_, v)| v.as_str())
    }
    pub fn path(&self) -> &str {
        self.target.split('?').next().unwrap_or("")
    }
    pub fn query(&self) -> &str {
        self.target.split_once('?').map(|(_, q)| q).unwrap_or("")
    }
    pub fn body_text(&self) -> String {
        String::from_utf8_lossy(&self.body).into_owned()
    }
    pub fn body_json(&self) -> serde_json::Value {
        serde_json::from_slice(&self.body).expect("request body is JSON")
    }
    /// Decode an urlencoded body or query string into pairs.
    pub fn form(&self) -> Vec<(String, String)> {
        decode_pairs(&self.body_text())
    }
    pub fn query_pairs(&self) -> Vec<(String, String)> {
        decode_pairs(self.query())
    }
}

fn decode_pairs(s: &str) -> Vec<(String, String)> {
    s.split('&')
        .filter(|p| !p.is_empty())
        .map(|p| {
            let (k, v) = p.split_once('=').unwrap_or((p, ""));
            (crate::percent_decode(k), crate::percent_decode(v))
        })
        .collect()
}

pub enum MockBody {
    Full(Vec<u8>),
    /// Sent with chunked transfer-encoding; sleeps `delay` before each chunk.
    Chunks(Vec<(Duration, Vec<u8>)>),
    /// Headers are sent announcing a body 1000 bytes longer than the bytes given (so
    /// `Content-Length` is `len + 1000`), then those bytes are sent and the connection is cut.
    TruncatedAfter(Vec<u8>),
}

pub struct MockResponse {
    pub status: u16,
    pub headers: Vec<(String, String)>,
    pub body: MockBody,
    /// Sleep before writing anything (to exercise timeouts).
    pub delay: Duration,
}

impl MockResponse {
    pub fn new(status: u16) -> Self {
        Self {
            status,
            headers: Vec::new(),
            body: MockBody::Full(Vec::new()),
            delay: Duration::ZERO,
        }
    }
    pub fn json(status: u16, value: &serde_json::Value) -> Self {
        Self::new(status)
            .header("Content-Type", "application/json")
            .body(serde_json::to_vec(value).unwrap())
    }
    pub fn text(status: u16, text: &str) -> Self {
        Self::new(status)
            .header("Content-Type", "text/plain")
            .body(text.as_bytes().to_vec())
    }
    pub fn sse(frames: &[&str]) -> Self {
        Self::new(200)
            .header("Content-Type", "text/event-stream")
            .chunks(
                frames
                    .iter()
                    .map(|f| (Duration::ZERO, f.as_bytes().to_vec()))
                    .collect(),
            )
    }
    pub fn header(mut self, k: &str, v: &str) -> Self {
        self.headers.push((k.into(), v.into()));
        self
    }
    pub fn body(mut self, b: Vec<u8>) -> Self {
        self.body = MockBody::Full(b);
        self
    }
    pub fn chunks(mut self, c: Vec<(Duration, Vec<u8>)>) -> Self {
        self.body = MockBody::Chunks(c);
        self
    }
    pub fn delay(mut self, d: Duration) -> Self {
        self.delay = d;
        self
    }
}

/// A response that writes exactly `body` while announcing a caller-supplied `Content-Length`.
/// Served through [`MockServer::start_declared_length`]. Use it for a precise interrupted
/// download (declare the full length, send a prefix) or a HEAD response (declare the length,
/// send no body). The server generates the one `Content-Length` header; do not add another.
pub struct DeclaredLengthResponse {
    status: u16,
    headers: Vec<(String, String)>,
    body: Vec<u8>,
    content_length: usize,
}

impl DeclaredLengthResponse {
    pub fn new(status: u16, body: Vec<u8>, content_length: usize) -> Self {
        Self {
            status,
            headers: Vec::new(),
            body,
            content_length,
        }
    }
    pub fn header(mut self, k: &str, v: &str) -> Self {
        self.headers.push((k.into(), v.into()));
        self
    }
}

/// What a handler produced, before it is written to the wire.
enum Outgoing {
    Mock(MockResponse),
    Declared(DeclaredLengthResponse),
}

type Handler = dyn Fn(&RecordedRequest, usize) -> Outgoing + Send + Sync;

pub struct MockServer {
    addr: std::net::SocketAddr,
    requests: Arc<Mutex<Vec<RecordedRequest>>>,
    stop: Arc<AtomicBool>,
    thread: Option<JoinHandle<()>>,
}

impl MockServer {
    /// `handler(request, zero_based_request_index)` produces each response.
    pub fn start(
        handler: impl Fn(&RecordedRequest, usize) -> MockResponse + Send + Sync + 'static,
    ) -> Self {
        Self::spawn(Arc::new(move |request: &RecordedRequest, index: usize| {
            Outgoing::Mock(handler(request, index))
        }))
    }

    /// Like [`MockServer::start`], but each response is a [`DeclaredLengthResponse`].
    pub fn start_declared_length(
        handler: impl Fn(&RecordedRequest, usize) -> DeclaredLengthResponse + Send + Sync + 'static,
    ) -> Self {
        Self::spawn(Arc::new(move |request: &RecordedRequest, index: usize| {
            Outgoing::Declared(handler(request, index))
        }))
    }

    fn spawn(handler: Arc<Handler>) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock server");
        let addr = listener.local_addr().unwrap();
        let requests = Arc::new(Mutex::new(Vec::new()));
        let stop = Arc::new(AtomicBool::new(false));
        let (reqs, stop2) = (requests.clone(), stop.clone());
        let thread = thread::spawn(move || {
            for conn in listener.incoming() {
                if stop2.load(Ordering::SeqCst) {
                    break;
                }
                let Ok(stream) = conn else { continue };
                let (handler, reqs) = (handler.clone(), reqs.clone());
                thread::spawn(move || {
                    let _ = serve(stream, &*handler, &reqs);
                });
            }
        });
        Self {
            addr,
            requests,
            stop,
            thread: Some(thread),
        }
    }

    pub fn url(&self, path: &str) -> String {
        format!("http://{}{}", self.addr, path)
    }
    pub fn base_url(&self) -> String {
        format!("http://{}", self.addr)
    }
    pub fn port(&self) -> u16 {
        self.addr.port()
    }
    pub fn requests(&self) -> Vec<RecordedRequest> {
        self.requests.lock().unwrap().clone()
    }
    pub fn request_count(&self) -> usize {
        self.requests.lock().unwrap().len()
    }
}

impl Drop for MockServer {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::SeqCst);
        let _ = TcpStream::connect(self.addr);
        if let Some(t) = self.thread.take() {
            let _ = t.join();
        }
    }
}

/// A response body as it is written to the connection.
enum Wire {
    Mock(MockBody),
    Declared(Vec<u8>, usize),
}

fn serve(
    stream: TcpStream,
    handler: &Handler,
    reqs: &Mutex<Vec<RecordedRequest>>,
) -> std::io::Result<()> {
    stream.set_read_timeout(Some(Duration::from_secs(10)))?;
    let mut reader = BufReader::new(stream.try_clone()?);
    let mut line = String::new();
    if reader.read_line(&mut line)? == 0 {
        return Ok(());
    }
    let mut parts = line.split_whitespace();
    let method = parts.next().unwrap_or("").to_string();
    let target = parts.next().unwrap_or("").to_string();
    let mut headers = Vec::new();
    loop {
        let mut h = String::new();
        reader.read_line(&mut h)?;
        let h = h.trim_end();
        if h.is_empty() {
            break;
        }
        if let Some((k, v)) = h.split_once(':') {
            headers.push((k.trim().to_ascii_lowercase(), v.trim().to_string()));
        }
    }
    let len = headers
        .iter()
        .find(|(k, _)| k == "content-length")
        .and_then(|(_, v)| v.parse::<usize>().ok())
        .unwrap_or(0);
    let mut body = vec![0u8; len];
    reader.read_exact(&mut body)?;
    let request = RecordedRequest {
        method,
        target,
        headers,
        body,
    };
    let index = {
        let mut g = reqs.lock().unwrap();
        g.push(request.clone());
        g.len() - 1
    };
    let (status, headers, delay, wire) = match handler(&request, index) {
        Outgoing::Mock(r) => (r.status, r.headers, r.delay, Wire::Mock(r.body)),
        Outgoing::Declared(d) => (
            d.status,
            d.headers,
            Duration::ZERO,
            Wire::Declared(d.body, d.content_length),
        ),
    };
    if !delay.is_zero() {
        thread::sleep(delay);
    }
    let mut out = stream;
    let reason = match status {
        200 => "OK",
        400 => "Bad Request",
        401 => "Unauthorized",
        429 => "Too Many Requests",
        500 => "Internal Server Error",
        503 => "Service Unavailable",
        _ => "Status",
    };
    let mut head = format!("HTTP/1.1 {status} {reason}\r\nConnection: close\r\n");
    for (k, v) in &headers {
        head.push_str(&format!("{k}: {v}\r\n"));
    }
    match wire {
        Wire::Mock(MockBody::Full(b)) => {
            head.push_str(&format!("Content-Length: {}\r\n\r\n", b.len()));
            out.write_all(head.as_bytes())?;
            out.write_all(&b)?;
        }
        Wire::Mock(MockBody::Chunks(chunks)) => {
            head.push_str("Transfer-Encoding: chunked\r\n\r\n");
            out.write_all(head.as_bytes())?;
            out.flush()?;
            for (delay, c) in chunks {
                if !delay.is_zero() {
                    thread::sleep(delay);
                }
                out.write_all(format!("{:x}\r\n", c.len()).as_bytes())?;
                out.write_all(&c)?;
                out.write_all(b"\r\n")?;
                out.flush()?;
            }
            out.write_all(b"0\r\n\r\n")?;
        }
        Wire::Mock(MockBody::TruncatedAfter(b)) => {
            head.push_str(&format!("Content-Length: {}\r\n\r\n", b.len() + 1000));
            out.write_all(head.as_bytes())?;
            out.write_all(&b)?;
        }
        Wire::Declared(body, content_length) => {
            head.push_str(&format!("Content-Length: {content_length}\r\n\r\n"));
            out.write_all(head.as_bytes())?;
            out.write_all(&body)?;
        }
    }
    out.flush()?;
    Ok(())
}

#[cfg(test)]
mod declared_length_tests {
    use super::*;
    #[test]
    fn declared_length_is_exact_while_body_is_truncated() {
        let server = MockServer::start_declared_length(|_, _| {
            DeclaredLengthResponse::new(200, b"abc".to_vec(), 37)
        });
        let mut stream = TcpStream::connect(("127.0.0.1", server.port())).unwrap();
        stream
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        stream
            .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n")
            .unwrap();
        let mut wire = Vec::new();
        stream.read_to_end(&mut wire).unwrap();
        let split = wire.windows(4).position(|w| w == b"\r\n\r\n").unwrap();
        let head = std::str::from_utf8(&wire[..split]).unwrap();
        assert_eq!(
            head.lines()
                .filter(|line| line.starts_with("Content-Length:"))
                .collect::<Vec<_>>(),
            vec!["Content-Length: 37"]
        );
        assert_eq!(&wire[split + 4..], b"abc");
    }
}