//! Slowing down guessing (oxygenhub#221): password sign-in, the 2FA code

//! step, password reset requests and device codes.

//!

//! Each key (an account or an address, within one action) counts its

//! failures. A few are free; after that each one doubles the wait before the

//! next try, up to a short lockout. A quiet hour forgets the count, and a

//! success clears the account's. Kept in memory: a restart forgets it, which

//! costs an attacker a restart they cannot cause.

use std::collections::HashMap;

/// Failures an account may have before waits start.

const ACCOUNT_FREE: u32 = 5;

/// Failures an address may have before waits start: higher, since many

/// people can share one address.

const ADDRESS_FREE: u32 = 20;

/// The first wait.

const FIRST_WAIT_MS: i64 = 2_000;

/// The longest wait: the lockout.

pub const LOCKOUT_MS: i64 = 15 * 60_000;

/// A key with no failure for this long starts again.

const FORGET_MS: i64 = 60 * 60_000;

/// Keys kept before stale ones are swept.

const SWEEP_AT: usize = 50_000;

#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]

pub enum Kind {

    Account,

    Address,

}

/// What is counted: an account or an address, within one action.

#[derive(Clone, PartialEq, Eq, Hash, Debug)]

pub struct Key {

    action: &'static str,

    kind: Kind,

    id: String,

}

impl Key {

    /// An account, by name; names compare without case.

    pub fn account(action: &'static str, name: &str) -> Key {

        Key { action, kind: Kind::Account, id: name.trim().to_lowercase() }

    }

    /// An address. An empty one (unknown) is never counted.

    pub fn address(action: &'static str, addr: &str) -> Key {

        Key { action, kind: Kind::Address, id: addr.to_string() }

    }

    fn free(&self) -> u32 {

        match self.kind {

            Kind::Account => ACCOUNT_FREE,

            Kind::Address => ADDRESS_FREE,

        }

    }

}

#[derive(Default)]

struct Count {

    failures: u32,

    last_ms: i64,

    until_ms: i64,

}

#[derive(Default)]

pub struct Throttle {

    counts: HashMap<Key, Count>,

}

impl Throttle {

    /// How long, in milliseconds, before any of these keys may try again;

    /// `None` when all may try now.

    pub fn wait(&self, keys: &[Key], now_ms: i64) -> Option<i64> {

        keys.iter()

            .filter_map(|k| self.counts.get(k))

            .filter(|c| now_ms - c.last_ms < FORGET_MS)

            .map(|c| c.until_ms - now_ms)

            .filter(|w| *w > 0)

            .max()

    }

    /// Count a failure (or, for an action with no failure, an attempt)

    /// against each key; the wait it leaves, if any.

    pub fn failed(&mut self, keys: &[Key], now_ms: i64) -> Option<i64> {

        if self.counts.len() >= SWEEP_AT {

            self.counts.retain(|_, c| now_ms - c.last_ms < FORGET_MS);

        }

        for key in keys.iter().filter(|k| !k.id.is_empty()) {

            let count = self.counts.entry(key.clone()).or_default();

            if now_ms - count.last_ms >= FORGET_MS {

                *count = Count::default();

            }

            count.failures += 1;

            count.last_ms = now_ms;

            if count.failures > key.free() {

                let doublings = (count.failures - key.free() - 1).min(20);

                count.until_ms = now_ms + (FIRST_WAIT_MS << doublings).min(LOCKOUT_MS);

            }

        }

        self.wait(keys, now_ms)

    }

    /// A success: the key starts again.

    pub fn cleared(&mut self, key: &Key) {

        self.counts.remove(key);

    }

}

/// "Too many tries: try again in 2 minutes." for a wait.

pub fn message(wait_ms: i64) -> String {

    let secs = (wait_ms + 999) / 1000;

    let when = match secs {

        ..=1 => "a second".to_string(),

        2..=59 => format!("{secs} seconds"),

        _ => {

            let mins = (secs + 59) / 60;

            if mins == 1 {

                "a minute".to_string()

            } else {

                format!("{mins} minutes")

            }

        }

    };

    format!("Too many tries: try again in {when}.")

}

#[cfg(test)]

mod tests {

    use super::*;

    const T0: i64 = 1_000_000_000;

    #[test]

    fn an_account_gets_a_few_free_failures_then_doubling_waits_up_to_the_lockout() {

        let mut t = Throttle::default();

        let k = [Key::account("login", "Alice")];

        for _ in 0..ACCOUNT_FREE {

            assert_eq!(t.wait(&k, T0), None);

            t.failed(&k, T0);

        }

        assert_eq!(t.wait(&k, T0), None, "the free failures cost nothing");

        assert_eq!(t.failed(&k, T0), Some(2_000), "the failure says the wait it starts");

        t.failed(&k, T0);

        assert_eq!(t.wait(&k, T0), Some(4_000));

        for _ in 0..20 {

            t.failed(&k, T0);

        }

        assert_eq!(t.wait(&k, T0), Some(LOCKOUT_MS), "capped at the lockout");

        assert_eq!(t.wait(&k, T0 + LOCKOUT_MS), None, "and it ends");

        assert_eq!(t.wait(&[Key::account("login", " alice ")], T0), Some(LOCKOUT_MS), "names compare without case");

    }

    #[test]

    fn a_success_clears_the_account_and_a_quiet_hour_forgets() {

        let mut t = Throttle::default();

        let k = [Key::account("login", "bob")];

        for _ in 0..10 {

            t.failed(&k, T0);

        }

        assert!(t.wait(&k, T0).is_some());

        t.cleared(&k[0]);

        assert_eq!(t.wait(&k, T0), None);

        for _ in 0..10 {

            t.failed(&k, T0);

        }

        let later = T0 + FORGET_MS;

        assert_eq!(t.wait(&k, later), None);

        t.failed(&k, later);

        assert_eq!(t.wait(&k, later), None, "the count started again");

    }

    #[test]

    fn addresses_get_more_room_and_actions_and_unknown_addresses_are_apart() {

        let mut t = Throttle::default();

        let addr = [Key::address("login", "203.0.113.9")];

        for _ in 0..ADDRESS_FREE {

            t.failed(&addr, T0);

        }

        assert_eq!(t.wait(&addr, T0), None);

        t.failed(&addr, T0);

        assert_eq!(t.wait(&addr, T0), Some(2_000));

        assert_eq!(t.wait(&[Key::address("reset", "203.0.113.9")], T0), None, "another action counts apart");

        let unknown = [Key::address("login", "")];

        for _ in 0..100 {

            t.failed(&unknown, T0);

        }

        assert_eq!(t.wait(&unknown, T0), None, "an unknown address is never counted");

        // The longest wait of the keys given is the one that counts.

        let both = [Key::account("login", "carol"), addr[0].clone()];

        assert_eq!(t.wait(&both, T0), Some(2_000));

    }

    #[test]

    fn the_message_says_when() {

        assert_eq!(message(900), "Too many tries: try again in a second.");

        assert_eq!(message(4_000), "Too many tries: try again in 4 seconds.");

        assert_eq!(message(61_000), "Too many tries: try again in 2 minutes.");

        assert_eq!(message(LOCKOUT_MS), "Too many tries: try again in 15 minutes.");

    }

}