//! TCP accept loop, per-connection handling, and the background purge loop. use std::net::TcpStream; use std::sync::Arc; use std::time::Duration; use crate::bind::TransportBindValues; use crate::channel::{handshake, Channel}; use crate::crypto::{b32, derive_public}; use crate::proto::{FETCH_BUDGET, HEADER_TIMEOUT_SECS, KEY_LEN, MALFORMED, PURGE_INTERVAL_SECS}; use crate::ratelimit::RateLimiter; use crate::session::{ServerConfig, Session}; use crate::store::Store; pub struct ServeArgs { pub key_path: String, pub db_path: String, pub host: String, pub port: u16, pub max_envelope: usize, pub quota: i64, pub requests_quota: i64, pub retention_days: i64, pub requests_retention_days: i64, pub max_tokens: u16, pub invite_token: Option, pub rate_connections: u32, pub rate_sends: u32, pub rate_tokens: u32, } pub fn run(args: ServeArgs) -> anyhow::Result<()> { let static_key = std::fs::read(&args.key_path) .map_err(|e| anyhow::anyhow!("cannot read server key: {e}"))?; anyhow::ensure!( static_key.len() == KEY_LEN, "server key must be {KEY_LEN} raw bytes, got {}", static_key.len() ); let key_array: [u8; KEY_LEN] = static_key.clone().try_into().unwrap(); let server_static = derive_public(&key_array); let config = Arc::new(ServerConfig { max_envelope: args.max_envelope, fetch_budget: FETCH_BUDGET, main_quota: args.quota, requests_quota: args.requests_quota, max_tokens: args.max_tokens, invite_token: args.invite_token.map(String::into_bytes), conn_limiter: RateLimiter::new(args.rate_connections), send_limiter: RateLimiter::new(args.rate_sends), token_limiter: RateLimiter::new(args.rate_tokens), }); let main_retention_secs = args.retention_days * 86400; let requests_retention_secs = args.requests_retention_days * 86400; let purge_db_path = args.db_path.clone(); std::thread::spawn(move || { purge_loop(purge_db_path, main_retention_secs, requests_retention_secs) }); let listener = std::net::TcpListener::bind((args.host.as_str(), args.port))?; log::info!("listening on {}:{}", args.host, args.port); log::info!("server public key: {}", b32(&server_static)); for incoming in listener.incoming() { let stream = match incoming { Ok(s) => s, Err(e) => { log::warn!("accept error: {e}"); continue; } }; let peer_ip = stream .peer_addr() .map(|a| a.ip().to_string()) .unwrap_or_else(|_| "unknown".to_string()); if !config.conn_limiter.allow(&peer_ip) { log::warn!("rate limited {peer_ip}"); continue; } let config = Arc::clone(&config); let static_key = key_array; let db_path = args.db_path.clone(); std::thread::spawn(move || { if let Err(e) = handle_connection(stream, &config, &static_key, &db_path, &peer_ip) { log::info!("connection error from {peer_ip}: {e}"); } }); } Ok(()) } /// Also the hostile harness's entry point: it drives real connections /// through the same accept/handshake/session path the TCP carrier serves. pub(crate) fn handle_connection( mut stream: TcpStream, config: &ServerConfig, static_key: &[u8; KEY_LEN], db_path: &str, peer_ip: &str, ) -> anyhow::Result<()> { // The handshake runs under the header budget; once it completes, // Channel::read_noise picks between the header and idle deadlines per // frame state. stream.set_read_timeout(Some(Duration::from_secs(HEADER_TIMEOUT_SECS)))?; let handshake_started = std::time::Instant::now(); let (transport, handshake_hash) = match handshake(&mut stream, static_key) { Ok(v) => v, Err(e) => { log::info!("handshake failed from {peer_ip}: {e}"); return Ok(()); } }; log::debug!( "handshake from {peer_ip} in {:?}", handshake_started.elapsed() ); // Drop logs the summary once, whatever return path closes the session. let summary = SessionSummary { peer_ip, started: std::time::Instant::now(), ops: std::cell::Cell::new(0), }; let store = Store::open(db_path)?; let server_static = derive_public(static_key); let bind = TransportBindValues::tcp(&handshake_hash, &server_static)?; let mut session = Session::new(config, store, peer_ip.to_string(), bind); let mut channel = Channel::new(stream, transport); loop { let (op, body) = match channel.read_frame() { Ok(v) => v, Err(e) => { if is_eof_like(&e) { return Ok(()); } log::info!("bad frame from {peer_ip}: {e}"); let _ = channel.write_frame(0, &[MALFORMED]); return Ok(()); } }; summary.ops.set(summary.ops.get() + 1); let (status, payload) = session.dispatch(op, &body); let mut response = Vec::with_capacity(1 + payload.len()); response.push(status); response.extend_from_slice(&payload); if channel.write_frame(op, &response).is_err() { return Ok(()); } } } struct SessionSummary<'a> { peer_ip: &'a str, started: std::time::Instant, ops: std::cell::Cell, } impl Drop for SessionSummary<'_> { fn drop(&mut self) { log::debug!( "session {} closed: {} op(s) in {:?}", self.peer_ip, self.ops.get(), self.started.elapsed() ); } } fn is_eof_like(e: &anyhow::Error) -> bool { if let Some(io_err) = e.downcast_ref::() { return matches!( io_err.kind(), std::io::ErrorKind::UnexpectedEof | std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::BrokenPipe | std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock ); } false } fn purge_loop(db_path: String, main_retention_secs: i64, requests_retention_secs: i64) { let store = match Store::open(&db_path) { Ok(s) => s, Err(e) => { log::error!("purge thread failed to open store: {e}"); return; } }; loop { std::thread::sleep(Duration::from_secs(PURGE_INTERVAL_SECS)); let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs() as i64; match store.purge(now - main_retention_secs, now - requests_retention_secs) { Ok(0) => {} Ok(n) => log::info!("expired {n} message(s)"), Err(e) => log::error!("purge failed: {e}"), } } }