arcis-compiler 0.11.0

use crate::{
    core::actually_used_field::ActuallyUsedField,
    traits::{
        FromF25519,
        GetSharedRescueKey,
        MxeRescueKey,
        MxeX25519PrivateKey,
        Reveal,
        ToMontgomery,
    },
    utils::{
        crypto::{
            key::{RescueKey, X25519PrivateKey, X25519PublicKey},
            rescue_desc::{RescueArg, RescueDesc},
        },
        curve_point::Curve,
        elliptic_curve::F25519,
        matrix::Matrix,
    },
};
use std::ops::Mul;

/// The Arcis Rescue cipher. This is an MPC implementation of the Marvellous Rescue cipher,
/// see <https://tosc.iacr.org/index.php/ToSC/article/view/8695/8287>. We use it in Counter (CTR)
/// mode, see <https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38a.pdf> (Section 6.5).
#[allow(dead_code)]
pub struct RescueCipher<F: ActuallyUsedField, T: RescueArg<F>> {
    desc: RescueDesc<F, T>,
}

impl<F: ActuallyUsedField, T: RescueArg<F>> RescueCipher<F, T> {
    pub fn new(key: RescueKey<T>) -> Self {
        RescueCipher {
            desc: RescueDesc::new_cipher_desc(Matrix::from(key)),
        }
    }

    /// Given a client pubkey, load the cached shared Rescue key if the pubkey is a public circuit
    /// input and compute the Rescue key otherwise.
    /// The Rescue key is computed as follows:
    /// - perform the x25519 key exchange with the MXE private key
    /// - convert the output to `Vec<T>`
    /// - perform a key derivation, following [Section 4, Option 1.](https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-56Cr2.pdf),
    ///   with the Arcis Rescue-Prime hash function.
    pub fn new_with_client<
        Base: F25519,
        S: Clone + Copy + MxeX25519PrivateKey + Mul<C, Output = C>,
        C: Curve + ToMontgomery<Output = Base>,
    >(
        pubkey: X25519PublicKey<C>,
    ) -> Self
    where
        T: FromF25519<Base> + GetSharedRescueKey<C>,
    {
        Self::new(RescueKey::new_with_client::<Base, S, C>(pubkey))
    }

    /// Given a private key and a client public key:
    /// - perform the x25519 key exchange
    /// - convert the output to `Vec<T>`
    /// - perform a key derivation, following [Section 4, Option 1.](https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-56Cr2.pdf),
    ///   with the Arcis Rescue-Prime hash function.
    pub fn new_with_client_from_keys<
        Base: F25519,
        S: Clone + Copy + Mul<C, Output = C>,
        C: Curve + ToMontgomery<Output = Base>,
    >(
        private_key: X25519PrivateKey<S>,
        client_pubkey: X25519PublicKey<C>,
    ) -> Self
    where
        T: FromF25519<Base>,
    {
        Self::new(RescueKey::new_with_client_from_keys::<F, Base, S, C>(
            private_key,
            client_pubkey,
        ))
    }

    pub fn new_for_mxe() -> Self
    where
        T: MxeRescueKey,
    {
        Self::new(RescueKey::<T>::mxe_rescue_key())
    }

    fn get_counter(nonce: T, n_blocks: usize) -> Vec<T> {
        (0..n_blocks).fold(Vec::new(), |mut acc, i| {
            acc.extend([
                nonce,
                T::from(F::from(i as u64)),
                T::from(F::ZERO),
                T::from(F::ZERO),
                T::from(F::ZERO),
            ]);
            acc
        })
    }

    fn encrypt_batch(desc: &RescueDesc<F, T>, ptxt: &[T], counter: &[T], output: &mut Vec<T>)
    where
        T: Reveal,
    {
        let n_ptxt = ptxt.len();
        let encrypted_counter = desc.permute(&Matrix::from(counter.to_vec()));
        output.extend(
            (Matrix::from(ptxt)
                + Matrix::from(
                    encrypted_counter
                        .into_iter()
                        .take(n_ptxt)
                        .collect::<Vec<T>>(),
                ))
            .into_iter()
            .map(|c| c.reveal()),
        )
    }
    /// Encrypt the masked plaintext vector in Counter (CTR) mode.
    pub fn encrypt(&self, masked_plaintext: Vec<T>, nonce: T) -> Vec<T>
    where
        T: Reveal,
    {
        // According to https://crypto.stackexchange.com/questions/1666/can-i-safely-replace-xor-with-add-in-a-stream-cipher
        // we can mask `ptxt` by `desc.permute(counter)` via addition (as opposed to XORing)
        // without any loss of security.
        let plaintext = masked_plaintext;
        let n_blocks = plaintext.len().div_ceil(self.desc.m);
        let counter = RescueCipher::get_counter(nonce, n_blocks);

        let mut result = Vec::new();
        for i in 0..n_blocks {
            let cnt = self.desc.m * i;
            Self::encrypt_batch(
                &self.desc,
                &plaintext[cnt..(cnt + self.desc.m).min(plaintext.len())],
                &counter[cnt..cnt + self.desc.m],
                &mut result,
            )
        }
        result
    }

    fn decrypt_batch(desc: &RescueDesc<F, T>, ctxt: &[T], counter: &[T], output: &mut Vec<T>) {
        let n_ctxt = ctxt.len();
        let encrypted_counter = desc.permute(&Matrix::from(counter.to_vec()));
        output.extend(
            Matrix::from(ctxt)
                - Matrix::from(
                    encrypted_counter
                        .into_iter()
                        .take(n_ctxt)
                        .collect::<Vec<T>>(),
                ),
        )
    }

    /// Decrypt the ciphertext vector in Counter (CTR) mode.
    pub fn decrypt(&self, ciphertext: Vec<T>, nonce: T) -> Vec<T> {
        // According to https://crypto.stackexchange.com/questions/1666/can-i-safely-replace-xor-with-add-in-a-stream-cipher
        // we can unmask `ctxt` by `desc.permute(counter)` via subtraction (as opposed to
        // XORing) without any loss of security.
        let n_blocks = ciphertext.len().div_ceil(self.desc.m);
        let counter = RescueCipher::get_counter(nonce, n_blocks);

        let mut decrypted = Vec::new();
        for i in 0..n_blocks {
            let cnt = self.desc.m * i;
            Self::decrypt_batch(
                &self.desc,
                &ciphertext[cnt..(cnt + self.desc.m).min(ciphertext.len())],
                &counter[cnt..cnt + self.desc.m],
                &mut decrypted,
            );
        }
        decrypted
    }
}