feat: implement Smol Mail server in Rust with nix flake deployment

src/server.rs

@@ -0,0 +1,163 @@
+//! 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::channel::{handshake, Channel};
+use crate::crypto::{b32, derive_public};
+use crate::proto::{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 retention_days: i64,
+    pub invite_token: Option<String>,
+    pub rate_connections: u32,
+    pub rate_sends: 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 config = Arc::new(ServerConfig {
+        max_envelope: args.max_envelope,
+        quota: args.quota,
+        invite_token: args.invite_token.map(String::into_bytes),
+        conn_limiter: RateLimiter::new(args.rate_connections),
+        send_limiter: RateLimiter::new(args.rate_sends),
+    });
+
+    let retention_secs = args.retention_days * 86400;
+    let purge_db_path = args.db_path.clone();
+    std::thread::spawn(move || purge_loop(purge_db_path, retention_secs));
+
+    let listener = std::net::TcpListener::bind((args.host.as_str(), args.port))?;
+    log::info!("listening on {}:{}", args.host, args.port);
+
+    let key_array: [u8; KEY_LEN] = static_key.clone().try_into().unwrap();
+    let public = derive_public(&key_array);
+    log::info!("server public key: {}", b32(&public));
+
+    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 = static_key.clone();
+        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],
+    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 mut session = Session::new(config, store, peer_ip.to_string(), handshake_hash);
+    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, 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 - retention_secs) {
+            Ok(0) => {}
+            Ok(n) => log::info!("expired {n} message(s)"),
+            Err(e) => log::error!("purge failed: {e}"),
+        }
+    }
+}