mod ban;
mod handler;

use ban::{BanTracker, is_bannable_address, parse_ip_allowlist};
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};

/// True once a response actually came from a plan file, rather than being
/// the username echoed back unchanged or an InvalidInput rejection. A
/// "failure" (the negation) is timestamped against the client IP by the
/// caller; enough failures within the rolling window trips the ban.
fn is_plan_served(response: &str, username: &str) -> bool {
    response != username && !response.starts_with("InvalidInput:")
}

fn trim_trailing_crlf(s: &mut String) {
    while matches!(s.chars().last(), Some('\r') | Some('\n')) {
        s.pop();
    }
}

async fn handle_connection(
    mut socket: TcpStream,
    client_addr: String,
    trackable: bool,
    bans: Arc<Mutex<BanTracker>>,
    start: Instant,
) {
    let now = start.elapsed();

    // An IP that has racked up too many failed lookups (scanners, username
    // guessers, non-finger junk) is dropped without being read or answered.
    // Only globally-routable addresses are tracked: behind Docker's bridge
    // every client is SNAT'd to the gateway, so banning there would block
    // everyone at once (see is_bannable_address()).
    if trackable && bans.lock().unwrap().is_blocked(&client_addr, now) {
        println!("finger drop from {client_addr}: blocked");
        return;
    }

    let mut data = [0u8; 1024];
    let bytes_read = match socket.read(&mut data).await {
        // Client hung up before sending a request: health checks (which
        // connect and immediately close), port scanners, and reset
        // connections all land here. This is normal -- don't log it as an
        // error.
        Ok(0) | Err(_) => return,
        Ok(n) => n,
    };

    let mut username = String::from_utf8_lossy(&data[..bytes_read]).into_owned();
    trim_trailing_crlf(&mut username);

    println!("finger request from {client_addr} for user '{username}'");
    let response = handler::process(&username);

    // A "failure" is simply any request that does not resolve to a readable
    // plan file: an unknown user, rejected input, or non-finger junk. Each
    // failure is timestamped against the client IP; once an IP exceeds the
    // threshold within the rolling window, the is_blocked() check above
    // starts dropping its connections. This also frustrates username
    // guessing.
    if !is_plan_served(&response, &username) {
        if trackable {
            let res = bans.lock().unwrap().record_offense(&client_addr, now);
            let suffix = if res.blocked { " -- now blocked" } else { "" };
            println!(
                "finger miss from {client_addr} for '{username}' ({} failures in window){suffix}",
                res.count
            );
        } else {
            println!("finger miss from {client_addr} for '{username}' (not tracked)");
        }
        // Best-effort reply; ignore write errors (the client may have
        // already gone away).
        let _ = socket.write_all(b"No plan found\r\n").await;
        return;
    }

    let _ = socket.write_all(response.as_bytes()).await;
}

async fn listener(
    bans: Arc<Mutex<BanTracker>>,
    allowlist: Arc<HashSet<String>>,
    start: Instant,
) -> std::io::Result<()> {
    let acceptor = TcpListener::bind(("0.0.0.0", 79)).await?;
    loop {
        let (socket, peer) = acceptor.accept().await?;
        let client_addr = peer.ip().to_string();
        // Allowlisted IPs (trusted aggregating front-ends like the
        // finger-web proxy) are never tracked, so their bursts neither block
        // them nor count as offenses.
        let trackable = is_bannable_address(peer.ip()) && !allowlist.contains(&client_addr);
        let bans = bans.clone();
        tokio::spawn(handle_connection(socket, client_addr, trackable, bans, start));
    }
}

// Periodically prune offense records that have aged out of the window so
// the tracker's memory stays bounded even for IPs that never reconnect.
async fn sweeper(bans: Arc<Mutex<BanTracker>>, start: Instant) {
    let mut interval = tokio::time::interval(Duration::from_secs(600));
    interval.tick().await; // first tick fires immediately; skip it
    loop {
        interval.tick().await;
        bans.lock().unwrap().sweep(start.elapsed());
    }
}

async fn wait_for_shutdown_signal() {
    let ctrl_c = tokio::signal::ctrl_c();
    #[cfg(unix)]
    {
        let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler");
        tokio::select! {
            _ = ctrl_c => {}
            _ = term.recv() => {}
        }
    }
    #[cfg(not(unix))]
    {
        let _ = ctrl_c.await;
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let start = Instant::now();
    let bans = Arc::new(Mutex::new(BanTracker::new()));

    let allow_env = std::env::var("FINGER_BAN_ALLOWLIST").unwrap_or_default();
    let allowlist = Arc::new(parse_ip_allowlist(&allow_env));
    for ip in allowlist.iter() {
        println!("ban allowlist: {ip} (never tracked or blocked)");
    }

    tokio::spawn(sweeper(bans.clone(), start));

    tokio::select! {
        res = listener(bans, allowlist, start) => {
            if let Err(e) = res {
                println!("fatal exception: {e}");
            }
        }
        _ = wait_for_shutdown_signal() => {}
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn plan_served_when_response_differs_and_not_invalid() {
        assert!(is_plan_served("Out to lunch.\r\n", "pete"));
    }

    #[test]
    fn plan_not_served_when_response_echoes_username() {
        assert!(!is_plan_served("pete", "pete"));
    }

    #[test]
    fn plan_not_served_on_invalid_input() {
        assert!(!is_plan_served(
            "InvalidInput: Path detected in username\r\n",
            "user/name"
        ));
    }

    #[test]
    fn trims_trailing_cr_and_lf() {
        let mut s = String::from("pete\r\n");
        trim_trailing_crlf(&mut s);
        assert_eq!(s, "pete");
    }

    #[test]
    fn trims_bare_lf_only() {
        let mut s = String::from("pete\n");
        trim_trailing_crlf(&mut s);
        assert_eq!(s, "pete");
    }

    #[test]
    fn leaves_string_without_trailing_crlf_untouched() {
        let mut s = String::from("pete");
        trim_trailing_crlf(&mut s);
        assert_eq!(s, "pete");
    }
}