arcis-compiler 0.9.1

use crate::{
    core::{
        circuits::boolean::{
            boolean_value::{Boolean, BooleanValue},
            byte::Byte,
        },
        expressions::{expr::Expr, other_expr::OtherExpr},
        global_value::value::FieldValue,
    },
    traits::{
        FromLeBits,
        FromLeBytes,
        MxeRescueKey,
        MxeX25519PrivateKey,
        Random,
        Reveal,
        ToMontgomery,
    },
    utils::{
        curve_point::{Curve, CurvePoint},
        elliptic_curve::{is_valid_x25519_public_key, AffineEdwardsPoint, F25519},
        field::{BaseField, ScalarField},
        matrix::Matrix,
    },
    MxeBitInput,
    MxeInput,
};
use std::ops::Mul;

pub const X25519_PRIVATE_KEY_COUNT: usize = 1;
pub const X25519_PUBLIC_KEY_COUNT: usize = 1;
pub const X25519_PUBLIC_KEY_BYTES_COUNT: usize = 32;
pub const RESCUE_KEY_COUNT: usize = 5;
pub const AES_128_KEY_COUNT: usize = 128;
pub const AES_192_KEY_COUNT: usize = 192;
pub const AES_256_KEY_COUNT: usize = 256;
pub const ED25519_SECRET_KEY_COUNT: usize = 256;
pub const ED25519_SIGNING_KEY_S_COUNT: usize = 1;
pub const ED25519_SIGNING_KEY_HASH_PREFIX_COUNT: usize = 256;
pub const ED25519_VERIFYING_KEY_COUNT: usize = 32;
pub const ELGAMAL_SECRET_KEY_COUNT: usize = 1;
pub const ELGAMAL_PUBKEY_COUNT: usize = 1;
pub const ELGAMAL_PUBKEY_BYTES_COUNT: usize = 32;

/// The Arcis Rescue key type. A key is an array of length 5 of secret-shared
/// elements over a finite field. The client rescue key is obtained as
/// key = rescue_prime_hash(get_shared_secret(mxe_x25519_private_key, pubkey)),
/// while the mxe rescue key corresponds to 5 randomly generated secret-shared
/// elements of the finite field.
#[derive(Copy, Clone)]
pub struct RescueKey<T: Clone + Copy + Random>([T; RESCUE_KEY_COUNT]);

impl<T: Clone + Copy + Random> RescueKey<T> {
    pub fn new_from_inner(a: [T; RESCUE_KEY_COUNT]) -> Self {
        Self(a)
    }

    pub fn inner(&self) -> [T; RESCUE_KEY_COUNT] {
        self.0
    }

    /// Loads the MXE Rescue key.
    pub fn mxe_rescue_key() -> Self
    where
        T: MxeRescueKey,
    {
        Self(
            (0..RESCUE_KEY_COUNT)
                .map(|i| T::mxe_rescue_key(i))
                .collect::<Vec<T>>()
                .try_into()
                .unwrap_or_else(|v: Vec<T>| {
                    panic!(
                        "Expected a Vec of length {} (found {})",
                        RESCUE_KEY_COUNT,
                        v.len()
                    )
                }),
        )
    }
}

impl<T: Clone + Copy + Random> Random for RescueKey<T> {
    fn random() -> Self {
        Self(
            (0..RESCUE_KEY_COUNT)
                .map(|_| T::random())
                .collect::<Vec<T>>()
                .try_into()
                .unwrap_or_else(|v: Vec<T>| {
                    panic!(
                        "Expected a Vec of length {} (found {})",
                        RESCUE_KEY_COUNT,
                        v.len()
                    )
                }),
        )
    }
}

impl<T: Clone + Copy + Random> From<RescueKey<T>> for Matrix<T> {
    fn from(key: RescueKey<T>) -> Self {
        Matrix::from(key.0.to_vec())
    }
}

macro_rules! impl_aes_key {
    ($t: ident, $byte_len:expr, $key_func_trait: ident, $key_func:ident, $mxe_key: ident) => {
        pub trait $key_func_trait {
            fn $key_func(i: usize) -> Self;
        }

        impl $key_func_trait for BooleanValue {
            fn $key_func(i: usize) -> Self {
                assert!(i < 8 * $byte_len);
                Self::from_expr(Expr::Other(OtherExpr::MxeKey(MxeInput::Bit(
                    MxeBitInput::$t(i),
                ))))
            }
        }

        impl $key_func_trait for bool {
            fn $key_func(i: usize) -> Self {
                debug_assert!(i < 8 * $byte_len);
                ($mxe_key[i / 8] >> (i % 8)) & 1u8 == 1u8
            }
        }

        /// The Arcis AES key type.
        #[derive(Clone, Copy)]
        pub struct $t<B: Boolean>([Byte<B>; $byte_len]);

        impl<B: Boolean> $t<B> {
            pub fn new_from_inner(a: [Byte<B>; $byte_len]) -> Self {
                Self(a)
            }

            pub fn inner(&self) -> [Byte<B>; $byte_len] {
                self.0
            }

            /// Loads the MXE AES key.
            pub fn mxe_aes_key() -> Self
            where
                B: $key_func_trait,
            {
                Self(
                    (0..8 * $byte_len)
                        .map(|i| B::$key_func(i))
                        .collect::<Vec<B>>()
                        .chunks(8)
                        .map(|chunk| {
                            Byte::new(chunk.to_vec().try_into().unwrap_or_else(|v: Vec<B>| {
                                panic!("Expected a Vec of length 8 (found {})", v.len())
                            }))
                        })
                        .collect::<Vec<Byte<B>>>()
                        .try_into()
                        .unwrap_or_else(|v: Vec<Byte<B>>| {
                            panic!("Expected a Vec of length {} (found {})", $byte_len, v.len())
                        }),
                )
            }
        }

        impl<B: Boolean> Random for $t<B> {
            fn random() -> Self {
                Self(
                    (0..$byte_len)
                        .map(|_| Byte::<B>::random())
                        .collect::<Vec<Byte<B>>>()
                        .try_into()
                        .unwrap_or_else(|v: Vec<Byte<B>>| {
                            panic!("Expected a Vec of length {} (found {})", $byte_len, v.len())
                        }),
                )
            }
        }

        impl $t<BooleanValue> {
            pub fn compress(&self) -> [FieldValue<BaseField>; ($byte_len as usize).div_ceil(16)] {
                self.0
                    .chunks(16)
                    .map(|bytes| {
                        FieldValue::<BaseField>::from_le_bits(
                            bytes
                                .into_iter()
                                .flat_map(|byte| byte.to_vec())
                                .collect::<Vec<BooleanValue>>(),
                            false,
                        )
                    })
                    .collect::<Vec<FieldValue<BaseField>>>()
                    .try_into()
                    .unwrap_or_else(|v: Vec<FieldValue<BaseField>>| {
                        panic!(
                            "Expected a Vec of length {} (found {})",
                            ($byte_len + 15) / 16,
                            v.len()
                        )
                    })
            }
        }
    };
}

impl_aes_key!(AES128Key, 16, MxeAES128Key, mxe_aes_128_key, MXE_AES128_KEY);
impl_aes_key!(AES192Key, 24, MxeAES192Key, mxe_aes_192_key, MXE_AES192_KEY);
impl_aes_key!(AES256Key, 32, MxeAES256Key, mxe_aes_256_key, MXE_AES256_KEY);

/// The Arcis x25519 private key type. A private key is a random element of the
/// finite field with ell = 2^252 + 27742317777372353535851937790883648493 elements.
#[derive(Copy, Clone)]
pub struct X25519PrivateKey<S: Clone + Copy> {
    key: S,
    pub is_expected_non_zero: bool,
}

impl<S: Clone + Copy> X25519PrivateKey<S> {
    pub fn new(value: S, is_expected_non_zero: bool) -> Self {
        Self {
            key: value,
            is_expected_non_zero,
        }
    }

    pub fn inner(&self) -> S {
        self.key
    }

    /// Generate a random x25519 private key.
    /// We generate private keys similar to the standard way of generating a x25519 private key:
    /// generate 32 random bytes and perform clamped multiplication with the generator/the client's
    /// pubkey. The reason for doing so is that one generally wants to avoid the expensive check if
    /// the client pubkey is a valid point of order \ell, and the clamped multiplication
    /// naturally prevents from small-subgroup attacks. Since we do check on the node that the
    /// client pubkey is a valid point of order \ell, and input it as a CurveValue, this comes
    /// down to generating 251 random bits and return \sum_{i=3}^{253} b_i*2^i + 2^254 mod \ell.
    /// Note: this is by no means better or worse than generating a random ScalarField element,
    /// it just has become a convention.
    pub fn random<B: Boolean>() -> Self
    where
        S: FromLeBits<B>,
    {
        let mut bits = vec![B::from(false); 3];
        for _ in 0..251 {
            bits.push(B::random());
        }
        bits.push(B::from(true));
        Self {
            key: S::from_le_bits(bits, false),
            is_expected_non_zero: true,
        }
    }

    /// Loads the MXE x25519 private key.
    pub fn mxe_private_key() -> Self
    where
        S: MxeX25519PrivateKey,
    {
        Self {
            key: S::mxe_x25519_private_key(),
            is_expected_non_zero: true,
        }
    }
}

impl<S: Clone + Copy + From<ScalarField>> FromLeBytes for X25519PrivateKey<S> {
    fn from_le_bytes(bytes: [u8; 32]) -> Self {
        // clamp the bytes
        let mut key_bytes = bytes;
        key_bytes[0] &= 0b1111_1000;
        key_bytes[31] &= 0b0111_1111;
        key_bytes[31] |= 0b0100_0000;
        // Note that the key cannot be zero since the clamped bytes represent a 255-bit
        // multiple of 8 and the lcm of 8 and \ell is 8 * \ell, which is 256 bits long.
        X25519PrivateKey::new(
            S::from(ScalarField::from_le_bits(
                key_bytes
                    .into_iter()
                    .flat_map(|byte| Byte::from(byte).to_vec())
                    .collect::<Vec<bool>>(),
                false,
            )),
            true,
        )
    }
}

impl Default for X25519PrivateKey<ScalarField> {
    fn default() -> Self {
        Self::from_le_bytes([0u8; 32])
    }
}

/// The Arcis x25519 public key type.
/// A public key is not necessarily public.
#[derive(Copy, Clone)]
pub struct X25519PublicKey<C: Curve> {
    point: C,
    pub is_expected_non_identity: bool,
}

impl<C: Curve> X25519PublicKey<C> {
    pub fn new(value: C, is_expected_non_identity: bool) -> Self {
        Self {
            point: value,
            is_expected_non_identity,
        }
    }

    pub fn new_from_private_key<S: Clone + Copy + Mul<C, Output = C>>(
        private_key: X25519PrivateKey<S>,
    ) -> Self {
        Self {
            point: private_key.inner() * C::generator(),
            is_expected_non_identity: private_key.is_expected_non_zero,
        }
    }

    pub fn inner(&self) -> C {
        self.point
    }

    pub fn to_montgomery<Base: F25519>(self) -> (Base, Base)
    where
        C: ToMontgomery<Output = Base>,
    {
        self.point.to_montgomery(self.is_expected_non_identity)
    }
}

impl X25519PublicKey<CurvePoint> {
    /// A X25519PublicKey is serialized by converting the inner CurvePoint to Montgomery coordinates
    /// and then serializing the u-coordinate (a BaseField element).
    pub fn to_le_bytes(&self) -> [u8; 32] {
        self.to_montgomery().0.to_le_bytes().map(u8::from)
    }

    /// Given the serialized Montgomery u-coordinate of a x25519 pubkey, checks the validity of the
    /// bytes and constructs a X25519PublicKey.
    pub fn from_le_bytes(bytes: [u8; 32]) -> Option<Self> {
        is_valid_x25519_public_key(bytes).map(|montgomery_point| {
            // convert the affine Montgomery point to a CurvePoint
            let edwards_point = AffineEdwardsPoint::from_montgomery(montgomery_point);
            Self {
                point: CurvePoint::from(edwards_point),
                is_expected_non_identity: true,
            }
        })
    }
}

impl<C: Curve> Reveal for X25519PublicKey<C> {
    fn reveal(self) -> Self {
        Self {
            point: self.point.reveal(),
            is_expected_non_identity: self.is_expected_non_identity,
        }
    }
}

impl Default for X25519PublicKey<CurvePoint> {
    fn default() -> Self {
        Self::new_from_private_key(X25519PrivateKey::<ScalarField>::default())
    }
}

// Randomly generated element of the scalar field.
pub(crate) const MXE_X25519_PRIVATE_KEY: [u8; 32] = [
    207, 40, 181, 230, 45, 204, 46, 17, 8, 19, 251, 241, 43, 129, 216, 23, 86, 169, 218, 248, 95,
    114, 111, 9, 188, 159, 223, 16, 124, 98, 41, 1,
];

// Five randomly generated elements of the base field.
pub(crate) const MXE_RESCUE_BASE_FIELD_KEY: [[u8; 32]; 5] = [
    [
        124, 219, 118, 100, 151, 174, 173, 201, 180, 159, 95, 202, 109, 154, 90, 104, 99, 221, 206,
        79, 44, 221, 182, 198, 143, 180, 180, 121, 78, 223, 238, 12,
    ],
    [
        158, 10, 32, 122, 234, 85, 113, 2, 69, 115, 151, 149, 163, 189, 216, 108, 160, 21, 118,
        154, 185, 199, 198, 251, 142, 193, 168, 98, 218, 59, 20, 9,
    ],
    [
        45, 115, 67, 23, 196, 46, 150, 202, 33, 22, 44, 144, 204, 34, 166, 30, 183, 63, 38, 213,
        166, 150, 234, 191, 201, 13, 79, 86, 171, 100, 140, 15,
    ],
    [
        209, 1, 108, 251, 175, 105, 199, 246, 83, 186, 72, 0, 15, 236, 105, 110, 5, 109, 41, 216,
        148, 98, 208, 128, 32, 47, 224, 93, 90, 176, 33, 2,
    ],
    [
        179, 142, 132, 221, 113, 147, 206, 83, 22, 121, 245, 155, 239, 204, 18, 158, 119, 190, 54,
        17, 28, 17, 247, 191, 151, 147, 118, 151, 38, 169, 21, 4,
    ],
];

// Five randomly generated elements of the scalar field.
pub(crate) const MXE_RESCUE_SCALAR_FIELD_KEY: [[u8; 32]; 5] = [
    [
        72, 239, 104, 144, 184, 215, 25, 76, 171, 209, 107, 34, 146, 39, 35, 134, 103, 10, 38, 148,
        251, 155, 69, 207, 216, 234, 104, 220, 253, 63, 217, 14,
    ],
    [
        11, 232, 48, 222, 61, 142, 29, 189, 73, 204, 120, 1, 111, 20, 233, 1, 101, 49, 169, 220,
        66, 250, 21, 125, 132, 247, 11, 21, 128, 217, 231, 12,
    ],
    [
        33, 164, 174, 155, 214, 77, 209, 93, 24, 230, 196, 38, 29, 5, 89, 76, 52, 85, 219, 66, 90,
        7, 190, 27, 181, 51, 206, 187, 99, 72, 251, 2,
    ],
    [
        7, 178, 214, 42, 8, 10, 134, 249, 244, 69, 192, 63, 96, 204, 229, 10, 126, 3, 33, 49, 174,
        26, 149, 5, 254, 118, 230, 75, 149, 105, 145, 5,
    ],
    [
        22, 94, 36, 119, 174, 65, 70, 75, 230, 11, 102, 41, 31, 194, 55, 255, 24, 25, 113, 181,
        206, 213, 243, 203, 115, 152, 90, 177, 232, 228, 145, 3,
    ],
];

// 16 randomly generated bytes.
pub(crate) const MXE_AES128_KEY: [u8; 16] = [
    124, 219, 118, 100, 151, 174, 173, 201, 180, 159, 95, 202, 109, 154, 90, 104,
];

// 24 randomly generated bytes.
pub(crate) const MXE_AES192_KEY: [u8; 24] = [
    159, 234, 203, 164, 191, 151, 2, 150, 165, 97, 217, 48, 4, 227, 91, 13, 244, 2, 222, 234, 226,
    187, 253, 127,
];

// 32 randomly generated bytes.
pub(crate) const MXE_AES256_KEY: [u8; 32] = [
    196, 120, 191, 34, 144, 89, 36, 185, 242, 159, 162, 158, 170, 37, 234, 191, 33, 141, 52, 147,
    253, 3, 58, 49, 169, 146, 53, 185, 1, 55, 60, 87,
];

// 32 randomly generated bytes.
pub(crate) const MXE_ED25519_SECRET_KEY: [u8; 32] = [
    13, 192, 52, 244, 242, 224, 239, 161, 194, 4, 155, 74, 60, 148, 207, 43, 188, 194, 251, 154,
    142, 189, 227, 93, 25, 249, 65, 205, 121, 53, 203, 22,
];

// Signing key s derived from MXE_ED25519_SECRET_KEY.
pub(crate) const MXE_ED25519_SIGNING_KEY_S: [u8; 32] = [
    60, 248, 182, 157, 91, 202, 9, 247, 187, 176, 179, 140, 115, 51, 98, 253, 44, 82, 56, 212, 19,
    155, 154, 68, 254, 17, 132, 223, 74, 130, 165, 12,
];

// Signing key hash-prefix derived from MXE_ED25519_SECRET_KEY.
pub(crate) const MXE_ED25519_SIGNING_KEY_HASH_PREFIX: [u8; 32] = [
    217, 141, 36, 192, 41, 135, 248, 83, 212, 208, 84, 27, 188, 50, 69, 245, 68, 207, 121, 207, 6,
    115, 30, 27, 94, 187, 91, 113, 18, 219, 27, 64,
];

// Verifying key derived from MXE_ED25519_SECRET_KEY.
pub(crate) const MXE_ED25519_VERIFYING_KEY: [u8; 32] = [
    14, 32, 126, 133, 202, 216, 28, 182, 228, 95, 141, 94, 34, 155, 36, 238, 53, 102, 82, 241, 81,
    180, 136, 239, 169, 22, 226, 133, 195, 171, 252, 7,
];

// Randomly generated element of the scalar field.
pub(crate) const MXE_ELGAMAL_SECRET_KEY: [u8; 32] = [
    57, 71, 187, 92, 7, 239, 222, 40, 252, 52, 135, 191, 198, 199, 173, 77, 31, 119, 254, 244, 152,
    71, 131, 255, 183, 183, 182, 191, 87, 99, 182, 3,
];

// Public key derived from MXE_ELGAMAL_SECRET_KEY, stored as CompressedRistretto.
pub(crate) const MXE_ELGAMAL_PUBKEY: [u8; 32] = [
    168, 163, 137, 87, 0, 80, 192, 23, 16, 133, 172, 238, 187, 164, 133, 148, 115, 126, 183, 16,
    220, 95, 125, 82, 16, 175, 99, 114, 226, 131, 197, 111,
];