Lumbridge Compute — Apache-2.0

src/gateway.rs

@@ -0,0 +1,102 @@
+//! Streaming-safe TCP gateway for OpenAI-compatible model servers.
+//!
+//! The gateway deliberately stays below HTTP: it forwards bytes unchanged, so
+//! chunked responses and server-sent-event token streams retain their timing and
+//! semantics while clients keep one stable Lumbridge Compute address.
+
+use anyhow::{Context, Result};
+use std::io;
+use std::net::{Shutdown, TcpListener, TcpStream};
+use std::thread;
+
+pub fn run(listen: &str, upstream: &str) -> Result<()> {
+    let listener = TcpListener::bind(listen)
+        .with_context(|| format!("binding Lumbridge gateway at {listen}"))?;
+    println!("Lumbridge gateway {listen} -> {upstream}");
+    serve_listener(listener, upstream, None)
+}
+
+fn serve_listener(
+    listener: TcpListener,
+    upstream: &str,
+    max_connections: Option<usize>,
+) -> Result<()> {
+    let mut accepted = 0usize;
+    for incoming in listener.incoming() {
+        let client = match incoming {
+            Ok(client) => client,
+            Err(error) => {
+                eprintln!("gateway accept failed: {error}");
+                continue;
+            }
+        };
+        let upstream = upstream.to_string();
+        thread::spawn(move || {
+            if let Err(error) = proxy(client, &upstream) {
+                eprintln!("gateway request failed: {error:#}");
+            }
+        });
+        accepted += 1;
+        if max_connections.is_some_and(|limit| accepted >= limit) {
+            break;
+        }
+    }
+    Ok(())
+}
+
+fn proxy(mut client: TcpStream, upstream_addr: &str) -> Result<()> {
+    client.set_nodelay(true).ok();
+    let mut upstream = TcpStream::connect(upstream_addr)
+        .with_context(|| format!("connecting gateway upstream {upstream_addr}"))?;
+    upstream.set_nodelay(true).ok();
+
+    let mut client_reader = client.try_clone()?;
+    let mut upstream_writer = upstream.try_clone()?;
+    let request = thread::spawn(move || -> io::Result<u64> {
+        let copied = io::copy(&mut client_reader, &mut upstream_writer)?;
+        upstream_writer.shutdown(Shutdown::Write).ok();
+        Ok(copied)
+    });
+
+    io::copy(&mut upstream, &mut client)?;
+    client.shutdown(Shutdown::Write).ok();
+    request
+        .join()
+        .map_err(|_| anyhow::anyhow!("gateway request-copy thread panicked"))??;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::{Read, Write};
+
+    #[test]
+    fn gateway_forwards_bidirectional_bytes_without_buffering_protocols() {
+        let upstream = TcpListener::bind("127.0.0.1:0").unwrap();
+        let upstream_addr = upstream.local_addr().unwrap();
+        let upstream_thread = thread::spawn(move || {
+            let (mut socket, _) = upstream.accept().unwrap();
+            let mut request = [0u8; 4];
+            socket.read_exact(&mut request).unwrap();
+            assert_eq!(&request, b"ping");
+            socket.write_all(b"pong").unwrap();
+        });
+
+        let gateway = TcpListener::bind("127.0.0.1:0").unwrap();
+        let gateway_addr = gateway.local_addr().unwrap();
+        let gateway_thread = thread::spawn(move || {
+            serve_listener(gateway, &upstream_addr.to_string(), Some(1)).unwrap();
+        });
+
+        let mut client = TcpStream::connect(gateway_addr).unwrap();
+        client.write_all(b"ping").unwrap();
+        client.shutdown(Shutdown::Write).unwrap();
+        let mut response = Vec::new();
+        client.read_to_end(&mut response).unwrap();
+        assert_eq!(response, b"pong");
+
+        upstream_thread.join().unwrap();
+        gateway_thread.join().unwrap();
+    }
+}