rxwing 0.1.0-draft11

pub use rkyber::fips202::{sha3_256, shake256};
pub use rkyber::fips203::sharedsecret_to_bytes;
pub use rkyber::fips203::{
    MLKEM768Params, CIPHERTEXTM768, MLKEM768DECAPSKEY, MLKEM768ENCAPSKEY, SHAREDSECRET,
};
use rkyber::mlkem_internal::{
    encaps_deterministic, encapsulation_key_from_bytes, encapsulation_key_to_bytes, keys_from_seed,
    mlkem768_decaps, COINS,
};

use rkyber::randombytes::randombytes;

use core::fmt;
use core::fmt::Debug;
use hybrid_array::{
    sizes::{U1088, U1120, U1184, U1216, U32, U6, U64, U96},
    Array, ArrayN,
};
use x25519_dalek::{PublicKey, SharedSecret, StaticSecret};

#[cfg(feature = "zeroize")]
use zeroize::{Zeroize, ZeroizeOnDrop};

const X_WING_LABEL: &[u8; 6] = br"\.//^\";

/// Size in bytes of the `EncapsulationKey`.
pub const ENCAPSULATION_KEY_SIZE: usize = 1216;
/// Size in bytes of the `DecapsulationKey`.
pub const DECAPSULATION_KEY_SIZE: usize = 32;
/// Size in bytes of the `Ciphertext`.
pub const CIPHERTEXT_SIZE: usize = 1120;
/// Number of bytes necessary to encapsulate a key
pub const ENCAPSULATION_RANDOMNESS_SIZE: usize = 64;

/// X-Wing ciphertext
#[derive(Debug, PartialEq)]
pub struct XWINGCIPHERTEXT {
    pub ct_m: ArrayN<u8, 1088>,
    pub ct_x: PublicKey,
}

impl Default for XWINGCIPHERTEXT {
    fn default() -> Self {
        XWINGCIPHERTEXT {
            ct_m: Default::default(),
            ct_x: PublicKey::from([0; 32]),
        }
    }
}

/**
Fills byte array of 1120 bytes size from input `XWINGCIPHERTEXT`
# Arguments
* `output`: output cipher text byte array (of 1120 bytes size)
* `ct`: input cipher text `XWINGCIPHERTEXT`
*/
pub fn ciphertext_to_bytes(outbytes: &mut Array<u8, U1120>, inct: &XWINGCIPHERTEXT) {
    for i in 0..CIPHERTEXT_SIZE - 32 {
        outbytes[i] = inct.ct_m[i];
    }
    for i in 0..32 {
        outbytes[CIPHERTEXT_SIZE - 32 + i] = inct.ct_x.to_bytes()[i];
    }
}

/**
Generates cipher text from input byte array of 1120 bytes size
# Arguments
* `ct`: output ciphertext (already allocated)
* `input`: input ciphertext byte array (of 1120 bytes size)
*/
pub fn ciphertext_from_bytes(outct: &mut XWINGCIPHERTEXT, inbytes: &Array<u8, U1120>) {
    let (m_bytes, x_bytes) = inbytes.split_ref::<U1088>();
    for i in 0..CIPHERTEXT_SIZE - 32 {
        outct.ct_m[i] = m_bytes[i];
    }
    outct.ct_x = PublicKey::from(x_bytes.0);
}

// The naming convention of variables matches the RFC.
// ss -> Shared Secret
// ct -> Cipher Text
// ek -> Ephemeral Key
// pk -> Public Key
// sk -> Secret Key
// Postfixes:
// _m -> ML-Kem related key
// _x -> x25519 related key

/// X-Wing encapsulation or public key.
#[derive(Clone, Debug, PartialEq)]
pub struct XWINGENCAPSKEY {
    pub pk_m: MLKEM768ENCAPSKEY,
    pub pk_x: PublicKey,
}

impl Default for XWINGENCAPSKEY {
    fn default() -> Self {
        XWINGENCAPSKEY {
            pk_m: Default::default(),
            pk_x: PublicKey::from([0; 32]),
        }
    }
}

/**
Fills byte array of 1216 bytes size from input `XWINGENCAPSKEY`
# Arguments
* `output`: output encapsulation key byte array (of 1216 bytes size)
* `pk`: input encapsulation key `XWINGENCAPSKEY`
*/
pub fn xwingencapskey_to_bytes(outbytes: &mut Array<u8, U1216>, ek: &XWINGENCAPSKEY) {
    encapsulation_key_to_bytes::<MLKEM768Params>(outbytes, &ek.pk_m);
    for i in 0..32 {
        outbytes[ENCAPSULATION_KEY_SIZE - 32 + i] = ek.pk_x.to_bytes()[i];
    }
}

/**
Generates public key from input byte array of 1216 bytes size
# Arguments
* `pk`: output `XWINGENCAPSKEY` (already allocated)
* `input`: input public key byte array (of 1216 bytes size)
*/
pub fn xwingencapskey_from_bytes(ek: &mut XWINGENCAPSKEY, inbytes: &Array<u8, U1216>) {
    let (m_bytes, x_bytes) = inbytes.split_ref::<U1184>();
    encapsulation_key_from_bytes::<MLKEM768Params>(&mut ek.pk_m, m_bytes);
    ek.pk_x = PublicKey::from(x_bytes.0);
}

/// X-Wing decapsulation key or private key.
#[derive(Clone, Default, PartialEq)]
pub struct XWINGDECAPSKEY {
    sk: [u8; DECAPSULATION_KEY_SIZE],
    pub ek: XWINGENCAPSKEY,
}

pub fn xwingdecapskey_to_bytes(outbytes: &mut [u8; DECAPSULATION_KEY_SIZE], dk: &XWINGDECAPSKEY) {
    for i in 0..DECAPSULATION_KEY_SIZE {
        outbytes[i] = dk.sk[i];
    }
}

pub fn xwingdecapskey_from_bytes(dk: &mut XWINGDECAPSKEY, inbytes: &[u8; DECAPSULATION_KEY_SIZE]) {
    let mut _ek: XWINGENCAPSKEY = Default::default();
    generate_keypair_derand(&mut _ek, dk, inbytes);
}

impl Debug for XWINGDECAPSKEY {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("DecapsulationKey")
            .field("ek", &self.ek)
            .finish_non_exhaustive()
    }
}

#[cfg(feature = "zeroize")]
impl ZeroizeOnDrop for XWINGDECAPSKEY {}

/** generate X-Wing keypair
# Arguments
* `ek`: X-Wing encapsulation key
* `dk`: X-Wing decapsulation key
*/
pub fn generate_keypair(ek: &mut XWINGENCAPSKEY, dk: &mut XWINGDECAPSKEY) {
    let mut coins: Array<u8, U32> = Default::default();
    randombytes(&mut coins, 32);
    generate_keypair_derand(ek, dk, (&coins).into());
    #[cfg(feature = "zeroize")]
    coins.zeroize();
}

pub fn generate_keypair_derand(
    ek: &mut XWINGENCAPSKEY,
    dk: &mut XWINGDECAPSKEY,
    coins: &[u8; DECAPSULATION_KEY_SIZE],
) {
    dk.sk = *coins;
    let mut dk_m: MLKEM768DECAPSKEY = Default::default();
    (ek.pk_x, _) = expand_key(&mut ek.pk_m, &mut dk_m, coins);
    #[cfg(feature = "zeroize")]
    dk_m.zeroize();
}

fn expand_key(
    pk_m: &mut MLKEM768ENCAPSKEY,
    sk_m: &mut MLKEM768DECAPSKEY,
    sk: &[u8; DECAPSULATION_KEY_SIZE],
) -> (PublicKey, StaticSecret) {
    let mut seed: Array<u8, U96> = Default::default();

    // SHAKE3-256 auf sk bis genug
    shake256(&mut seed, 96, sk, 32);
    let (mut rand_m, mut rand_x) = seed.split::<U64>();

    // generate ML-KEM keys
    keys_from_seed(pk_m, sk_m, &rand_m);

    // X25519 keys
    let sk_x = StaticSecret::from(rand_x.0);
    let pk_x = PublicKey::from(&sk_x);
    #[cfg(feature = "zeroize")]
    seed.zeroize();
    #[cfg(feature = "zeroize")]
    rand_m.zeroize();
    #[cfg(feature = "zeroize")]
    rand_x.zeroize();
    (pk_x, sk_x)
}

fn combiner(
    outhash: &mut SHAREDSECRET,
    ss_m: &SHAREDSECRET,
    ss_x: &SharedSecret,
    ct_x: &PublicKey,
    pk_x: &PublicKey,
) {
    sha3_256(
        &mut outhash.ss_hat,
        &ss_m
            .ss_hat
            .concat(Array::<u8, U32>(ss_x.to_bytes()))
            .concat(Array::<u8, U32>(ct_x.to_bytes()))
            .concat(Array::<u8, U32>(pk_x.to_bytes()))
            .concat(Array::<u8, U6>(*X_WING_LABEL)),
        134,
    );
}

/** Generates cipher text and shared secret for given encapsulation key
# Arguments
* `ss`: output shared secret (already allocated)
* `ct`: output ciphertext (already allocated)
* `pk`: input X-Wing encapsulation key
*/
pub fn encaps(ss: &mut SHAREDSECRET, ct: &mut XWINGCIPHERTEXT, pk: &XWINGENCAPSKEY) {
    let mut coins: Array<u8, U64> = Default::default();
    randombytes(&mut coins, ENCAPSULATION_RANDOMNESS_SIZE);
    encaps_derand(ss, ct, pk, &coins);
    #[cfg(feature = "zeroize")]
    coins.zeroize();
}

pub fn encaps_derand(
    ss: &mut SHAREDSECRET,
    ct: &mut XWINGCIPHERTEXT,
    pk: &XWINGENCAPSKEY,
    eseed: &Array<u8, U64>,
) {
    let (mut rand_m, mut rand_x) = eseed.split::<U32>();
    let ek_x = StaticSecret::from(rand_x.0);
    let mut coins: COINS = Default::default();
    coins.c_hat = rand_m;
    let mut ct_m: CIPHERTEXTM768 = Default::default();
    let mut ss_m: SHAREDSECRET = Default::default();
    encaps_deterministic(&mut ct_m, &mut ss_m, &pk.pk_m, &coins);
    let ct_x = PublicKey::from(&ek_x);
    let ss_x = ek_x.diffie_hellman(&pk.pk_x);
    combiner(ss, &ss_m, &ss_x, &ct_x, &pk.pk_x);
    ct.ct_m = ct_m.ct_hat;
    ct.ct_x = ct_x;
    #[cfg(feature = "zeroize")]
    rand_m.zeroize();
    #[cfg(feature = "zeroize")]
    coins.zeroize();
    #[cfg(feature = "zeroize")]
    rand_x.zeroize();
    #[cfg(feature = "zeroize")]
    ss_m.zeroize();
    //ss_x is ZeroizeOnDrop
}

/** Generates shared secret for given ciphertext and decapsulation key
# Arguments
* `ss`: output shared secret (already allocated)
* `ct`: input ciphertext
* `sk`: input X-Wing decapsulation key
*/
pub fn decaps(ss: &mut SHAREDSECRET, ct: &XWINGCIPHERTEXT, sk: &XWINGDECAPSKEY) {
    let mut pk_m: MLKEM768ENCAPSKEY = Default::default();
    let mut sk_m: MLKEM768DECAPSKEY = Default::default();
    let (pk_x, sk_x) = expand_key(&mut pk_m, &mut sk_m, (&sk.sk).into());
    let mut ct_m: CIPHERTEXTM768 = Default::default();
    ct_m.ct_hat = ct.ct_m;
    let ct_x = ct.ct_x;
    let mut ss_m: SHAREDSECRET = Default::default();
    mlkem768_decaps(&mut ss_m, &sk_m, &ct_m);
    let ss_x = sk_x.diffie_hellman(&ct_x);
    combiner(ss, &ss_m, &ss_x, &ct_x, &pk_x);
}