relay-crypto 0.2.0-alpha.2

//! Tested in the integration test `e2e.rs`

use chacha20poly1305::{
    KeyInit as _, XChaCha20Poly1305, XNonce,
    aead::{Aead as _, Payload},
};
use hkdf::Hkdf;
use rand_core::{CryptoRng, RngCore};

use super::CryptoError;
use crate::{kid::kid_from_x25519_pub, record::KeyRecord};
use relay_core::prelude::CryptoMeta;

const INFO: &[u8] = b"relay-mail/e2e/private/v1";
pub const ALG: &str = "x25519-hkdf-sha256-xchacha20poly1305";

/// Encrypts a message for a recipient using X25519 key agreement, HKDF key derivation, and
/// XChaCha20-Poly1305 AEAD encryption.
pub fn encrypt(
    mut rng: impl RngCore + CryptoRng,
    recipient: &KeyRecord,
    message: &[u8],
    aad: &[u8],
) -> Result<(CryptoMeta, Vec<u8>), CryptoError> {
    if recipient.is_expired() {
        return Err(CryptoError::ExpiredKeyRecord);
    }

    let ephemeral = x25519_dalek::EphemeralSecret::random_from_rng(&mut rng);
    let ephemeral_pub = x25519_dalek::PublicKey::from(&ephemeral);
    let shared = ephemeral.diffie_hellman(&recipient.public_key());

    let hk = Hkdf::<sha2::Sha256>::new(None, shared.as_bytes());
    let mut key = [0u8; 32];
    hk.expand(INFO, &mut key)
        .map_err(|_| CryptoError::KeyDerivationError)?;

    let aead = XChaCha20Poly1305::new(&key.into());

    let mut nonce_bytes = [0u8; 24];
    rng.fill_bytes(&mut nonce_bytes);
    let nonce = XNonce::from_slice(&nonce_bytes);
    let ciphertext = aead
        .encrypt(nonce, Payload { msg: message, aad })
        .map_err(|e| CryptoError::Encrypt(e.to_string()))?;

    Ok((
        CryptoMeta {
            alg: ALG.to_string(),
            recipient: kid_from_x25519_pub(&recipient.x25519),
            sender_ephemeral: ephemeral_pub.as_bytes().to_vec(),
            nonce: nonce_bytes.to_vec(),
        },
        ciphertext,
    ))
}

/// Decrypts a message using the recipient's X25519 static secret key, along with the provided
/// metadata and additional authenticated data (AAD).
pub fn decrypt(
    secret: &x25519_dalek::StaticSecret,
    meta: &CryptoMeta,
    ciphertext: &[u8],
    aad: &[u8],
) -> Result<Vec<u8>, CryptoError> {
    if meta.alg != ALG {
        return Err(CryptoError::InvalidMeta(format!(
            "Unsupported algorithm: {}",
            meta.alg
        )));
    }

    if meta.sender_ephemeral.len() != 32 {
        return Err(CryptoError::InvalidMeta(
            "Invalid sender ephemeral key length".into(),
        ));
    }
    let sender_ephemeral_raw: &[u8; 32] = meta.sender_ephemeral.as_slice().try_into().unwrap();
    let sender_ephemeral_pub = x25519_dalek::PublicKey::from(*sender_ephemeral_raw);
    let shared = secret.diffie_hellman(&sender_ephemeral_pub);

    let hk = Hkdf::<sha2::Sha256>::new(None, shared.as_bytes());
    let mut key = [0u8; 32];
    hk.expand(INFO, &mut key)
        .map_err(|_| CryptoError::KeyDerivationError)?;

    let aead = XChaCha20Poly1305::new(&key.into());

    if meta.nonce.len() != 24 {
        return Err(CryptoError::InvalidMeta("Invalid nonce length".into()));
    }
    let nonce = XNonce::from_slice(&meta.nonce);

    let payload = aead
        .decrypt(
            nonce,
            Payload {
                msg: ciphertext,
                aad,
            },
        )
        .map_err(|e| CryptoError::Decrypt(e.to_string()))?;

    Ok(payload)
}