Add the inferproxy byte forwarder with its two limits

Implemented-By: Laguna S 2.1 (OpenCode)

crates/inferproxy/src/lib.rs

@@ -1 +1,135 @@
 //! Forwards bytes between `infer.sock` and the llama-server router. It logs nothing.
+
+use std::io::{self, copy};
+use std::net::{Shutdown, TcpStream};
+use std::os::unix::net::{UnixListener, UnixStream};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::thread;
+use std::time::Instant;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct Limits {
+    pub max_connections: usize,
+    pub burst: u32,
+    pub per_second: u32,
+}
+
+impl Default for Limits {
+    fn default() -> Self {
+        Limits {
+            max_connections: 8,
+            burst: 10,
+            per_second: 2,
+        }
+    }
+}
+
+/// Holds at most `burst` tokens and gains `per_second` tokens each second.
+pub struct TokenBucket {
+    tokens: u64,
+    burst: u64,
+    per_second: u32,
+    last: Instant,
+}
+
+impl TokenBucket {
+    pub fn new(burst: u32, per_second: u32, now: Instant) -> Self {
+        let burst = u64::from(burst) * 1000;
+        TokenBucket {
+            tokens: burst,
+            burst,
+            per_second,
+            last: now,
+        }
+    }
+
+    pub fn take(&mut self, now: Instant) -> bool {
+        let elapsed = now.saturating_duration_since(self.last);
+        let earned = refill_thousandths(self.per_second, elapsed);
+        let total = u128::from(self.tokens).saturating_add(earned);
+        let capped = total.min(u128::from(self.burst));
+        self.tokens = u64::try_from(capped).unwrap_or(self.burst);
+        self.last = now;
+        if self.tokens >= 1000 {
+            self.tokens -= 1000;
+            true
+        } else {
+            false
+        }
+    }
+}
+
+/// Thousandths of a token earned over `elapsed` at `per_second` tokens per second.
+/// per_second * elapsed_nanoseconds / 1_000_000 keeps partial refills additive without
+/// floating point, so checks 100 ms apart at 2/s accumulate correctly.
+fn refill_thousandths(per_second: u32, elapsed: std::time::Duration) -> u128 {
+    let nanos = u128::from(per_second).saturating_mul(elapsed.as_nanos());
+    nanos.div_euclid(1_000_000)
+}
+
+struct OpenGuard {
+    open: Arc<AtomicUsize>,
+}
+
+impl OpenGuard {
+    fn new(open: Arc<AtomicUsize>) -> Self {
+        open.fetch_add(1, Ordering::SeqCst);
+        OpenGuard { open }
+    }
+}
+
+impl Drop for OpenGuard {
+    fn drop(&mut self) {
+        self.open.fetch_sub(1, Ordering::SeqCst);
+    }
+}
+
+/// Accepts connections forever. Returns only if `accept` itself fails.
+pub fn serve(listener: UnixListener, upstream: String, limits: Limits) -> io::Result<()> {
+    let open = Arc::new(AtomicUsize::new(0));
+    let mut bucket = TokenBucket::new(limits.burst, limits.per_second, Instant::now());
+    loop {
+        let (client, _) = listener.accept()?;
+        if !bucket.take(Instant::now()) {
+            eprintln!("inferproxy: refused (rate limit)");
+            drop(client);
+            continue;
+        }
+        if open.load(Ordering::SeqCst) >= limits.max_connections {
+            eprintln!("inferproxy: refused (connection limit)");
+            drop(client);
+            continue;
+        }
+        let open = Arc::clone(&open);
+        let up = upstream.clone();
+        thread::spawn(move || handle(client, up, open));
+    }
+}
+
+fn handle(client: UnixStream, upstream: String, open: Arc<AtomicUsize>) {
+    let _guard = OpenGuard::new(open);
+    let server = match TcpStream::connect(&upstream) {
+        Ok(s) => s,
+        Err(_) => return,
+    };
+    let _ = forward(&client, &server);
+    drop(_guard);
+}
+
+fn forward(client: &UnixStream, server: &TcpStream) -> io::Result<()> {
+    let mut c2s_in = client.try_clone()?;
+    let mut s2c_out = client.try_clone()?;
+    let mut s2c_in = server.try_clone()?;
+    let mut c2s_out = server.try_clone()?;
+    thread::scope(|s| {
+        s.spawn(move || {
+            let _ = copy(&mut c2s_in, &mut c2s_out);
+            let _ = c2s_out.shutdown(Shutdown::Write);
+        });
+        s.spawn(move || {
+            let _ = copy(&mut s2c_in, &mut s2c_out);
+        });
+    });
+    Ok(())
+}

crates/inferproxy/src/main.rs

@@ -1,4 +1,68 @@
+use std::env;
+use std::fs;
+use std::os::unix::fs::PermissionsExt;
+use std::path::PathBuf;
+use std::process;
+
+use inferproxy::{Limits, serve};
+
 fn main() {
-    eprintln!("inferproxy: not implemented until M2");
-    std::process::exit(2);
+    let args: Vec<String> = env::args().collect();
+    let mut listen: Option<String> = None;
+    let mut upstream: Option<String> = None;
+    let mut i = 1;
+    while i < args.len() {
+        match args[i].as_str() {
+            "--listen" => {
+                i += 1;
+                listen = args.get(i).cloned();
+            }
+            "--upstream" => {
+                i += 1;
+                upstream = args.get(i).cloned();
+            }
+            _ => usage(),
+        }
+        i += 1;
+    }
+    let listen = match listen {
+        Some(p) => p,
+        None => usage(),
+    };
+    let upstream = match upstream {
+        Some(p) => p,
+        None => usage(),
+    };
+
+    let path: PathBuf = listen.into();
+    match fs::remove_file(&path) {
+        Ok(()) => {}
+        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
+        Err(e) => {
+            eprintln!("inferproxy: cannot remove stale socket: {e}");
+            process::exit(1);
+        }
+    }
+    let listener = match std::os::unix::net::UnixListener::bind(&path) {
+        Ok(l) => l,
+        Err(e) => {
+            eprintln!("inferproxy: bind failed: {e}");
+            process::exit(1);
+        }
+    };
+    let perms = fs::Permissions::from_mode(0o600);
+    if let Err(e) = fs::set_permissions(&path, perms) {
+        eprintln!("inferproxy: chmod failed: {e}");
+        process::exit(1);
+    }
+    eprintln!("inferproxy: {} -> {}", path.display(), upstream);
+    if let Err(e) = serve(listener, upstream, Limits::default()) {
+        eprintln!("inferproxy: {e}");
+        process::exit(1);
+    }
+}
+
+fn usage() -> ! {
+    eprintln!("usage: inferproxy --listen <socket path> --upstream <host:port>");
+    process::exit(2);
 }

crates/inferproxy/tests/bucket.rs

@@ -0,0 +1,70 @@
+//! Tests for the accept-rate limiter. Do not edit: these define the required behaviour.
+
+use inferproxy::{Limits, TokenBucket};
+use std::time::{Duration, Instant};
+
+#[test]
+fn default_limits() {
+    let want = Limits {
+        max_connections: 8,
+        burst: 10,
+        per_second: 2,
+    };
+    assert_eq!(Limits::default(), want);
+}
+
+#[test]
+fn starts_full_and_empties() {
+    let t0 = Instant::now();
+    let mut b = TokenBucket::new(3, 2, t0);
+    assert!(b.take(t0));
+    assert!(b.take(t0));
+    assert!(b.take(t0));
+    assert!(!b.take(t0), "a fourth take at the same instant must fail");
+}
+
+#[test]
+fn refills_at_the_given_rate() {
+    let t0 = Instant::now();
+    let mut b = TokenBucket::new(1, 2, t0);
+    assert!(b.take(t0));
+    assert!(
+        !b.take(t0 + Duration::from_millis(499)),
+        "half a second at 2/s is one token"
+    );
+    assert!(b.take(t0 + Duration::from_millis(500)));
+    assert!(!b.take(t0 + Duration::from_millis(500)));
+}
+
+#[test]
+fn partial_refills_add_up() {
+    // Ten checks 100 ms apart at 2/s must add up to two tokens, not zero.
+    let t0 = Instant::now();
+    let mut b = TokenBucket::new(1, 2, t0);
+    assert!(b.take(t0));
+    let granted = (1..=10)
+        .filter(|i| b.take(t0 + Duration::from_millis(100 * i)))
+        .count();
+    assert_eq!(granted, 2);
+}
+
+#[test]
+fn never_holds_more_than_the_burst() {
+    let t0 = Instant::now();
+    let mut b = TokenBucket::new(2, 100, t0);
+    let later = t0 + Duration::from_secs(3600);
+    assert!(b.take(later));
+    assert!(b.take(later));
+    assert!(
+        !b.take(later),
+        "an hour of refill must still stop at the burst size"
+    );
+}
+
+#[test]
+fn a_clock_that_does_not_advance_is_harmless() {
+    let t0 = Instant::now();
+    let mut b = TokenBucket::new(1, 2, t0 + Duration::from_secs(10));
+    assert!(b.take(t0), "an earlier `now` must not panic or underflow");
+    assert!(!b.take(t0));
+}

crates/inferproxy/tests/forward.rs

@@ -0,0 +1,156 @@
+//! 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));
+}