#![cfg(target_os = "linux")]
use std::fs::{self, OpenOptions};
use std::io::{Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::os::fd::AsRawFd;
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use tungstenite::handshake::client::{Request, generate_key};
use tungstenite::protocol::frame::coding::CloseCode;
use tungstenite::{Message, WebSocket};
static COUNTER: AtomicUsize = AtomicUsize::new(0);
struct TestDir {
path: PathBuf,
}
impl TestDir {
fn new(label: &str) -> Self {
let id = COUNTER.fetch_add(1, Ordering::Relaxed);
let path = std::env::temp_dir().join(format!(
"ks-fhserver-test-{}-{}-{}",
std::process::id(),
label,
id
));
let _ = fs::remove_dir_all(&path);
fs::create_dir_all(&path).unwrap();
Self { path }
}
fn frames_path(&self) -> PathBuf {
self.path.join("frames.bin")
}
}
impl Drop for TestDir {
fn drop(&mut self) {
let _ = fs::remove_dir_all(&self.path);
}
}
struct RunningServer {
child: Child,
addr: std::net::SocketAddr,
}
impl RunningServer {
fn spawn(frames_path: &Path) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let raw_fd = listener.as_raw_fd();
let bin = env!("CARGO_BIN_EXE_knot-static-firehose");
let mut cmd = Command::new(bin);
cmd.arg(frames_path)
.env("LISTEN_FDS", "1")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null());
unsafe {
use std::os::unix::process::CommandExt;
cmd.pre_exec(move || {
if raw_fd != 3 {
if libc::dup2(raw_fd, 3) < 0 {
return Err(std::io::Error::last_os_error());
}
libc::close(raw_fd);
} else {
// Clear FD_CLOEXEC by duping to temporary fd and dup2 back to 3
let tmp = libc::dup2(3, 100);
if tmp < 0 {
return Err(std::io::Error::last_os_error());
}
if libc::dup2(tmp, 3) < 0 {
return Err(std::io::Error::last_os_error());
}
libc::close(tmp);
}
Ok(())
});
}
let child = cmd.spawn().unwrap();
drop(listener);
Self { child, addr }
}
fn connect(&self, cursor: Option<&str>) -> WebSocket<TcpStream> {
let stream = TcpStream::connect(self.addr).unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(15)))
.unwrap();
stream
.set_write_timeout(Some(Duration::from_secs(5)))
.unwrap();
let key = generate_key();
let mut builder = Request::builder()
.method("GET")
.uri("ws://localhost/xrpc/com.atproto.sync.subscribeRepos")
.header("Host", "localhost")
.header("Connection", "Upgrade")
.header("Upgrade", "websocket")
.header("Sec-WebSocket-Version", "13")
.header("Sec-WebSocket-Key", &key)
.header("Ks-Key", &key);
if let Some(c) = cursor {
builder = builder.header("Ks-Cursor", c);
}
let req = builder.body(()).unwrap();
let (ws, resp) = tungstenite::client(req, stream).unwrap();
assert_eq!(resp.status().as_u16(), 101);
ws
}
}
impl Drop for RunningServer {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
fn make_ws_binary_frame(payload: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(10 + payload.len());
out.push(0x82);
if payload.len() < 126 {
out.push(payload.len() as u8);
} else if payload.len() <= 65535 {
out.push(126);
out.extend_from_slice(&(payload.len() as u16).to_be_bytes());
} else {
out.push(127);
out.extend_from_slice(&(payload.len() as u64).to_be_bytes());
}
out.extend_from_slice(payload);
out
}
#[test]
fn replay_from_zero_middle_tip_and_empty_cursor() {
let dir = TestDir::new("replay");
let frames_path = dir.frames_path();
let payloads: Vec<Vec<u8>> = (1..=6u8)
.map(|seq| vec![0xa1, 0x63, b's', b'e', b'q', seq])
.collect();
let mut file_bytes = Vec::new();
for p in &payloads {
file_bytes.extend_from_slice(&make_ws_binary_frame(p));
}
fs::write(&frames_path, &file_bytes).unwrap();
let srv = RunningServer::spawn(&frames_path);
// 1. Replay from 0 -> receives all 6 frames
let mut ws0 = srv.connect(Some("0"));
for expected in &payloads {
match ws0.read().unwrap() {
Message::Binary(b) => assert_eq!(&b, expected),
other => panic!("unexpected message: {other:?}"),
}
}
// 2. Replay from middle (cursor = 3 -> frame index 3 onward, i.e. seq 4, 5, 6)
let mut ws3 = srv.connect(Some("3"));
for expected in &payloads[3..] {
match ws3.read().unwrap() {
Message::Binary(b) => assert_eq!(&b, expected),
other => panic!("unexpected message: {other:?}"),
}
}
// 3. Replay from tip (cursor = 6), empty cursor (""), and missing Ks-Cursor (None) -> live only
let mut ws_tip = srv.connect(Some("6"));
let mut ws_empty = srv.connect(Some(""));
let mut ws_absent = srv.connect(None);
ws_tip
.get_ref()
.set_read_timeout(Some(Duration::from_millis(150)))
.unwrap();
assert!(
ws_tip.read().is_err(),
"tip cursor should not replay old frames"
);
ws_empty
.get_ref()
.set_read_timeout(Some(Duration::from_millis(150)))
.unwrap();
assert!(
ws_empty.read().is_err(),
"empty cursor should not replay old frames"
);
ws_absent
.get_ref()
.set_read_timeout(Some(Duration::from_millis(150)))
.unwrap();
assert!(
ws_absent.read().is_err(),
"missing Ks-Cursor should not replay old frames"
);
// Append frame 7 and verify both ws_tip and ws_empty receive it live
let frame7_payload = vec![0xa1, 0x63, b's', b'e', b'q', 7];
let mut f = OpenOptions::new().append(true).open(&frames_path).unwrap();
f.write_all(&make_ws_binary_frame(&frame7_payload)).unwrap();
f.flush().unwrap();
ws_tip
.get_ref()
.set_read_timeout(Some(Duration::from_secs(2)))
.unwrap();
ws_empty
.get_ref()
.set_read_timeout(Some(Duration::from_secs(2)))
.unwrap();
match ws_tip.read().unwrap() {
Message::Binary(b) => assert_eq!(b, frame7_payload),
other => panic!("unexpected on ws_tip: {other:?}"),
}
match ws_empty.read().unwrap() {
Message::Binary(b) => assert_eq!(b, frame7_payload),
other => panic!("unexpected on ws_empty: {other:?}"),
}
ws_absent
.get_ref()
.set_read_timeout(Some(Duration::from_secs(2)))
.unwrap();
match ws_absent.read().unwrap() {
Message::Binary(b) => assert_eq!(b, frame7_payload),
other => panic!("unexpected on ws_absent: {other:?}"),
}
}
#[test]
fn waits_for_missing_file_and_withholds_incomplete_frames_until_whole() {
let dir = TestDir::new("incomplete");
let frames_path = dir.frames_path();
// Start server BEFORE frames_path exists
let srv = RunningServer::spawn(&frames_path);
let mut ws = srv.connect(Some(""));
// Create the file and write only the first half of a 1000-byte frame
let full_payload: Vec<u8> = (0..1000usize).map(|i| (i % 251) as u8).collect();
let full_frame = make_ws_binary_frame(&full_payload);
let split = 400;
let mut f = OpenOptions::new()
.create(true)
.write(true)
.append(true)
.open(&frames_path)
.unwrap();
f.write_all(&full_frame[..split]).unwrap();
f.flush().unwrap();
// Verify no bytes of the incomplete frame are sent to the client
ws.get_ref()
.set_read_timeout(Some(Duration::from_millis(300)))
.unwrap();
assert!(
ws.read().is_err(),
"incomplete frame must not be sent before all bytes are written"
);
// Now write the rest of the frame
f.write_all(&full_frame[split..]).unwrap();
f.flush().unwrap();
ws.get_ref()
.set_read_timeout(Some(Duration::from_secs(2)))
.unwrap();
match ws.read().unwrap() {
Message::Binary(b) => assert_eq!(b, full_payload),
other => panic!("expected complete binary frame, got {other:?}"),
}
}
#[test]
fn future_cursor_sends_error_frame_and_closes_1008() {
let dir = TestDir::new("future");
let frames_path = dir.frames_path();
let p1 = vec![1, 2, 3, 4];
fs::write(&frames_path, make_ws_binary_frame(&p1)).unwrap();
let srv = RunningServer::spawn(&frames_path);
// Only 1 complete frame exists; cursor = 2 is in the future
let mut ws = srv.connect(Some("2"));
let first = ws.read().unwrap();
let expected_cbor: &[u8] = &[
0xa1, 0x62, b'o', b'p', 0x20, 0xa1, 0x65, b'e', b'r', b'r', b'o', b'r', 0x6c, b'F', b'u',
b't', b'u', b'r', b'e', b'C', b'u', b'r', b's', b'o', b'r',
];
assert_eq!(first, Message::Binary(expected_cbor.to_vec()));
let second = ws.read().unwrap();
match second {
Message::Close(Some(cf)) => {
assert_eq!(cf.code, CloseCode::Policy);
}
other => panic!("expected Close(1008), got {other:?}"),
}
}
#[test]
fn garbage_and_oversized_requests_are_closed_and_client_junk_is_ignored() {
let dir = TestDir::new("garbage");
let frames_path = dir.frames_path();
fs::write(&frames_path, b"").unwrap();
let srv = RunningServer::spawn(&frames_path);
// 1. Garbage request line
{
let mut s = TcpStream::connect(srv.addr).unwrap();
s.set_read_timeout(Some(Duration::from_secs(2))).unwrap();
s.write_all(
b"GET /wrong HTTP/1.1\r\nKs-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nKs-Cursor: 0\r\n\r\n",
)
.unwrap();
let mut buf = [0u8; 64];
let n = s.read(&mut buf).unwrap_or(0);
assert_eq!(n, 0, "garbage request must be closed without reply");
}
// 2. Oversized request (> 4 KiB)
{
let mut s = TcpStream::connect(srv.addr).unwrap();
s.set_read_timeout(Some(Duration::from_secs(2))).unwrap();
let mut req = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: 0\r\n\
X-Pad: "
.to_vec();
req.extend(std::iter::repeat(b'a').take(4100));
req.extend_from_slice(b"\r\n\r\n");
let _ = s.write_all(&req);
let mut buf = [0u8; 64];
let n = s.read(&mut buf).unwrap_or(0);
assert_eq!(n, 0, "oversized request must be closed without reply");
}
// 3. Valid connection sends 64 KiB of unparsed client junk; server drains it and still streams frames
let mut ws = srv.connect(Some(""));
let junk = vec![0xffu8; 65536];
ws.get_mut().write_all(&junk).unwrap();
let payload = b"after-client-junk".to_vec();
let mut f = OpenOptions::new().append(true).open(&frames_path).unwrap();
f.write_all(&make_ws_binary_frame(&payload)).unwrap();
f.flush().unwrap();
match ws.read().unwrap() {
Message::Binary(b) => assert_eq!(b, payload),
other => panic!("expected binary frame after client junk, got {other:?}"),
}
}
#[test]
fn slot_exhaustion_accepts_and_closes_when_full() {
let _ = rustix::process::setrlimit(
rustix::process::Resource::Nofile,
rustix::process::Rlimit {
current: Some(8192),
maximum: Some(8192),
},
);
let dir = TestDir::new("slots");
let frames_path = dir.frames_path();
fs::write(&frames_path, b"").unwrap();
let srv = RunningServer::spawn(&frames_path);
let mut conns = Vec::with_capacity(1024);
for _ in 0..1024 {
conns.push(srv.connect(Some("")));
}
// 1025th connection must be accepted and immediately closed
let mut extra = TcpStream::connect(srv.addr).unwrap();
extra
.set_read_timeout(Some(Duration::from_secs(2)))
.unwrap();
let _ = extra.write_all(
b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: \r\n\r\n",
);
let mut buf = [0u8; 32];
let n = extra.read(&mut buf).unwrap_or(0);
assert_eq!(
n, 0,
"connection past 1024 slots must be closed immediately"
);
// Dropping one slot frees it for a new connection
conns
.pop()
.unwrap()
.get_mut()
.shutdown(Shutdown::Both)
.unwrap();
std::thread::sleep(Duration::from_millis(50));
let mut recovered = srv.connect(Some(""));
let payload = b"slot-reused".to_vec();
fs::write(&frames_path, make_ws_binary_frame(&payload)).unwrap();
match recovered.read().unwrap() {
Message::Binary(b) => assert_eq!(b, payload),
other => panic!("expected frame on recovered slot, got {other:?}"),
}
}
#[test]
fn ping_cadence_and_pings_never_mid_frame_on_large_frames() {
let dir = TestDir::new("pings");
let frames_path = dir.frames_path();
let sock_path = dir.path.join("fh.sock");
// Write a 2 MiB binary frame so sendfile stops mid-frame when the unix socket buffer fills
let big_payload: Vec<u8> = (0..(2 * 1024 * 1024))
.map(|i| (i * 31 % 253) as u8)
.collect();
let big_frame = make_ws_binary_frame(&big_payload);
fs::write(&frames_path, &big_frame).unwrap();
// Also verifies systemd socket activation with a Unix stream listener on fd 3 (prod setup)
let listener = UnixListener::bind(&sock_path).unwrap();
let raw_fd = listener.as_raw_fd();
let bin = env!("CARGO_BIN_EXE_knot-static-firehose");
let mut cmd = Command::new(bin);
cmd.arg(&frames_path)
.env("LISTEN_FDS", "1")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null());
unsafe {
use std::os::unix::process::CommandExt;
cmd.pre_exec(move || {
if raw_fd != 3 {
if libc::dup2(raw_fd, 3) < 0 {
return Err(std::io::Error::last_os_error());
}
libc::close(raw_fd);
}
Ok(())
});
}
let mut child = cmd.spawn().unwrap();
drop(listener);
let stream = UnixStream::connect(&sock_path).unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(15)))
.unwrap();
let key = generate_key();
let req = Request::builder()
.method("GET")
.uri("ws://localhost/xrpc/com.atproto.sync.subscribeRepos")
.header("Host", "localhost")
.header("Connection", "Upgrade")
.header("Upgrade", "websocket")
.header("Sec-WebSocket-Version", "13")
.header("Sec-WebSocket-Key", &key)
.header("Ks-Key", &key)
.header("Ks-Cursor", "0")
.body(())
.unwrap();
let (mut ws, _) = tungstenite::client(req, stream).unwrap();
// Wait 10.5 seconds while the 2 MiB frame is partially in flight in kernel socket buffers.
// The 10s ping timer will fire while the server is mid-frame (offset < frame_end).
// If the server ever injected the 2-byte ping mid-frame, the 2 MiB frame payload would be
// corrupted and shifted by 2 bytes, followed by garbage at the end.
let start = Instant::now();
std::thread::sleep(Duration::from_millis(10_500));
// Now read from the WebSocket: we must receive the exact, uncorrupted 2 MiB binary frame
// first, followed immediately by the queued Ping(empty) at the frame boundary!
let msg1 = ws.read().unwrap();
match msg1 {
Message::Binary(b) => {
assert_eq!(b.len(), big_payload.len());
assert_eq!(b, big_payload, "ping must never be interleaved mid-frame!");
}
other => panic!("expected 2 MiB binary frame first, got {other:?}"),
}
let msg2 = ws.read().unwrap();
assert_eq!(
msg2,
Message::Ping(Vec::new()),
"expected empty ping right at the frame boundary after 10s"
);
assert!(start.elapsed() >= Duration::from_secs(10));
let _ = child.kill();
let _ = child.wait();
}
fn read_voluntary_ctxt_switches(pid: u32) -> u64 {
let status = fs::read_to_string(format!("/proc/{pid}/status")).unwrap();
for line in status.lines() {
if let Some(rest) = line.strip_prefix("voluntary_ctxt_switches:") {
return rest.trim().parse().unwrap();
}
}
panic!("voluntary_ctxt_switches not found in /proc/{pid}/status");
}
#[test]
fn idle_server_with_connected_client_sleeps_until_deadline() {
let dir = TestDir::new("idle_wakeups");
let frames_path = dir.frames_path();
fs::write(&frames_path, b"").unwrap();
let srv = RunningServer::spawn(&frames_path);
let _ws = srv.connect(None);
// Let handshake settle, then measure voluntary context switches over 3 seconds of idle time
std::thread::sleep(Duration::from_millis(50));
let pid = srv.child.id();
let before = read_voluntary_ctxt_switches(pid);
std::thread::sleep(Duration::from_secs(3));
let after = read_voluntary_ctxt_switches(pid);
let delta = after.saturating_sub(before);
assert!(
delta <= 2,
"idle server with one connected client woke up {delta} times in 3s (before={before}, after={after})"
);
}