opencord 0.2.1

//! End-to-end encryption primitives for opencord.
//!
//! This module handles:
//! - **X25519 key exchange**: Each client generates an ephemeral keypair. When two
//!   peers exchange public keys, they can derive the same shared secret independently
//!   (Diffie-Hellman). This is how Alice and Bob agree on a secret without ever
//!   sending it over the wire.
//! - **HKDF key derivation**: Raw DH output isn't uniformly random, so we run it
//!   through HKDF (HMAC-based Key Derivation Function) to produce a proper AES key.
//! - **AES-256-GCM encryption**: Symmetric encryption for messages. GCM provides
//!   both confidentiality (can't read it) and integrity (can't tamper with it).
//! - **Group key distribution**: The host generates a random AES key and encrypts
//!   it individually for each peer using their pairwise DH-derived key.
//!
//! The server never sees any of this — it just relays opaque ciphertext.

use aes_gcm::aead::{Aead, KeyInit};
use aes_gcm::{Aes256Gcm, Nonce};
use base64::engine::general_purpose::STANDARD as BASE64;
use base64::Engine;
use hkdf::Hkdf;
use rand::RngCore;
use sha2::Sha256;
use x25519_dalek::{PublicKey, StaticSecret};
use zeroize::Zeroize;

use std::fmt;

// ---------------------------------------------------------------------------
// Error type
// ---------------------------------------------------------------------------

/// Errors that can occur during cryptographic operations.
///
/// We implement `Display` and `Error` manually rather than pulling in the
/// `thiserror` crate — keeps our dependency tree small.
#[derive(Debug)]
pub enum CryptoError {
    /// base64 string couldn't be decoded.
    InvalidBase64,
    /// Decoded bytes were the wrong length (e.g. expected 32, got something else).
    InvalidKeyLength,
    /// AES-GCM decryption failed — wrong key, corrupted data, or tampered ciphertext.
    DecryptionFailed,
}

/// `Display` is Rust's "human-readable formatting" trait. It's what gets called
/// when you do `format!("{}", err)` or `println!("{err}")`. We implement it so
/// our errors print nicely.
impl fmt::Display for CryptoError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CryptoError::InvalidBase64 => write!(f, "invalid base64 encoding"),
            CryptoError::InvalidKeyLength => write!(f, "invalid key length"),
            CryptoError::DecryptionFailed => write!(f, "decryption failed"),
        }
    }
}

/// `std::error::Error` is a marker trait that says "this type is an error".
/// Implementing it lets our error work with `?`, `Box<dyn Error>`, and the
/// wider Rust error ecosystem.
impl std::error::Error for CryptoError {}

// ---------------------------------------------------------------------------
// Keypair (X25519)
// ---------------------------------------------------------------------------

/// An ephemeral X25519 keypair.
///
/// `StaticSecret` is a 32-byte secret scalar. Despite the name "static", we use
/// it for ephemeral (one-session) keys. `PublicKey` is the corresponding point
/// on the Curve25519 elliptic curve.
///
/// Think of it like this:
/// - The secret key is a random number only you know
/// - The public key is derived from it and safe to share with anyone
/// - Two people can combine (their secret + other's public) to get the same shared secret
pub struct Keypair {
    secret: StaticSecret,
    public: PublicKey,
}

impl Keypair {
    /// Generate a new random keypair using the OS random number generator.
    ///
    /// We generate 32 random bytes and construct the `StaticSecret` from them
    /// via `From<[u8; 32]>`, rather than using `random_from_rng`, to avoid
    /// version conflicts between rand_core 0.6 (used by x25519-dalek) and
    /// rand 0.9 (used by our crate).
    pub fn generate() -> Self {
        let mut bytes = [0u8; 32];
        rand::rng().fill_bytes(&mut bytes);
        let secret = StaticSecret::from(bytes);
        // Zeroize the seed bytes now that StaticSecret has its own copy.
        // This prevents the plaintext key material from lingering on the stack.
        bytes.zeroize();
        let public = PublicKey::from(&secret);
        Keypair { secret, public }
    }

    /// Get a reference to our public key.
    pub fn public_key(&self) -> &PublicKey {
        &self.public
    }

    /// Perform Diffie-Hellman key exchange with a peer's public key, then run
    /// the result through HKDF to produce a proper 32-byte AES key.
    ///
    /// **Why HKDF?** The raw DH output has mathematical structure — certain bit
    /// patterns are more likely than others. HKDF "extracts" the entropy and
    /// "expands" it into a uniformly random key. Using raw DH output directly
    /// as an AES key would be subtly insecure.
    ///
    /// `[u8; 32]` is a fixed-size array of 32 bytes — Rust's way of saying
    /// "exactly 32 bytes, known at compile time". This is different from `Vec<u8>`
    /// which is a heap-allocated, variable-length byte vector.
    pub fn derive_shared_key(&self, peer_public: &PublicKey) -> [u8; 32] {
        // DH: my_secret * their_public = shared_point
        let shared_secret = self.secret.diffie_hellman(peer_public);

        // HKDF: extract + expand into a 32-byte key
        // The "salt" and "info" parameters add domain separation — even if two
        // protocols used the same DH output, different info strings would produce
        // different keys.
        let hkdf = Hkdf::<Sha256>::new(Some(b"opencord-v1"), shared_secret.as_bytes());
        let mut key = [0u8; 32];
        hkdf.expand(b"pairwise-key", &mut key)
            .expect("32 bytes is a valid HKDF output length");
        key
    }
}

// ---------------------------------------------------------------------------
// Public key encoding (for JSON transport)
// ---------------------------------------------------------------------------

/// Encode an X25519 public key as a base64 string for sending over JSON.
///
/// Public keys are 32 bytes. Base64 turns arbitrary bytes into printable ASCII
/// characters, making them safe to embed in JSON strings.
pub fn encode_public_key(key: &PublicKey) -> String {
    BASE64.encode(key.as_bytes())
}

/// Decode a base64 string back into an X25519 public key.
///
/// `try_into()` converts a `Vec<u8>` (variable length) into a `[u8; 32]`
/// (fixed length). It fails if the Vec isn't exactly 32 bytes, which is why
/// we map the error to `CryptoError::InvalidKeyLength`.
pub fn decode_public_key(encoded: &str) -> Result<PublicKey, CryptoError> {
    let bytes = BASE64.decode(encoded).map_err(|_| CryptoError::InvalidBase64)?;
    let array: [u8; 32] = bytes
        .try_into()
        .map_err(|_| CryptoError::InvalidKeyLength)?;
    Ok(PublicKey::from(array))
}

// ---------------------------------------------------------------------------
// Group key generation
// ---------------------------------------------------------------------------

/// Generate a random 32-byte AES-256 group key.
///
/// This is the symmetric key that all peers will share. Only the host generates
/// it, then distributes it to each peer encrypted with their pairwise key.
pub fn generate_group_key() -> [u8; 32] {
    let mut key = [0u8; 32];
    rand::rng().fill_bytes(&mut key);
    key
}

// ---------------------------------------------------------------------------
// AES-256-GCM encryption / decryption
// ---------------------------------------------------------------------------

/// Encrypt plaintext with AES-256-GCM, returning (nonce_base64, ciphertext_base64).
///
/// **AES-256-GCM** is an "authenticated encryption" scheme:
/// - AES-256 = the cipher (256-bit key)
/// - GCM = Galois/Counter Mode (provides both encryption AND a MAC for integrity)
///
/// The **nonce** (number used once) is a random 12-byte value. Each encryption
/// MUST use a unique nonce with the same key. With random 12-byte nonces,
/// the collision probability after 2^32 messages is about 2^-33 — negligible
/// for ephemeral chat sessions that last minutes to hours.
pub fn encrypt(key: &[u8; 32], plaintext: &[u8]) -> (String, String) {
    let cipher = Aes256Gcm::new_from_slice(key).expect("32-byte key is valid for AES-256");

    // Generate a random 12-byte nonce
    let nonce_bytes: [u8; 12] = {
        let mut buf = [0u8; 12];
        rand::rng().fill_bytes(&mut buf);
        buf
    };
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .expect("encryption should not fail with valid key and nonce");

    (BASE64.encode(nonce_bytes), BASE64.encode(ciphertext))
}

/// Decrypt ciphertext with AES-256-GCM.
///
/// Returns the plaintext bytes, or `DecryptionFailed` if the key is wrong,
/// the nonce is wrong, or the ciphertext has been tampered with. GCM's
/// authentication tag catches all of these cases.
pub fn decrypt(
    key: &[u8; 32],
    nonce_b64: &str,
    ciphertext_b64: &str,
) -> Result<Vec<u8>, CryptoError> {
    let cipher = Aes256Gcm::new_from_slice(key).expect("32-byte key is valid for AES-256");

    let nonce_bytes = BASE64
        .decode(nonce_b64)
        .map_err(|_| CryptoError::InvalidBase64)?;
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = BASE64
        .decode(ciphertext_b64)
        .map_err(|_| CryptoError::InvalidBase64)?;

    cipher
        .decrypt(nonce, ciphertext.as_ref())
        .map_err(|_| CryptoError::DecryptionFailed)
}

// ---------------------------------------------------------------------------
// Group key wrapping (encrypt/decrypt the group key with a pairwise key)
// ---------------------------------------------------------------------------

/// Encrypt the group key using a pairwise DH-derived key.
///
/// Returns a base64 string containing nonce + ciphertext concatenated.
/// The group key is 32 bytes; after AES-GCM encryption it becomes 32 + 16 (tag) = 48 bytes,
/// plus the 12-byte nonce = 60 bytes total before base64 encoding.
pub fn encrypt_group_key(pairwise_key: &[u8; 32], group_key: &[u8; 32]) -> String {
    let (nonce_b64, ciphertext_b64) = encrypt(pairwise_key, group_key);
    // Pack both into a single string separated by `:` for easy transport
    format!("{nonce_b64}:{ciphertext_b64}")
}

/// Decrypt the group key using a pairwise DH-derived key.
///
/// `mut` on the return value lets us zeroize it when we're done — good practice
/// for key material even though our sessions are ephemeral.
pub fn decrypt_group_key(
    pairwise_key: &[u8; 32],
    encrypted: &str,
) -> Result<[u8; 32], CryptoError> {
    let parts: Vec<&str> = encrypted.splitn(2, ':').collect();
    if parts.len() != 2 {
        return Err(CryptoError::InvalidBase64);
    }

    let mut plaintext = decrypt(pairwise_key, parts[0], parts[1])?;

    // `try_into()` converts Vec<u8> → [u8; 32]. Fails if not exactly 32 bytes.
    let key: [u8; 32] = plaintext
        .as_slice()
        .try_into()
        .map_err(|_| CryptoError::InvalidKeyLength)?;

    // Zeroize the temporary Vec so the key doesn't linger in memory
    plaintext.zeroize();

    Ok(key)
}

// ---------------------------------------------------------------------------
// Memory locking (prevent swap)
// ---------------------------------------------------------------------------

/// Lock a 32-byte key into physical RAM so the OS won't swap it to disk.
///
/// `mlock` is a Unix system call that tells the kernel: "keep these pages in
/// physical memory, don't write them to the swap partition." This matters because
/// swap is on-disk, and disk contents survive power cycles — an attacker with
/// physical access could recover key material from swap.
///
/// Returns `true` if locking succeeded. Failure is non-fatal — it just means
/// the OS might swap the page (common on systems with low `RLIMIT_MEMLOCK`).
///
/// On non-Unix platforms this is a no-op that returns `false`.
#[cfg(unix)]
pub fn mlock_key(key: &[u8; 32]) -> bool {
    // SAFETY: we're passing a valid pointer and length to mlock.
    // The pointer is to our own stack/heap memory and the length is exact.
    unsafe { libc::mlock(key.as_ptr().cast(), 32) == 0 }
}

#[cfg(not(unix))]
pub fn mlock_key(_key: &[u8; 32]) -> bool {
    false
}

/// Unlock a previously mlocked key, allowing the OS to swap it again.
/// Call this right before or after zeroizing the key.
#[cfg(unix)]
pub fn munlock_key(key: &[u8; 32]) {
    // SAFETY: same as mlock_key — valid pointer and length.
    unsafe {
        libc::munlock(key.as_ptr().cast(), 32);
    }
}

#[cfg(not(unix))]
pub fn munlock_key(_key: &[u8; 32]) {}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn dh_symmetry() {
        // Alice and Bob each generate keypairs. When they do DH with each
        // other's public keys, they should get the same shared secret.
        let alice = Keypair::generate();
        let bob = Keypair::generate();

        let alice_shared = alice.derive_shared_key(bob.public_key());
        let bob_shared = bob.derive_shared_key(alice.public_key());

        assert_eq!(alice_shared, bob_shared, "DH should be symmetric");
    }

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let key = generate_group_key();
        let plaintext = b"hello opencord!";

        let (nonce, ciphertext) = encrypt(&key, plaintext);
        let decrypted = decrypt(&key, &nonce, &ciphertext).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn wrong_key_rejected() {
        let key1 = generate_group_key();
        let key2 = generate_group_key();
        let plaintext = b"secret stuff";

        let (nonce, ciphertext) = encrypt(&key1, plaintext);
        let result = decrypt(&key2, &nonce, &ciphertext);

        assert!(result.is_err(), "decryption with wrong key should fail");
        assert!(
            matches!(result.unwrap_err(), CryptoError::DecryptionFailed),
            "should be DecryptionFailed"
        );
    }

    #[test]
    fn group_key_roundtrip() {
        let alice = Keypair::generate();
        let bob = Keypair::generate();

        let pairwise = alice.derive_shared_key(bob.public_key());
        let group_key = generate_group_key();

        let encrypted = encrypt_group_key(&pairwise, &group_key);

        // Bob derives the same pairwise key and decrypts
        let bob_pairwise = bob.derive_shared_key(alice.public_key());
        let decrypted = decrypt_group_key(&bob_pairwise, &encrypted).unwrap();

        assert_eq!(decrypted, group_key);
    }

    #[test]
    fn public_key_encode_decode_roundtrip() {
        let kp = Keypair::generate();
        let encoded = encode_public_key(kp.public_key());
        let decoded = decode_public_key(&encoded).unwrap();

        assert_eq!(decoded.as_bytes(), kp.public_key().as_bytes());
    }

    #[test]
    fn invalid_base64_rejected() {
        let result = decode_public_key("not-valid-base64!!!");
        assert!(matches!(result, Err(CryptoError::InvalidBase64)));
    }

    #[test]
    fn wrong_length_key_rejected() {
        // Valid base64 but only 16 bytes, not 32
        let short = BASE64.encode([0u8; 16]);
        let result = decode_public_key(&short);
        assert!(matches!(result, Err(CryptoError::InvalidKeyLength)));
    }
}