#![cfg_attr(all(target_os = "linux", panic = "abort"), no_std)]
#![cfg_attr(all(target_os = "linux", panic = "abort"), no_main)]
#[cfg(target_os = "linux")]
mod linux_main {
use core::ffi::CStr;
use core::mem::MaybeUninit;
use heapless::Vec as HVec;
use knot_static_firehose::{
CursorTarget, FUTURE_CURSOR_AND_CLOSE, FrameHeader, FrameIndex, HANDSHAKE_RESPONSE_LEN,
MAX_REQ_BYTES, PING_FRAME, ParseStatus, format_handshake_response, parse_request,
parse_ws_frame_header,
};
use rustix::event::epoll::{self, CreateFlags as EpollCreate, EventData, EventFlags};
use rustix::fd::{AsFd as _, BorrowedFd, OwnedFd};
use rustix::fs::{
CWD, Mode, OFlags, fstat,
inotify::{self, CreateFlags as InotifyCreate, WatchFlags},
openat, sendfile,
};
use rustix::io::{Errno, ioctl_fionbio, pread, read, retry_on_intr, write};
use rustix::net::{SocketFlags, accept_with};
use rustix::process::{Resource, Rlimit, setrlimit};
use rustix::time::{ClockId, Timespec, clock_gettime};
const LISTEN_RAW_FD: i32 = 3;
const MAX_CONNS: usize = 1024;
const TAG_LISTEN: u64 = u64::MAX;
const TAG_INOTIFY: u64 = u64::MAX - 1;
const HANDSHAKE_TIMEOUT_MS: u64 = 5_000;
const PING_INTERVAL_MS: u64 = 10_000;
const SLOW_CLIENT_TIMEOUT_MS: u64 = 60_000;
const _: () = assert!(HANDSHAKE_RESPONSE_LEN <= 160 && FUTURE_CURSOR_AND_CLOSE.len() <= 160);
const EPOLL_ERR_MASK: EventFlags = EventFlags::ERR
.union(EventFlags::HUP)
.union(EventFlags::RDHUP);
const EPOLL_IN_RDHUP: EventFlags = EventFlags::IN.union(EventFlags::RDHUP);
#[derive(Clone, Copy, PartialEq, Eq)]
enum ConnState {
ReadingRequest,
SendingHandshake,
SendingFutureCursor,
Streaming,
}
struct Conn {
fd: OwnedFd,
state: ConnState,
epoll_out: bool,
ping_pending: bool,
out_pos: usize,
out_buf: HVec<u8, 160>,
cursor: Option<u64>,
offset: u64,
frame_end: u64,
last_ping_ms: u64,
last_advance_ms: u64,
accepted_ms: u64,
}
impl Conn {
fn new(fd: OwnedFd, now: u64) -> Self {
Self {
fd,
state: ConnState::ReadingRequest,
epoll_out: false,
ping_pending: false,
out_pos: 0,
out_buf: HVec::new(),
cursor: None,
offset: 0,
frame_end: 0,
last_ping_ms: 0,
last_advance_ms: 0,
accepted_ms: now,
}
}
fn queue_out(&mut self, bytes: &[u8]) {
self.out_buf.clear();
let _ = self.out_buf.extend_from_slice(bytes);
self.out_pos = 0;
}
fn has_pending_out(&self) -> bool {
self.out_pos < self.out_buf.len()
}
}
struct ServerState {
conns: [Option<Conn>; MAX_CONNS],
req_bufs: [HVec<u8, MAX_REQ_BYTES>; MAX_CONNS],
index: FrameIndex,
scratch: [u8; 16384],
frames_fd: Option<OwnedFd>,
frames_ino: u64,
slot_bound: usize,
}
static mut STATE: ServerState = ServerState {
conns: [const { None }; MAX_CONNS],
req_bufs: [const { HVec::new() }; MAX_CONNS],
index: FrameIndex::new(),
scratch: [0; 16384],
frames_fd: None,
frames_ino: 0,
slot_bound: 0,
};
fn now_ms() -> u64 {
let ts = clock_gettime(ClockId::Monotonic);
(ts.tv_sec as u64)
.saturating_mul(1000)
.saturating_add((ts.tv_nsec as u64) / 1_000_000)
}
fn fstat_ino_size(fd: BorrowedFd<'_>) -> Option<(u64, u64)> {
let st = fstat(fd).ok()?;
(st.st_size >= 0).then_some((st.st_ino as u64, st.st_size as u64))
}
fn pread_all(fd: BorrowedFd<'_>, mut offset: u64, buf: &mut [u8]) -> usize {
let mut total = 0;
while total < buf.len() {
match retry_on_intr(|| pread(fd, &mut buf[total..], offset)) {
Ok(0) | Err(_) => break,
Ok(n) => {
total += n;
offset += n as u64;
}
}
}
total
}
fn setup_inotify_dir_watch(inotify_fd: BorrowedFd<'_>, frames_path: &CStr) {
let path_bytes = frames_path.to_bytes();
let mut dir_buf = HVec::<u8, 4096>::new();
let dir_cstr: &CStr = match path_bytes.iter().rposition(|&b| b == b'/') {
None => c".",
Some(0) => c"/",
Some(idx) => {
if dir_buf.extend_from_slice(&path_bytes[..idx]).is_err()
|| dir_buf.push(0).is_err()
{
return;
}
let Ok(cstr) = CStr::from_bytes_with_nul(&dir_buf) else {
return;
};
cstr
}
};
let _ = inotify::add_watch(
inotify_fd,
dir_cstr,
WatchFlags::MODIFY | WatchFlags::CREATE | WatchFlags::MOVED_TO | WatchFlags::ATTRIB,
);
}
fn refresh_frames_file(st: &mut ServerState, inotify_fd: BorrowedFd<'_>, frames_path: &CStr) {
if let Ok(new_fd) = openat(
CWD,
frames_path,
OFlags::RDONLY | OFlags::CLOEXEC,
Mode::empty(),
) && let Some((new_ino, _)) = fstat_ino_size(new_fd.as_fd())
&& (st.frames_fd.is_none() || new_ino != st.frames_ino)
{
st.frames_fd = Some(new_fd);
st.frames_ino = new_ino;
st.index.reset();
let _ = inotify::add_watch(
inotify_fd,
frames_path,
WatchFlags::MODIFY | WatchFlags::ATTRIB,
);
}
if let Some(fd) = st.frames_fd.as_ref()
&& let Some((_, size)) = fstat_ino_size(fd.as_fd())
{
if size < st.index.tip_offset {
st.index.reset();
}
let bfd = fd.as_fd();
st.index.advance(size, |off, buf| pread_all(bfd, off, buf));
}
}
fn close_slot(st: &mut ServerState, epoll_fd: BorrowedFd<'_>, slot: usize) {
if let Some(conn) = st.conns[slot].take() {
let _ = epoll::delete(epoll_fd, &conn.fd);
}
st.req_bufs[slot].clear();
while st.slot_bound > 0 && st.conns[st.slot_bound - 1].is_none() {
st.slot_bound -= 1;
}
}
fn update_slot_epoll(
st: &mut ServerState,
epoll_fd: BorrowedFd<'_>,
slot: usize,
want_out: bool,
) {
let Some(conn) = st.conns[slot].as_mut() else {
return;
};
if conn.epoll_out == want_out {
return;
}
conn.epoll_out = want_out;
let flags = if want_out {
EPOLL_IN_RDHUP | EventFlags::OUT
} else {
EPOLL_IN_RDHUP
};
let _ = epoll::modify(epoll_fd, &conn.fd, EventData::new_u64(slot as u64), flags);
}
fn flush_out_buf(conn: &mut Conn) -> Option<bool> {
while conn.out_pos < conn.out_buf.len() {
match retry_on_intr(|| write(&conn.fd, &conn.out_buf[conn.out_pos..])) {
Ok(n) if n > 0 => conn.out_pos += n,
Err(Errno::AGAIN) => return Some(false),
_ => return None,
}
}
conn.out_buf.clear();
conn.out_pos = 0;
Some(true)
}
fn drive_write(st: &mut ServerState, epoll_fd: BorrowedFd<'_>, slot: usize, now: u64) {
loop {
let Some(state) = st.conns[slot].as_ref().map(|c| c.state) else {
return;
};
match state {
ConnState::ReadingRequest => return,
ConnState::SendingHandshake => {
let flushed = flush_out_buf(st.conns[slot].as_mut().unwrap());
match flushed {
None => {
close_slot(st, epoll_fd, slot);
return;
}
Some(false) => {
update_slot_epoll(st, epoll_fd, slot, true);
return;
}
Some(true) => {
let cursor = st.conns[slot].as_ref().unwrap().cursor;
let target = match st.frames_fd.as_ref() {
None => match cursor {
None | Some(0) => CursorTarget::StreamFrom { offset: 0 },
Some(_) => CursorTarget::FutureCursor,
},
Some(fd) => {
let bfd = fd.as_fd();
st.index
.resolve_cursor(cursor, |off, buf| pread_all(bfd, off, buf))
}
};
let conn = st.conns[slot].as_mut().unwrap();
match target {
CursorTarget::InvalidStream => {
close_slot(st, epoll_fd, slot);
return;
}
CursorTarget::FutureCursor => {
conn.queue_out(&FUTURE_CURSOR_AND_CLOSE);
conn.state = ConnState::SendingFutureCursor;
}
CursorTarget::StreamFrom { offset } => {
conn.offset = offset;
conn.frame_end = offset;
conn.last_ping_ms = now;
conn.last_advance_ms = now;
conn.state = ConnState::Streaming;
}
}
}
}
}
ConnState::SendingFutureCursor => {
match flush_out_buf(st.conns[slot].as_mut().unwrap()) {
None | Some(true) => {
close_slot(st, epoll_fd, slot);
return;
}
Some(false) => {
update_slot_epoll(st, epoll_fd, slot, true);
return;
}
}
}
ConnState::Streaming => {
if st.conns[slot].as_ref().unwrap().has_pending_out() {
match flush_out_buf(st.conns[slot].as_mut().unwrap()) {
None => {
close_slot(st, epoll_fd, slot);
return;
}
Some(false) => {
update_slot_epoll(st, epoll_fd, slot, true);
return;
}
Some(true) => {}
}
}
{
let conn = st.conns[slot].as_mut().unwrap();
if conn.offset == conn.frame_end && conn.ping_pending {
conn.ping_pending = false;
conn.last_ping_ms = now;
conn.queue_out(&PING_FRAME);
match flush_out_buf(conn) {
None => {
close_slot(st, epoll_fd, slot);
return;
}
Some(false) => {
update_slot_epoll(st, epoll_fd, slot, true);
return;
}
Some(true) => {}
}
}
}
let (offset, frame_end) = {
let c = st.conns[slot].as_ref().unwrap();
(c.offset, c.frame_end)
};
if offset == frame_end {
if let Some(fd) = st.frames_fd.as_ref()
&& offset < st.index.tip_offset
{
let mut hdr = [0u8; 10];
let n = pread_all(fd.as_fd(), offset, &mut hdr);
match parse_ws_frame_header(&hdr[..n]) {
FrameHeader::Complete { total_len, .. } => {
let Some(end) = offset.checked_add(total_len) else {
close_slot(st, epoll_fd, slot);
return;
};
if end > st.index.tip_offset {
close_slot(st, epoll_fd, slot);
return;
}
st.conns[slot].as_mut().unwrap().frame_end = end;
}
FrameHeader::Incomplete | FrameHeader::Invalid => {
close_slot(st, epoll_fd, slot);
return;
}
}
} else {
st.conns[slot].as_mut().unwrap().last_advance_ms = now;
update_slot_epoll(st, epoll_fd, slot, false);
return;
}
}
let Some(frames_fd) = st.frames_fd.as_ref() else {
close_slot(st, epoll_fd, slot);
return;
};
let conn = st.conns[slot].as_mut().unwrap();
let rem = conn.frame_end - conn.offset;
let chunk = (rem as usize).min(1 << 20);
let mut off = conn.offset;
match retry_on_intr(|| sendfile(&conn.fd, frames_fd, Some(&mut off), chunk)) {
Ok(n) if n > 0 => {
conn.offset = off;
conn.last_advance_ms = now;
}
Ok(_) | Err(Errno::AGAIN) => {
update_slot_epoll(st, epoll_fd, slot, true);
return;
}
Err(_) => {
close_slot(st, epoll_fd, slot);
return;
}
}
}
}
}
}
fn handle_client_read(
st: &mut ServerState,
epoll_fd: BorrowedFd<'_>,
inotify_fd: BorrowedFd<'_>,
frames_path: &CStr,
slot: usize,
now: u64,
) {
let Some(state) = st.conns[slot].as_ref().map(|c| c.state) else {
return;
};
match state {
ConnState::ReadingRequest => loop {
if st.req_bufs[slot].is_full() {
close_slot(st, epoll_fd, slot);
return;
}
let rem = MAX_REQ_BYTES - st.req_bufs[slot].len();
let conn_fd = st.conns[slot].as_ref().unwrap().fd.as_fd();
match retry_on_intr(|| read(conn_fd, &mut st.scratch[..rem])) {
Ok(n) if n > 0 => {
let _ = st.req_bufs[slot].extend_from_slice(&st.scratch[..n]);
match parse_request(&st.req_bufs[slot]) {
ParseStatus::Incomplete => continue,
ParseStatus::Invalid => {
close_slot(st, epoll_fd, slot);
return;
}
ParseStatus::Complete(parsed) => {
let resp = format_handshake_response(&parsed.key);
let conn = st.conns[slot].as_mut().unwrap();
conn.queue_out(&resp);
conn.cursor = parsed.cursor;
conn.state = ConnState::SendingHandshake;
refresh_frames_file(st, inotify_fd, frames_path);
drive_write(st, epoll_fd, slot, now);
return;
}
}
}
Err(Errno::AGAIN) => return,
_ => {
close_slot(st, epoll_fd, slot);
return;
}
}
},
ConnState::SendingHandshake | ConnState::SendingFutureCursor | ConnState::Streaming => {
loop {
let conn_fd = st.conns[slot].as_ref().unwrap().fd.as_fd();
match retry_on_intr(|| read(conn_fd, &mut st.scratch)) {
Ok(n) if n > 0 => continue,
Err(Errno::AGAIN) => return,
_ => {
close_slot(st, epoll_fd, slot);
return;
}
}
}
}
}
}
fn accept_clients(
st: &mut ServerState,
listen_fd: BorrowedFd<'_>,
epoll_fd: BorrowedFd<'_>,
now: u64,
) {
while let Ok(client_fd) =
retry_on_intr(|| accept_with(listen_fd, SocketFlags::NONBLOCK | SocketFlags::CLOEXEC))
{
let Some(slot) = st.conns.iter().position(Option::is_none) else {
continue;
};
if epoll::add(
epoll_fd,
&client_fd,
EventData::new_u64(slot as u64),
EPOLL_IN_RDHUP,
)
.is_err()
{
continue;
}
st.req_bufs[slot].clear();
st.conns[slot] = Some(Conn::new(client_fd, now));
if slot + 1 > st.slot_bound {
st.slot_bound = slot + 1;
}
}
}
fn serve(
st: &mut ServerState,
listen_fd: BorrowedFd<'_>,
frames_path: &CStr,
) -> rustix::io::Result<core::convert::Infallible> {
let _ = setrlimit(
Resource::Nofile,
Rlimit {
current: Some(8192),
maximum: Some(8192),
},
);
let _ = ioctl_fionbio(listen_fd, true);
let epoll_fd = epoll::create(EpollCreate::CLOEXEC)?;
epoll::add(
&epoll_fd,
listen_fd,
EventData::new_u64(TAG_LISTEN),
EventFlags::IN,
)?;
let inotify_fd = inotify::init(InotifyCreate::NONBLOCK | InotifyCreate::CLOEXEC)?;
epoll::add(
&epoll_fd,
&inotify_fd,
EventData::new_u64(TAG_INOTIFY),
EventFlags::IN,
)?;
setup_inotify_dir_watch(inotify_fd.as_fd(), frames_path);
refresh_frames_file(st, inotify_fd.as_fd(), frames_path);
let mut events = [const { MaybeUninit::uninit() }; 256];
let mut timeout: Option<Timespec> = None;
loop {
let mut file_updated = false;
if let Ok((ready, _)) = epoll::wait(&epoll_fd, &mut events, timeout.as_ref()) {
let now = now_ms();
for ev in ready {
let flags = ev.flags;
let tag = ev.data.u64();
if tag == TAG_LISTEN {
accept_clients(st, listen_fd, epoll_fd.as_fd(), now);
} else if tag == TAG_INOTIFY {
while let Ok(n) = retry_on_intr(|| read(&inotify_fd, &mut st.scratch)) {
if n == 0 {
break;
}
}
file_updated = true;
} else {
let slot = tag as usize;
if slot >= MAX_CONNS || st.conns[slot].is_none() {
continue;
}
if flags.intersects(EPOLL_ERR_MASK) {
close_slot(st, epoll_fd.as_fd(), slot);
continue;
}
if flags.contains(EventFlags::IN) {
handle_client_read(
st,
epoll_fd.as_fd(),
inotify_fd.as_fd(),
frames_path,
slot,
now,
);
}
if st.conns[slot].is_some() && flags.contains(EventFlags::OUT) {
drive_write(st, epoll_fd.as_fd(), slot, now);
}
}
}
}
let now = now_ms();
let prev_tip = st.index.tip_offset;
if file_updated {
refresh_frames_file(st, inotify_fd.as_fd(), frames_path);
}
let tip_advanced = st.index.tip_offset > prev_tip;
let mut next_deadline_ms: u64 = u64::MAX;
for slot in 0..st.slot_bound {
let Some(state) = st.conns[slot].as_ref().map(|c| c.state) else {
continue;
};
match state {
ConnState::ReadingRequest
| ConnState::SendingHandshake
| ConnState::SendingFutureCursor => {
let dl = st.conns[slot]
.as_ref()
.unwrap()
.accepted_ms
.saturating_add(HANDSHAKE_TIMEOUT_MS);
if now >= dl {
close_slot(st, epoll_fd.as_fd(), slot);
continue;
}
next_deadline_ms = next_deadline_ms.min(dl);
}
ConnState::Streaming => {
let needs_drive = {
let conn = st.conns[slot].as_mut().unwrap();
let is_behind =
conn.offset < st.index.tip_offset || conn.has_pending_out();
if is_behind {
let slow_dl =
conn.last_advance_ms.saturating_add(SLOW_CLIENT_TIMEOUT_MS);
if now >= slow_dl {
close_slot(st, epoll_fd.as_fd(), slot);
continue;
}
} else {
conn.last_advance_ms = now;
}
if now.saturating_sub(conn.last_ping_ms) >= PING_INTERVAL_MS {
conn.ping_pending = true;
}
tip_advanced || (conn.ping_pending && conn.offset == conn.frame_end)
};
if needs_drive {
drive_write(st, epoll_fd.as_fd(), slot, now);
}
let Some(conn) = st.conns[slot].as_mut() else {
continue;
};
if conn.state != ConnState::Streaming {
continue;
}
if conn.offset < st.index.tip_offset || conn.has_pending_out() {
let slow_dl =
conn.last_advance_ms.saturating_add(SLOW_CLIENT_TIMEOUT_MS);
next_deadline_ms = next_deadline_ms.min(slow_dl);
} else {
conn.last_advance_ms = now;
}
if !conn.ping_pending {
let ping_dl = conn.last_ping_ms.saturating_add(PING_INTERVAL_MS);
next_deadline_ms = next_deadline_ms.min(ping_dl);
}
}
}
}
timeout = (next_deadline_ms != u64::MAX).then(|| {
let wait_ms = next_deadline_ms.saturating_sub(now);
Timespec {
tv_sec: (wait_ms / 1000) as _,
tv_nsec: ((wait_ms % 1000) * 1_000_000) as _,
}
});
}
}
pub unsafe fn run(argc: libc::c_int, argv: *const *const libc::c_char) -> ! {
unsafe {
if argc < 2 {
libc::_exit(2);
}
let argv1 = *argv.add(1);
if argv1.is_null() {
libc::_exit(2);
}
let frames_path = CStr::from_ptr(argv1);
let listen_fd = BorrowedFd::borrow_raw(LISTEN_RAW_FD);
let Err(_) = serve(&mut *(&raw mut STATE), listen_fd, frames_path);
libc::_exit(1);
}
}
#[cfg(panic = "abort")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn main(
argc: libc::c_int,
argv: *const *const libc::c_char,
) -> libc::c_int {
unsafe { run(argc, argv) }
}
#[cfg(panic = "abort")]
#[panic_handler]
fn panic(_: &core::panic::PanicInfo) -> ! {
unsafe { libc::_exit(1) }
}
#[cfg(panic = "abort")]
#[unsafe(no_mangle)]
extern "C" fn rust_eh_personality() {}
}
#[cfg(not(all(target_os = "linux", panic = "abort")))]
fn main() {
#[cfg(target_os = "linux")]
{
use std::os::unix::ffi::OsStrExt as _;
let mut args: Vec<std::ffi::CString> = std::env::args_os()
.map(|s| std::ffi::CString::new(s.as_bytes()).unwrap())
.collect();
let mut ptrs: Vec<*const libc::c_char> = args.iter_mut().map(|s| s.as_ptr()).collect();
let argc = ptrs.len() as libc::c_int;
ptrs.push(core::ptr::null());
unsafe { linux_main::run(argc, ptrs.as_ptr()) }
}
}