bunshin/src/server.rs

//! 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, IDLE_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<String>,
    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(())
}

fn handle_connection(
    mut stream: TcpStream,
    config: &ServerConfig,
    static_key: &[u8; KEY_LEN],
    db_path: &str,
    peer_ip: &str,
) -> anyhow::Result<()> {
    stream.set_read_timeout(Some(Duration::from_secs(IDLE_TIMEOUT_SECS)))?;

    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(());
        }
    };

    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(());
            }
        };

        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(());
        }
    }
}

fn is_eof_like(e: &anyhow::Error) -> bool {
    if let Some(io_err) = e.downcast_ref::<std::io::Error>() {
        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}"),
        }
    }
}