//! Swarm membership by pre-shared key.
//!
//! Mycelium authenticates the sender of every message (`src_pk` can't be forged), so we
//! always know *who* talks to us. With a swarm key configured, only nodes that also know
//! the key may take part: HELLO and HELLO_REPLY carry a tag proving it, and until a peer
//! has sent a valid one, everything else it sends is ignored.
//!
//! The tag is `HMAC-SHA256(key, label ‖ swarm_id ‖ from_pk ‖ to_pk)`. It needs no
//! challenge: it is bound to both keys, and only `from_pk` can send as `from_pk`, so a
//! captured tag is useless to anyone else. Different labels for HELLO and HELLO_REPLY
//! keep one from being passed off as the other. A HELLO has no swarm ID yet (a new node
//! doesn't know it), so its tag binds an empty one.
//!
//! Limits of a shared key: every member holds the same secret, so a single member can't
//! be revoked (forgetting it lasts until it says HELLO again), and changing the key
//! means restarting every node. A member list in the document is the proper successor.

use std::path::Path;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use anyhow::{bail, Context, Result};
use sha2::{Digest, Sha256};

const HELLO_LABEL: &[u8] = b"continuum/hello/1";
const HELLO_REPLY_LABEL: &[u8] = b"continuum/hello-reply/1";

/// The swarm's pre-shared key.
pub(crate) struct SwarmKey {
    key: Vec<u8>,
    rejections: Mutex<Rejections>,
}

/// For rate limiting the log of rejected handshakes.
#[derive(Default)]
struct Rejections {
    last_logged: Option<Instant>,
    suppressed: u64,
}

/// At most one rejection log line per this interval.
const REJECTION_LOG_INTERVAL: Duration = Duration::from_secs(10);

impl SwarmKey {
    /// Reads the key from a file; its bytes are the key, as they are. Generate one with
    /// e.g. `head -c 32 /dev/urandom > swarm.key`.
    pub(crate) fn from_file(path: &Path) -> Result<Self> {
        let key = std::fs::read(path)
            .with_context(|| format!("Reading the swarm key from {}", path.display()))?;
        if key.len() < 16 {
            bail!(
                "The swarm key in {} has {} bytes; use at least 16 (32 random bytes recommended)",
                path.display(),
                key.len()
            );
        }
        Ok(Self::new(key))
    }

    fn new(key: Vec<u8>) -> Self {
        Self {
            key,
            rejections: Mutex::default(),
        }
    }

    /// Records a rejected handshake. Returns how many went unlogged since the last log
    /// line if this one should be logged, None if it should be suppressed.
    pub(crate) fn note_rejection(&self) -> Option<u64> {
        let mut rejections = self
            .rejections
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        let due = rejections
            .last_logged
            .is_none_or(|last| last.elapsed() >= REJECTION_LOG_INTERVAL);
        if !due {
            rejections.suppressed += 1;
            return None;
        }
        rejections.last_logged = Some(Instant::now());
        Some(std::mem::take(&mut rejections.suppressed))
    }

    /// Tag for a HELLO from `from` to `to`.
    pub(crate) fn hello_tag(&self, from: &str, to: &str) -> Vec<u8> {
        self.tag(HELLO_LABEL, "", from, to).to_vec()
    }

    /// Tag for a HELLO_REPLY from `from` to `to`, naming the swarm.
    pub(crate) fn hello_reply_tag(&self, swarm_id: &str, from: &str, to: &str) -> Vec<u8> {
        self.tag(HELLO_REPLY_LABEL, swarm_id, from, to).to_vec()
    }

    pub(crate) fn verify_hello(&self, tag: &[u8], from: &str, to: &str) -> bool {
        constant_time_eq(tag, &self.tag(HELLO_LABEL, "", from, to))
    }

    pub(crate) fn verify_hello_reply(
        &self,
        tag: &[u8],
        swarm_id: &str,
        from: &str,
        to: &str,
    ) -> bool {
        constant_time_eq(tag, &self.tag(HELLO_REPLY_LABEL, swarm_id, from, to))
    }

    fn tag(&self, label: &[u8], swarm_id: &str, from: &str, to: &str) -> [u8; 32] {
        // Length-prefixed fields, so no two different inputs encode the same.
        let mut message = Vec::new();
        for field in [label, swarm_id.as_bytes(), from.as_bytes(), to.as_bytes()] {
            message.extend_from_slice(&(field.len() as u32).to_be_bytes());
            message.extend_from_slice(field);
        }
        hmac_sha256(&self.key, &message)
    }
}

/// HMAC-SHA256 (RFC 2104), on top of the `sha2` crate we already use.
fn hmac_sha256(key: &[u8], message: &[u8]) -> [u8; 32] {
    const BLOCK: usize = 64;

    let mut block_key = [0u8; BLOCK];
    if key.len() > BLOCK {
        block_key[..32].copy_from_slice(&Sha256::digest(key));
    } else {
        block_key[..key.len()].copy_from_slice(key);
    }

    let mut inner = Sha256::new();
    inner.update(block_key.map(|byte| byte ^ 0x36));
    inner.update(message);
    let inner = inner.finalize();

    let mut outer = Sha256::new();
    outer.update(block_key.map(|byte| byte ^ 0x5c));
    outer.update(inner);

    let mut tag = [0u8; 32];
    tag.copy_from_slice(&outer.finalize());
    tag
}

fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    a.len() == b.len() && a.iter().zip(b).fold(0u8, |diff, (x, y)| diff | (x ^ y)) == 0
}

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

    fn hex(bytes: &[u8]) -> String {
        bytes.iter().map(|byte| format!("{byte:02x}")).collect()
    }

    #[test]
    fn hmac_matches_rfc_4231() {
        // Test case 2.
        assert_eq!(
            hex(&hmac_sha256(b"Jefe", b"what do ya want for nothing?")),
            "5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843"
        );
        // Test case 6: a key longer than the block size.
        assert_eq!(
            hex(&hmac_sha256(
                &[0xaa; 131],
                b"Test Using Larger Than Block-Size Key - Hash Key First"
            )),
            "60e431591ee0b67f0d8a26aacbf5b77f8e0bc6213728c5140546040f0ee37f54"
        );
    }

    #[test]
    fn tags_bind_key_peers_and_kind() {
        let key = SwarmKey::new(vec![1; 32]);
        let other = SwarmKey::new(vec![2; 32]);

        let hello = key.hello_tag("A", "B");
        assert!(key.verify_hello(&hello, "A", "B"));
        assert!(!other.verify_hello(&hello, "A", "B"));
        assert!(!key.verify_hello(&hello, "B", "A"));
        assert!(!key.verify_hello(&hello, "A", "C"));
        assert!(!key.verify_hello(&[], "A", "B"));

        let reply = key.hello_reply_tag("swarm", "A", "B");
        assert!(key.verify_hello_reply(&reply, "swarm", "A", "B"));
        assert!(!key.verify_hello_reply(&reply, "other-swarm", "A", "B"));
        // A HELLO tag is no HELLO_REPLY tag, and the other way round.
        assert!(!key.verify_hello_reply(&hello, "", "A", "B"));
        assert!(!key.verify_hello(&key.hello_reply_tag("", "A", "B"), "A", "B"));
    }
}