boxmaker/crates/inferproxy/tests/forward.rs

//! End-to-end tests for the forwarder. Do not edit: these define the required behaviour.
//!
//! Each test starts a small TCP server as the upstream, runs `serve` on a Unix socket in a
//! temporary directory, and talks to it as a client would.

use inferproxy::{Limits, serve};
use std::io::{Read, Write};
use std::net::{Shutdown, TcpListener};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::PathBuf;
use std::sync::atomic::{AtomicU32, Ordering};
use std::thread;
use std::time::Duration;

static NEXT: AtomicU32 = AtomicU32::new(0);

fn socket_path() -> PathBuf {
    let n = NEXT.fetch_add(1, Ordering::SeqCst);
    let dir = std::env::temp_dir().join(format!("inferproxy-test-{}-{n}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    dir.join("infer.sock")
}

/// An upstream that reads until the client stops sending, then answers with what it read,
/// upper-cased, and closes. This is the shape of one HTTP exchange with `Connection: close`.
fn shouting_upstream() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap().to_string();
    thread::spawn(move || {
        for stream in listener.incoming() {
            let mut stream = stream.unwrap();
            thread::spawn(move || {
                let mut got = Vec::new();
                stream.read_to_end(&mut got).unwrap();
                stream.write_all(&got.to_ascii_uppercase()).unwrap();
            });
        }
    });
    addr
}

/// An upstream that accepts and then holds every connection open without reading or writing.
fn silent_upstream() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap().to_string();
    thread::spawn(move || {
        let mut held = Vec::new();
        for stream in listener.incoming() {
            held.push(stream.unwrap());
        }
    });
    addr
}

fn start(upstream: String, limits: Limits) -> PathBuf {
    let path = socket_path();
    let listener = UnixListener::bind(&path).unwrap();
    thread::spawn(move || serve(listener, upstream, limits));
    path
}

fn exchange(path: &PathBuf, request: &[u8]) -> Vec<u8> {
    let mut s = UnixStream::connect(path).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
    s.write_all(request).unwrap();
    s.shutdown(Shutdown::Write).unwrap();
    let mut got = Vec::new();
    s.read_to_end(&mut got).unwrap();
    got
}

/// True if the proxy closed the connection without sending anything.
fn was_refused(path: &PathBuf) -> bool {
    let mut s = UnixStream::connect(path).unwrap();
    s.set_read_timeout(Some(Duration::from_millis(500)))
        .unwrap();
    let mut buf = [0u8; 1];
    matches!(s.read(&mut buf), Ok(0))
}

#[test]
fn forwards_both_ways_and_passes_the_half_close_on() {
    let path = start(shouting_upstream(), Limits::default());
    assert_eq!(exchange(&path, b"hello, box"), b"HELLO, BOX");
}

#[test]
fn forwards_more_than_one_buffer() {
    let path = start(shouting_upstream(), Limits::default());
    let request: Vec<u8> = (0..300_000u32).map(|i| b'a' + (i % 26) as u8).collect();
    assert_eq!(exchange(&path, &request), request.to_ascii_uppercase());
}

#[test]
fn serves_connections_one_after_another() {
    let path = start(
        shouting_upstream(),
        Limits {
            max_connections: 1,
            burst: 100,
            per_second: 100,
        },
    );
    for i in 0..5 {
        let msg = format!("request {i}");
        assert_eq!(
            exchange(&path, msg.as_bytes()),
            msg.to_ascii_uppercase().as_bytes()
        );
    }
}

#[test]
fn refuses_connections_over_the_open_limit() {
    let limits = Limits {
        max_connections: 2,
        burst: 100,
        per_second: 100,
    };
    let path = start(silent_upstream(), limits);
    let _a = UnixStream::connect(&path).unwrap();
    let _b = UnixStream::connect(&path).unwrap();
    thread::sleep(Duration::from_millis(200)); // let the proxy take both
    assert!(
        was_refused(&path),
        "a third open connection must be closed at once"
    );
}

#[test]
fn refuses_connections_over_the_rate_limit() {
    let limits = Limits {
        max_connections: 100,
        burst: 3,
        per_second: 1,
    };
    let path = start(shouting_upstream(), limits);
    for i in 0..3 {
        assert_eq!(
            exchange(&path, b"x"),
            b"X",
            "connection {i} is within the burst"
        );
    }
    assert!(
        was_refused(&path),
        "the fourth connection in the same instant must be refused"
    );
}

#[test]
fn an_unreachable_upstream_closes_the_client() {
    // Port 1 on localhost refuses connections.
    let path = start("127.0.0.1:1".to_string(), Limits::default());
    assert!(was_refused(&path));
}

/// An upstream that answers at once and closes, whether or not the client has finished sending.
/// This is what `llama-server` does, and what a server that dies mid-answer looks like.
fn answer_and_close_upstream() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap().to_string();
    thread::spawn(move || {
        for stream in listener.incoming() {
            let mut stream = stream.unwrap();
            let mut first = [0u8; 1];
            let _ = stream.read(&mut first);
            stream.write_all(b"answer").unwrap();
            // dropping `stream` closes it
        }
    });
    addr
}

#[test]
fn the_upstreams_close_reaches_a_client_that_is_still_sending() {
    // The HTTP client in loopd never half-closes: it keeps its sending side open until it drops
    // the connection. When the server closes, the proxy must close towards the client at once,
    // or the client only learns of a dead server from its own timeout.
    let path = start(answer_and_close_upstream(), Limits::default());
    let mut s = UnixStream::connect(&path).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(2))).unwrap();
    s.write_all(b"request").unwrap();
    let started = std::time::Instant::now();
    let mut got = Vec::new();
    let result = s.read_to_end(&mut got);
    assert!(
        result.is_ok(),
        "the read must end with EOF, not a timeout: {result:?}"
    );
    assert_eq!(got, b"answer");
    assert!(
        started.elapsed() < Duration::from_millis(1000),
        "EOF took {:?}",
        started.elapsed()
    );
}