void_crypto/

pin.rs

//! PIN-based encryption for identity keys.
//!
//! Encrypts signing, recipient, and (optionally) Nostr secret keys using:
//! - Argon2id for PIN → key derivation (memory-hard, resistant to brute force)
//! - AES-256-GCM for authenticated encryption
//!
//! ## Format versions
//!
//! **v1** (legacy): `[1B version=1][16B salt][12B nonce][64B encrypted][16B tag]`
//! - Payload: 32B signing + 32B recipient = 64 bytes
//!
//! **v2** (current): `[1B version=2][16B salt][12B nonce][96B encrypted][16B tag]`
//! - Payload: 32B signing + 32B recipient + 32B nostr = 96 bytes
//!
//! Decryption is backward-compatible: v1 blobs return `nostr = None`.

use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm,
};
use argon2::Argon2;
use rand::RngCore;
use zeroize::Zeroize;

use crate::kdf::{AeadNonce, NostrSecretKey, RecipientSecretKey, SigningSecretKey};

/// Current format version (v2 includes Nostr key).
const VERSION_V2: u8 = 2;

/// Legacy format version (signing + recipient only).
const VERSION_V1: u8 = 1;

/// AAD for identity key encryption (shared across versions).
const AAD: &[u8] = b"void:identity-keys:v1";

/// Header sizes.
const VERSION_LEN: usize = 1;
const SALT_LEN: usize = 16;
const TAG_LEN: usize = 16;

/// Key payload sizes.
const KEYS_LEN_V1: usize = 64; // 32 signing + 32 recipient
const KEYS_LEN_V2: usize = 96; // 32 signing + 32 recipient + 32 nostr

/// Minimum valid encrypted blob size (v1).
const MIN_ENCRYPTED_LEN: usize =
    VERSION_LEN + SALT_LEN + AeadNonce::SIZE + KEYS_LEN_V1 + TAG_LEN;

/// Errors during PIN encryption/decryption.
#[derive(Debug, thiserror::Error)]
pub enum PinError {
    #[error("PIN must not be empty")]
    EmptyPin,

    #[error("key derivation failed: {0}")]
    DerivationFailed(String),

    #[error("encryption failed")]
    EncryptionFailed,

    #[error("decryption failed (wrong PIN or corrupted data)")]
    DecryptionFailed,

    #[error("invalid encrypted data: too short")]
    DataTooShort,

    #[error("unsupported format version: {0}")]
    UnsupportedVersion(u8),
}

/// Encrypt signing, recipient, and Nostr keys using a PIN (v2 format).
///
/// Returns a binary blob containing all data needed for decryption:
/// version, salt, nonce, and the encrypted+authenticated keys.
pub fn encrypt_identity_keys(
    signing: &SigningSecretKey,
    recipient: &RecipientSecretKey,
    nostr: &NostrSecretKey,
    pin: &str,
) -> Result<Vec<u8>, PinError> {
    if pin.is_empty() {
        return Err(PinError::EmptyPin);
    }

    // Generate random salt and nonce
    let mut salt = [0u8; SALT_LEN];
    rand::thread_rng().fill_bytes(&mut salt);
    let nonce = AeadNonce::generate();

    // Derive encryption key from PIN using Argon2id
    let mut derived_key = derive_key_from_pin(pin, &salt)?;

    // Concatenate keys for encryption (v2: 96 bytes)
    let mut plaintext = [0u8; KEYS_LEN_V2];
    plaintext[..32].copy_from_slice(signing.as_bytes());
    plaintext[32..64].copy_from_slice(recipient.as_bytes());
    plaintext[64..96].copy_from_slice(nostr.as_bytes());

    // Encrypt with AES-256-GCM
    let cipher =
        Aes256Gcm::new_from_slice(&derived_key).map_err(|_| PinError::EncryptionFailed)?;
    let gcm_nonce = aes_gcm::Nonce::from_slice(nonce.as_bytes());

    let ciphertext = match cipher.encrypt(
        gcm_nonce,
        aes_gcm::aead::Payload {
            msg: &plaintext,
            aad: AAD,
        },
    ) {
        Ok(ciphertext) => ciphertext,
        Err(_) => {
            derived_key.zeroize();
            plaintext.zeroize();
            return Err(PinError::EncryptionFailed);
        }
    };

    // Zero sensitive intermediates
    derived_key.zeroize();
    plaintext.zeroize();

    // Build output: version || salt || nonce || ciphertext (includes tag)
    let mut output =
        Vec::with_capacity(VERSION_LEN + SALT_LEN + AeadNonce::SIZE + ciphertext.len());
    output.push(VERSION_V2);
    output.extend_from_slice(&salt);
    output.extend_from_slice(nonce.as_bytes());
    output.extend_from_slice(&ciphertext);

    Ok(output)
}

/// Decrypt identity keys using a PIN.
///
/// Supports both v1 (signing + recipient) and v2 (signing + recipient + nostr) formats.
/// Returns `(signing, recipient, Option<nostr>)` — nostr is `None` for v1 blobs.
pub fn decrypt_identity_keys(
    encrypted: &[u8],
    pin: &str,
) -> Result<(SigningSecretKey, RecipientSecretKey, Option<NostrSecretKey>), PinError> {
    if pin.is_empty() {
        return Err(PinError::EmptyPin);
    }

    if encrypted.len() < MIN_ENCRYPTED_LEN {
        return Err(PinError::DataTooShort);
    }

    // Parse header
    let version = encrypted[0];
    let expected_payload_len = match version {
        VERSION_V1 => KEYS_LEN_V1,
        VERSION_V2 => KEYS_LEN_V2,
        _ => return Err(PinError::UnsupportedVersion(version)),
    };

    let salt = &encrypted[VERSION_LEN..VERSION_LEN + SALT_LEN];
    let nonce_start = VERSION_LEN + SALT_LEN;
    let nonce_end = nonce_start + AeadNonce::SIZE;
    let nonce = AeadNonce::from_bytes(&encrypted[nonce_start..nonce_end])
        .ok_or(PinError::DataTooShort)?;
    let ciphertext = &encrypted[nonce_end..];

    // Derive decryption key from PIN
    let mut derived_key = derive_key_from_pin(pin, salt)?;

    // Decrypt with AES-256-GCM
    let cipher =
        Aes256Gcm::new_from_slice(&derived_key).map_err(|_| PinError::DecryptionFailed)?;
    let gcm_nonce = aes_gcm::Nonce::from_slice(nonce.as_bytes());

    let mut plaintext = match cipher.decrypt(
        gcm_nonce,
        aes_gcm::aead::Payload {
            msg: ciphertext,
            aad: AAD,
        },
    ) {
        Ok(plaintext) => plaintext,
        Err(_) => {
            derived_key.zeroize();
            return Err(PinError::DecryptionFailed);
        }
    };

    // Zero derived key
    derived_key.zeroize();

    if plaintext.len() != expected_payload_len {
        plaintext.zeroize();
        return Err(PinError::DecryptionFailed);
    }

    // Extract keys
    let mut signing_bytes = [0u8; 32];
    let mut recipient_bytes = [0u8; 32];
    signing_bytes.copy_from_slice(&plaintext[..32]);
    recipient_bytes.copy_from_slice(&plaintext[32..64]);

    let nostr = if version == VERSION_V2 {
        let mut nostr_bytes = [0u8; 32];
        nostr_bytes.copy_from_slice(&plaintext[64..96]);
        Some(NostrSecretKey::from_bytes(nostr_bytes))
    } else {
        None
    };

    // Zero plaintext
    plaintext.zeroize();

    Ok((
        SigningSecretKey::from_bytes(signing_bytes),
        RecipientSecretKey::from_bytes(recipient_bytes),
        nostr,
    ))
}

/// Derive a 32-byte key from a PIN and salt using Argon2id.
///
/// Parameters: 64MB memory, 3 iterations, 4 lanes (~500ms on typical hardware).
fn derive_key_from_pin(pin: &str, salt: &[u8]) -> Result<[u8; 32], PinError> {
    let params = argon2::Params::new(
        64 * 1024, // 64 MB memory cost
        3,         // 3 iterations
        4,         // 4 lanes of parallelism
        Some(32),  // 32-byte output
    )
    .map_err(|e| PinError::DerivationFailed(e.to_string()))?;

    let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);

    let mut output = [0u8; 32];
    argon2
        .hash_password_into(pin.as_bytes(), salt, &mut output)
        .map_err(|e| PinError::DerivationFailed(e.to_string()))?;

    Ok(output)
}

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

    fn test_keys() -> (SigningSecretKey, RecipientSecretKey, NostrSecretKey) {
        (
            SigningSecretKey::from_bytes([0x42u8; 32]),
            RecipientSecretKey::from_bytes([0x99u8; 32]),
            NostrSecretKey::from_bytes([0xaa; 32]),
        )
    }

    #[test]
    fn roundtrip_encrypt_decrypt_v2() {
        let (signing, recipient, nostr) = test_keys();
        let pin = "1234";

        let encrypted = encrypt_identity_keys(&signing, &recipient, &nostr, pin).unwrap();
        assert_eq!(encrypted[0], VERSION_V2);

        let (dec_signing, dec_recipient, dec_nostr) =
            decrypt_identity_keys(&encrypted, pin).unwrap();

        assert_eq!(signing.as_bytes(), dec_signing.as_bytes());
        assert_eq!(recipient.as_bytes(), dec_recipient.as_bytes());
        assert!(dec_nostr.is_some());
        assert_eq!(nostr.as_bytes(), dec_nostr.unwrap().as_bytes());
    }

    #[test]
    fn backward_compat_v1_decrypt() {
        let signing = SigningSecretKey::from_bytes([0x42u8; 32]);
        let recipient = RecipientSecretKey::from_bytes([0x99u8; 32]);
        let pin = "test-pin";

        // Build v1 blob manually
        let mut salt = [0u8; SALT_LEN];
        rand::thread_rng().fill_bytes(&mut salt);
        let nonce = AeadNonce::generate();

        let derived_key = derive_key_from_pin(pin, &salt).unwrap();
        let cipher = Aes256Gcm::new_from_slice(&derived_key).unwrap();
        let gcm_nonce = aes_gcm::Nonce::from_slice(nonce.as_bytes());

        let mut plaintext = [0u8; KEYS_LEN_V1];
        plaintext[..32].copy_from_slice(signing.as_bytes());
        plaintext[32..].copy_from_slice(recipient.as_bytes());

        let ciphertext = cipher
            .encrypt(
                gcm_nonce,
                aes_gcm::aead::Payload {
                    msg: &plaintext,
                    aad: AAD,
                },
            )
            .unwrap();

        let mut v1_blob = Vec::new();
        v1_blob.push(VERSION_V1);
        v1_blob.extend_from_slice(&salt);
        v1_blob.extend_from_slice(nonce.as_bytes());
        v1_blob.extend_from_slice(&ciphertext);

        let (dec_signing, dec_recipient, dec_nostr) =
            decrypt_identity_keys(&v1_blob, pin).unwrap();

        assert_eq!(signing.as_bytes(), dec_signing.as_bytes());
        assert_eq!(recipient.as_bytes(), dec_recipient.as_bytes());
        assert!(dec_nostr.is_none());
    }

    #[test]
    fn wrong_pin_fails() {
        let (signing, recipient, nostr) = test_keys();

        let encrypted = encrypt_identity_keys(&signing, &recipient, &nostr, "correct").unwrap();
        let result = decrypt_identity_keys(&encrypted, "wrong");

        assert!(result.is_err());
        assert!(matches!(result, Err(PinError::DecryptionFailed)));
    }

    #[test]
    fn tampered_data_fails() {
        let (signing, recipient, nostr) = test_keys();

        let mut encrypted = encrypt_identity_keys(&signing, &recipient, &nostr, "pin").unwrap();
        let last = encrypted.len() - 1;
        encrypted[last] ^= 0xFF;

        let result = decrypt_identity_keys(&encrypted, "pin");
        assert!(result.is_err());
    }

    #[test]
    fn empty_pin_rejected() {
        let (signing, recipient, nostr) = test_keys();

        let result = encrypt_identity_keys(&signing, &recipient, &nostr, "");
        assert!(matches!(result, Err(PinError::EmptyPin)));

        let result = decrypt_identity_keys(&[0u8; MIN_ENCRYPTED_LEN], "");
        assert!(matches!(result, Err(PinError::EmptyPin)));
    }

    #[test]
    fn too_short_data_rejected() {
        let result = decrypt_identity_keys(&[1u8; 5], "pin");
        assert!(matches!(result, Err(PinError::DataTooShort)));
    }

    #[test]
    fn unsupported_version_rejected() {
        let (signing, recipient, nostr) = test_keys();

        let mut encrypted = encrypt_identity_keys(&signing, &recipient, &nostr, "pin").unwrap();
        encrypted[0] = 99;

        let result = decrypt_identity_keys(&encrypted, "pin");
        assert!(matches!(result, Err(PinError::UnsupportedVersion(99))));
    }

    #[test]
    fn different_encryptions_produce_different_output() {
        let (signing, recipient, nostr) = test_keys();
        let pin = "same-pin";

        let enc1 = encrypt_identity_keys(&signing, &recipient, &nostr, pin).unwrap();
        let enc2 = encrypt_identity_keys(&signing, &recipient, &nostr, pin).unwrap();

        assert_ne!(enc1, enc2);

        let (s1, r1, n1) = decrypt_identity_keys(&enc1, pin).unwrap();
        let (s2, r2, n2) = decrypt_identity_keys(&enc2, pin).unwrap();
        assert_eq!(s1.as_bytes(), s2.as_bytes());
        assert_eq!(r1.as_bytes(), r2.as_bytes());
        assert_eq!(n1.unwrap().as_bytes(), n2.unwrap().as_bytes());
    }
}