ecdh-omr 0.2.0

// SPDX-FileCopyrightText: 2024 eaon <eaon@posteo.net>
// SPDX-License-Identifier: EUPL-1.2

use std::marker::PhantomData;

use aead::{Aead, AeadCore, KeyInit, generic_array::typenum::marker_traits::Unsigned};
#[cfg(feature = "rustcrypto-ec")]
use elliptic_curve::{
    Curve, CurveArithmetic,
    point::PointCompression,
    sec1::{CompressedPoint, FromEncodedPoint, ModulusSize, ToEncodedPoint},
};
use rand_core::CryptoRngCore;
use sha3::{Digest, Sha3_256};

use crate::{BlindedPublicKey, Decoy, cipher_from_shared_secret, curves::KeyPair, error::*};
#[cfg(feature = "dalek")]
use crate::{
    blind_dalek,
    curves::{X25519, dalek},
    random_blind_dalek,
};
#[cfg(feature = "rustcrypto-ec")]
use crate::{
    blind_rcec,
    curves::{EllipticCurve, rcec},
    random_blind_rcec,
};

/// Semi generic implementations to create new [`Hint`]s
pub trait Hinting<K: KeyPair, const L: usize>: Sized {
    /// Create a new [`Hint`].
    fn new(
        blinded_public_key: &BlindedPublicKey<K>,
        message: &[u8; L],
        salt: &[u8],
        csprng: &mut impl CryptoRngCore,
    ) -> Result<Self>;

    /// Create a new [`Hint`] using a blinding factor secret.
    fn from_blinding_factor_secret(
        blinding_factor_secret: &K::SecretKey,
        blinded_public_key: &BlindedPublicKey<K>,
        message: &[u8; L],
        salt: &[u8],
    ) -> Result<Self>;

    /// Return the underlying [`Hint`]'s length when serialized
    // TODO If we can figure out to do math with the types provided by aead (AeadCore::NonceSize
    // etc) , this could be turned into an associated type and reused for GenericArray lengths for
    // from_bytes/to_bytes
    fn bytes_length() -> usize;

    /// Deserialize from byte slice
    fn from_bytes(bytes: &[u8]) -> Result<Self>;

    /// Serialize to byte vector
    fn to_bytes(self) -> Vec<u8>;
}

/// Message encrypted by a third party, decryptable by an anonymous recipient that doesn't know
/// whether it is addressed to them or not.
pub struct Hint<K: KeyPair, A: Aead + KeyInit, const L: usize> {
    /// Two thirds of a three part secret used to encrypt the hint's underlying message.
    ///
    /// The blinding factor of a [`BlindedPublicKey`] is blinded itself, allowing the entire hint
    /// to be randomized.
    pub(crate) blinded_blinding_factor: K::PublicKey,
    /// Ciphertext decryptable by trial decryption.
    pub(crate) ciphertext: Vec<u8>,
    /// Marker to allow us to generically use AEAD implementations
    _aead: PhantomData<A>,
}

fn encrypt<A: Aead + KeyInit>(
    shared_secret: impl AsRef<[u8]>,
    nonce: &[u8],
    message: &[u8],
) -> Result<Vec<u8>> {
    let cipher: A = cipher_from_shared_secret(shared_secret);
    let nonce_size = <A as AeadCore>::NonceSize::to_usize();
    let nonce = aead::Nonce::<A>::from_slice(&nonce[..nonce_size]);
    let ciphertext = cipher.encrypt(nonce, message)?;

    Ok(ciphertext)
}

#[cfg(feature = "dalek")]
fn construct_dalek_hint<A: Aead + KeyInit, const L: usize>(
    blinded_blinding_factor: dalek::PublicKey,
    blinding_factor_secret: &dalek::StaticSecret,
    blinded_public_key_inner: &dalek::PublicKey,
    message: &[u8; L],
    salt: &[u8],
) -> Result<Hint<X25519, A, L>> {
    let raw_shared_secret = blinding_factor_secret.diffie_hellman(blinded_public_key_inner);

    let mut hasher = <Sha3_256 as Digest>::new();
    hasher.update(raw_shared_secret.as_bytes());
    hasher.update(blinded_blinding_factor.as_bytes());
    hasher.update(salt);

    let shared_secret = hasher.finalize();

    let ciphertext = encrypt::<A>(
        shared_secret,
        blinded_blinding_factor.as_bytes(),
        message.as_ref(),
    )?;

    Ok(Hint {
        blinded_blinding_factor,
        ciphertext,
        _aead: PhantomData,
    })
}

#[cfg(feature = "dalek")]
impl<A: Aead + KeyInit, const L: usize> Hinting<X25519, L> for Hint<X25519, A, L> {
    fn new(
        blinded_public_key: &BlindedPublicKey<X25519>,
        message: &[u8; L],
        salt: &[u8],
        csprng: &mut impl CryptoRngCore,
    ) -> Result<Self> {
        let (blinded_blinding_factor, _, blinding_factor_secret) =
            random_blind_dalek(&blinded_public_key.blinding_factor, false, csprng);

        construct_dalek_hint(
            blinded_blinding_factor,
            &blinding_factor_secret,
            &blinded_public_key.inner,
            message,
            salt,
        )
    }

    fn from_blinding_factor_secret(
        blinding_factor_secret: &dalek::StaticSecret,
        blinded_public_key: &BlindedPublicKey<X25519>,
        message: &[u8; L],
        salt: &[u8],
    ) -> Result<Self> {
        let blinded_blinding_factor =
            blind_dalek(&blinded_public_key.blinding_factor, blinding_factor_secret);

        construct_dalek_hint(
            blinded_blinding_factor,
            blinding_factor_secret,
            &blinded_public_key.inner,
            message,
            salt,
        )
    }

    fn bytes_length() -> usize {
        32 + L + <A as AeadCore>::TagSize::to_usize()
    }

    fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != Self::bytes_length() {
            return Err(Error::Decoding);
        }

        let mut blinded_blinding_factor_bytes = [0u8; 32];
        blinded_blinding_factor_bytes.copy_from_slice(&bytes[0..32]);
        let blinded_blinding_factor = dalek::PublicKey::from(blinded_blinding_factor_bytes);
        let ciphertext = bytes[32..].to_vec();

        Ok(Self {
            blinded_blinding_factor,
            ciphertext,
            _aead: PhantomData,
        })
    }

    fn to_bytes(self) -> Vec<u8> {
        [
            self.blinded_blinding_factor.as_bytes(),
            self.ciphertext.as_slice(),
        ]
        .concat()
    }
}

impl<K: KeyPair, A: Aead + KeyInit, const L: usize> TryFrom<&[u8]> for Hint<K, A, L>
where
    Self: Hinting<K, L>,
{
    type Error = Error;

    fn try_from(bytes: &[u8]) -> Result<Self> {
        Self::from_bytes(bytes)
    }
}

#[cfg(feature = "rustcrypto-ec")]
fn construct_rustcrypto_ec_hint<A, C, const L: usize>(
    blinded_blinding_factor: rcec::PublicKey<C>,
    blinding_factor_secret: &rcec::SecretKey<C>,
    blinded_public_key_inner: &rcec::PublicKey<C>,
    message: &[u8; L],
    salt: &[u8],
) -> Result<Hint<EllipticCurve<C>, A, L>>
where
    A: Aead + KeyInit,
    C: CurveArithmetic + PointCompression,
    <C as Curve>::FieldBytesSize: ModulusSize,
    <C as CurveArithmetic>::AffinePoint: ToEncodedPoint<C> + FromEncodedPoint<C>,
{
    let raw_shared_secret = elliptic_curve::ecdh::diffie_hellman(
        blinding_factor_secret.to_nonzero_scalar(),
        blinded_public_key_inner.as_affine(),
    );

    let blinded_blinding_factor_cp = CompressedPoint::<C>::from(&blinded_blinding_factor);

    let mut hasher = <Sha3_256 as Digest>::new();
    hasher.update(raw_shared_secret.raw_secret_bytes());
    hasher.update(blinded_blinding_factor_cp.as_slice());
    hasher.update(salt);

    let shared_secret = hasher.finalize();

    let ciphertext = encrypt::<A>(
        shared_secret,
        blinded_blinding_factor_cp.as_slice(),
        message,
    )?;

    Ok(Hint {
        blinded_blinding_factor,
        ciphertext,
        _aead: PhantomData,
    })
}

#[cfg(feature = "rustcrypto-ec")]
impl<A: Aead + KeyInit, C: CurveArithmetic + PointCompression, const L: usize>
    Hinting<EllipticCurve<C>, L> for Hint<EllipticCurve<C>, A, L>
where
    <C as Curve>::FieldBytesSize: ModulusSize,
    <C as CurveArithmetic>::AffinePoint: ToEncodedPoint<C> + FromEncodedPoint<C>,
{
    fn new(
        blinded_public_key: &BlindedPublicKey<EllipticCurve<C>>,
        message: &[u8; L],
        salt: &[u8],
        csprng: &mut impl CryptoRngCore,
    ) -> Result<Self> {
        let (blinded_blinding_factor, _, blinding_factor_secret) =
            random_blind_rcec(&blinded_public_key.blinding_factor, false, csprng);

        construct_rustcrypto_ec_hint(
            blinded_blinding_factor,
            &blinding_factor_secret,
            &blinded_public_key.inner,
            message,
            salt,
        )
    }

    fn from_blinding_factor_secret(
        blinding_factor_secret: &rcec::SecretKey<C>,
        blinded_public_key: &BlindedPublicKey<EllipticCurve<C>>,
        message: &[u8; L],
        salt: &[u8],
    ) -> Result<Self> {
        let blinded_blinding_factor =
            blind_rcec(&blinded_public_key.blinding_factor, blinding_factor_secret);

        construct_rustcrypto_ec_hint(
            blinded_blinding_factor,
            blinding_factor_secret,
            &blinded_public_key.inner,
            message,
            salt,
        )
    }

    fn bytes_length() -> usize {
        33 + L + <A as AeadCore>::TagSize::to_usize()
    }

    fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != Self::bytes_length() {
            return Err(Error::Decoding);
        }

        let blinded_blinding_factor =
            rcec::PublicKey::<C>::from_sec1_bytes(&bytes[0..33]).map_err(|_| Error::Decoding)?;
        let ciphertext = bytes[33..].to_vec();

        Ok(Self {
            blinded_blinding_factor,
            ciphertext,
            _aead: PhantomData,
        })
    }

    fn to_bytes(self) -> Vec<u8> {
        let blinded_blinding_factor_cp = CompressedPoint::<C>::from(&self.blinded_blinding_factor);

        [
            blinded_blinding_factor_cp.as_slice(),
            self.ciphertext.as_slice(),
        ]
        .concat()
    }
}

/// Pairing of message contents and the [`BlindedPublicKey`] of its recipient.
///
/// Hint seeds are used to create [`Hints`](crate::Hints).
pub struct HintSeed<K: KeyPair, const L: usize> {
    /// Blinded Public Key
    pub blinded_public_key: BlindedPublicKey<K>,
    /// Message
    pub message: [u8; L],
}

impl<K: KeyPair, const L: usize> HintSeed<K, L> {
    /// Pair [`BlindedPublicKey`] with message contents.
    pub fn new(blinded_public_key: BlindedPublicKey<K>, message: [u8; L]) -> Self {
        Self {
            blinded_public_key,
            message,
        }
    }
}

impl<K: KeyPair, const L: usize> Decoy for HintSeed<K, L>
where
    K::PublicKey: Decoy,
{
    fn random_decoy(csprng: &mut impl CryptoRngCore) -> Self {
        Self {
            blinded_public_key: BlindedPublicKey {
                inner: K::PublicKey::random_decoy(csprng),
                blinding_factor: K::PublicKey::random_decoy(csprng),
            },
            message: [0u8; L],
        }
    }
}