relay-crypto 0.2.0-beta.4

use rand_core::{CryptoRng, RngCore};
use serde::{Serialize, de::DeserializeOwned};

use crate::record::{Key, PubRecord};
use relay_core::{
    info::{EncryptionInfo, SignatureInfo},
    prelude::{Payload, Signature},
};

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum AlgorithmError {
    #[error("Algorithm mismatch: expected {0}, got {1}")]
    AlgorithmMismatch(String, String),
    #[error("Invalid secrets: {0}")]
    InvalidSecrets(String),
    #[error("Invalid meta: {0}")]
    InvalidMeta(String),
    #[error("Signature is missing")]
    SignatureMissing,
    #[error("Key derivation failed: {0}")]
    KeyDerivation(String),
    #[error("Invalid nonce")]
    Nonce,
    #[error("Invalid ephemeral")]
    Ephemeral,
    #[error("Encryption failed: {0}")]
    Encrypt(String),
    #[error("Decryption failed: {0}")]
    Decrypt(String),
    #[error("Failed: {0}")]
    Error(String),
}

pub trait EncryptionContext {
    fn decrypt<T>(&self, payload: &Payload<T>, aad: &[u8]) -> Result<Vec<u8>, AlgorithmError>;
}

pub trait EncryptionAlgorithm {
    type Meta: Serialize + DeserializeOwned;
    type Secrets: Serialize + DeserializeOwned;
    type Context: EncryptionContext;

    fn alg() -> &'static str;
    fn info() -> Vec<u8> {
        format!("relay-{}", Self::alg()).into_bytes()
    }

    fn generate(
        &self,
        rng: impl RngCore + CryptoRng,
    ) -> Result<(Key, Self::Secrets), AlgorithmError>;

    fn open_inner(
        &self,
        meta: &Self::Meta,
        secrets: &Self::Secrets,
    ) -> Result<Self::Context, AlgorithmError>;

    fn encrypt_inner(
        &self,
        rng: impl RngCore + CryptoRng,
        recipient: &PubRecord,
        payload: &[u8],
        aad: &[u8],
    ) -> Result<(Self::Meta, Vec<u8>), AlgorithmError>;

    fn open(&self, info: &EncryptionInfo, secrets: &[u8]) -> Result<Self::Context, AlgorithmError> {
        if info.alg != Self::alg() {
            return Err(AlgorithmError::AlgorithmMismatch(
                info.alg.clone(),
                Self::alg().to_string(),
            ));
        }

        let meta: Self::Meta = serde_json::from_slice(&info.data)
            .map_err(|e| AlgorithmError::InvalidMeta(format!("Failed to parse meta: {e}")))?;
        let secrets: Self::Secrets = serde_json::from_slice(secrets)
            .map_err(|e| AlgorithmError::InvalidSecrets(format!("Failed to parse secrets: {e}")))?;
        self.open_inner(&meta, &secrets)
    }

    fn encrypt<T>(
        &self,
        rng: impl RngCore + CryptoRng,
        recipient: &PubRecord,
        payload: &[u8],
        aad: &[u8],
    ) -> Result<Payload<T>, AlgorithmError> {
        self.encrypt_inner(rng, recipient, payload, aad)
            .map(|(meta, ciphertext)| {
                let meta_bytes = serde_json::to_vec(&meta).map_err(|e| {
                    AlgorithmError::InvalidMeta(format!("Failed to serialize meta: {e}"))
                })?;
                Ok(Payload::new(
                    EncryptionInfo {
                        alg: Self::alg().to_string(),
                        data: meta_bytes,
                    },
                    None,
                    ciphertext,
                ))
            })?
    }
}

pub trait SigningAlgorithm: Send + Sync {
    type Meta: Serialize + DeserializeOwned;
    type Secrets: Serialize + DeserializeOwned;

    fn alg() -> &'static str;

    fn sign_inner(
        &self,
        payload: &[u8],
        secrets: &Self::Secrets,
    ) -> Result<(Self::Meta, Vec<u8>), AlgorithmError>;

    fn verify_inner(
        &self,
        meta: &Self::Meta,
        signature: &[u8],
        payload: &[u8],
    ) -> Result<(), AlgorithmError>;

    fn sign<T>(
        &self,
        payload: Payload<T>,
        raw_secrets: &[u8],
    ) -> Result<Payload<T>, AlgorithmError> {
        let secrets: Self::Secrets = serde_json::from_slice(raw_secrets)
            .map_err(|e| AlgorithmError::InvalidSecrets(format!("{e}")))?;

        let (meta, sig_bytes) = self.sign_inner(&payload.signing_bytes(), &secrets)?;

        let meta_bytes =
            serde_json::to_vec(&meta).map_err(|e| AlgorithmError::InvalidMeta(format!("{e}")))?;

        let mut signed_payload = payload;
        signed_payload.signature = Some(Signature {
            info: SignatureInfo {
                alg: Self::alg().to_string(),
                data: meta_bytes,
            },
            sig: sig_bytes,
        });

        Ok(signed_payload)
    }

    fn verify<T>(&self, payload: &Payload<T>) -> Result<(), AlgorithmError> {
        let signature = payload
            .signature
            .as_ref()
            .ok_or(AlgorithmError::SignatureMissing)?;

        if signature.info.alg != Self::alg() {
            return Err(AlgorithmError::AlgorithmMismatch(
                signature.info.alg.clone(),
                Self::alg().to_string(),
            ));
        }

        let meta: Self::Meta = serde_json::from_slice(&signature.info.data)
            .map_err(|e| AlgorithmError::InvalidMeta(format!("{e}")))?;

        self.verify_inner(&meta, &signature.sig, &payload.signing_bytes())
    }
}