relay-ciphers 0.2.0-beta.4

use chacha20poly1305::{
    KeyInit as _, XChaCha20Poly1305, XNonce,
    aead::{Aead as _, Payload as AeadPayload},
};
use relay_core::prelude::Payload;
use serde::{Deserialize, Serialize};
use serde_with::{base64::Base64, serde_as};
use x25519_dalek::{PublicKey, StaticSecret};

use relay_crypto::{
    alg::{AlgorithmError, EncryptionAlgorithm, EncryptionContext},
    hex,
    hkdf::Hkdf,
    kid::kid_from_key,
    rand_core::{CryptoRng, RngCore},
    record::{Key, PubRecord},
    sha2::Sha256,
};

#[serde_as]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Meta {
    #[serde_as(as = "Base64")]
    pub sender_ephemeral: Vec<u8>,
    #[serde_as(as = "Base64")]
    pub nonce: Vec<u8>,
}

#[derive(Debug, Clone)]
pub struct Context {
    key: [u8; 32],
    nonce: [u8; 24],
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Secrets {
    #[serde(with = "hex::serde")]
    pub secret: [u8; 32],
}

pub struct X25519XChaCha20Poly1305;

impl EncryptionAlgorithm for X25519XChaCha20Poly1305 {
    type Meta = Meta;
    type Secrets = Secrets;
    type Context = Context;

    fn alg() -> &'static str {
        "x25519-hkdf-sha256-xchacha20poly1305"
    }

    fn generate(
        &self,
        mut rng: impl RngCore + CryptoRng,
    ) -> Result<(Key, Self::Secrets), AlgorithmError> {
        let secret = StaticSecret::random_from_rng(&mut rng);
        let public = PublicKey::from(&secret);
        let kid = kid_from_key(public.as_bytes().to_vec());

        Ok((
            Key {
                alg: Self::alg().to_string(),
                kid,
                data: public.as_bytes().to_vec(),
            },
            Secrets {
                secret: secret.to_bytes(),
            },
        ))
    }

    fn encrypt_inner(
        &self,
        mut rng: impl RngCore + CryptoRng,
        recipient: &PubRecord,
        payload: &[u8],
        aad: &[u8],
    ) -> Result<(Self::Meta, Vec<u8>), AlgorithmError> {
        let recipient_public_key: &[u8; 32] = recipient
            .encryption
            .data
            .as_slice()
            .try_into()
            .map_err(|_| {
                AlgorithmError::InvalidMeta("recipient public key must be 32 bytes".into())
            })?;
        let recipient = PublicKey::from(*recipient_public_key);

        let ephemeral = x25519_dalek::EphemeralSecret::random_from_rng(&mut rng);
        let ephemeral_pub = PublicKey::from(&ephemeral);

        let shared = ephemeral.diffie_hellman(&recipient);

        let hkdf = Hkdf::<Sha256>::new(None, shared.as_bytes());
        let mut key = [0u8; 32];
        hkdf.expand(&Self::info(), &mut key)
            .map_err(|e| AlgorithmError::KeyDerivation(e.to_string()))?;

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

        let mut nonce = [0u8; 24];
        rng.fill_bytes(&mut nonce);

        let ciphertext = aead
            .encrypt(
                XNonce::from_slice(&nonce),
                AeadPayload { msg: payload, aad },
            )
            .map_err(|e| AlgorithmError::Encrypt(e.to_string()))?;

        Ok((
            Meta {
                sender_ephemeral: ephemeral_pub.as_bytes().to_vec(),
                nonce: nonce.to_vec(),
            },
            ciphertext,
        ))
    }

    fn open_inner(
        &self,
        meta: &Self::Meta,
        secrets: &Self::Secrets,
    ) -> Result<Self::Context, AlgorithmError> {
        let ephemeral: &[u8; 32] = meta
            .sender_ephemeral
            .as_slice()
            .try_into()
            .map_err(|_| AlgorithmError::Ephemeral)?;
        let sender = PublicKey::from(*ephemeral);

        let secret_key = StaticSecret::from(secrets.secret);
        let shared = secret_key.diffie_hellman(&sender);

        let hkdf = Hkdf::<Sha256>::new(None, shared.as_bytes());
        let mut key = [0u8; 32];
        hkdf.expand(&Self::info(), &mut key)
            .map_err(|e| AlgorithmError::KeyDerivation(e.to_string()))?;

        let nonce: [u8; 24] = meta
            .nonce
            .as_slice()
            .try_into()
            .map_err(|_| AlgorithmError::Nonce)?;

        Ok(Context { key, nonce })
    }
}

impl EncryptionContext for Context {
    fn decrypt<T>(&self, payload: &Payload<T>, aad: &[u8]) -> Result<Vec<u8>, AlgorithmError> {
        let aead = XChaCha20Poly1305::new(&self.key.into());

        aead.decrypt(
            XNonce::from_slice(&self.nonce),
            AeadPayload {
                msg: &payload.ciphertext,
                aad,
            },
        )
        .map_err(|e| AlgorithmError::Decrypt(e.to_string()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use relay_core::chrono::Utc;
    use relay_crypto::rng::os_rng_hkdf;

    #[test]
    fn encrypt_decrypt() {
        let mut rng = os_rng_hkdf(None, b"test").unwrap();
        let (key, secrets) = X25519XChaCha20Poly1305
            .generate(&mut rng)
            .expect("Generation failed");
        let record = PubRecord {
            id: "recipient".to_string(),
            created_at: Utc::now(),
            expires_at: None,
            encryption: key,
            signing: Key {
                alg: String::new(),
                kid: String::new(),
                data: vec![],
            },
        };

        let message = b"Hello, world!";
        let aad = b"Additional authenticated data";

        let payload: Payload<()> = X25519XChaCha20Poly1305
            .encrypt(rng, &record, message, aad)
            .expect("Encryption failed");

        let context = X25519XChaCha20Poly1305
            .open(&payload.info, &serde_json::to_vec(&secrets).unwrap())
            .unwrap();

        let decrypted_message = context.decrypt(&payload, aad).unwrap();
        assert_eq!(decrypted_message, message);
    }
}