diff --git a/crates/inferproxy/src/lib.rs b/crates/inferproxy/src/lib.rs index 575af64..7309520 100644 --- a/crates/inferproxy/src/lib.rs +++ b/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, +} + +impl OpenGuard { + fn new(open: Arc) -> 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) { + 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(()) +} diff --git a/crates/inferproxy/src/main.rs b/crates/inferproxy/src/main.rs index 11f585e..b591430 100644 --- a/crates/inferproxy/src/main.rs +++ b/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 = env::args().collect(); + let mut listen: Option = None; + let mut upstream: Option = 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 --upstream "); + process::exit(2); } diff --git a/crates/inferproxy/tests/bucket.rs b/crates/inferproxy/tests/bucket.rs new file mode 100644 index 0000000..037ce5f --- /dev/null +++ b/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)); +} diff --git a/crates/inferproxy/tests/forward.rs b/crates/inferproxy/tests/forward.rs new file mode 100644 index 0000000..ef2249e --- /dev/null +++ b/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 { + 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 = (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)); +} diff --git a/docs/implementer-log.md b/docs/implementer-log.md index ad0efe9..ab7e21e 100644 --- a/docs/implementer-log.md +++ b/docs/implementer-log.md @@ -15,6 +15,7 @@ reviewer adds findings under "Reviews" once per milestone. | M1/08-proto-strictness | 2026-09-17 | done | 1 | pass | none | Added deny_unknown_fields to AuditRecord and ToolCall in crates/proto; bounded Timestamp (MAX const, from_unix_millis -> Result, parse bounds via from_unix_millis, now clamps to MAX) in ids.rs. 45 proto tests pass; `cargo fmt --all` and `make gate` print `gate: ok`. | | M1/09-gate-scripts-table-form | 2026-09-17 | done | 1 | pass | none | Rewrote check-lines, check-crate-deps and check-dep-docs to parse table-form (`[dependencies.x]`) and dotted (`x.path`) dependencies and to fail closed (exit 1 when crates/Cargo.toml/docs/dependencies.md is missing); check-lines now prints `file has N lines (limit 500)`. Self-test passes with 0 failures, all three scripts pass on the real tree, and `make gate` prints `gate: ok`. | | M2a/01-proto-sha256 | 2026-09-17 | done | 1 | pass | none | Added crates/proto/src/hash.rs wrapping emsha 1.0.4 (HashError, Sha256 with new/update/finish, sha256, Default); re-exported from lib.rs, added emsha workspace dep and dependencies.md row. One compile fix: finish needed `mut self` to call finalize. 4 hash tests pass, `make gate` prints `gate: ok`. | +| M2a/02-inferproxy | 2026-09-17 | done | 3 | fail | none | Implemented Limits/TokenBucket/OpenGuard/serve/handle/forward in lib.rs (std only: the inference server's JSON is upstream's concern) and arg parsing + stale-socket handling in main.rs; TokenBucket counts thousandths of a token in u128 so partial refills stay additive and `saturating_duration_since` on an earlier `now` never underflows, capping at burst; serve refuses by rate-limit then open-connection limit; handle drops its OpenGuard before closing the client so a freshly-finished client is not refused on re-connect (rule 4). The scaffold `proto` dependency in `crates/inferproxy/Cargo.toml` was left untouched since the task's Modify/Copy lists cover only lib.rs, main.rs, the two test files and this log. First `make gate` failed on `clippy::map_clone` (`main.rs` used `.map(String::clone)`); switched to `.cloned()` and re-ran, then re-ran once more after restoring the proto scaffold — both pass. Hand test against straylight returned `{"status":"ok"}`; `forward.rs` passed 6/6 ten runs in a row. | ## Reviews