kyberlib 0.0.7

use crate::rng::randombytes;
use crate::{
    params::*, poly::*, polyvec::*, symmetric::*, CryptoRng,
    KyberLibError, RngCore,
};

/// Name:  pack_pk
///
/// Description: Serialize the public key as concatenation of the
///  serialized vector of polynomials pk
///  and the public seed used to generate the matrix A.
///
/// Arguments:   [u8] r:  the output serialized public key
///  const poly *pk:  the input public-key polynomial
///  const [u8] seed: the input public seed
fn pack_pk(r: &mut [u8], pk: &mut Polyvec, seed: &[u8]) {
    const END: usize = KYBER_SYM_BYTES + KYBER_POLYVEC_BYTES;
    polyvec_tobytes(r, pk);
    r[KYBER_POLYVEC_BYTES..END]
        .copy_from_slice(&seed[..KYBER_SYM_BYTES]);
}

/// Name:  unpack_pk
///
/// Description: De-serialize public key from a byte array;
///  approximate inverse of pack_pk
///
/// Arguments:   - Polyvec pk:  output public-key vector of polynomials
///  - [u8] seed:   output seed to generate matrix A
///  - const [u8] packedpk: input serialized public key
fn unpack_pk(pk: &mut Polyvec, seed: &mut [u8], packedpk: &[u8]) {
    const END: usize = KYBER_SYM_BYTES + KYBER_POLYVEC_BYTES;
    polyvec_frombytes(pk, packedpk);
    seed[..KYBER_SYM_BYTES]
        .copy_from_slice(&packedpk[KYBER_POLYVEC_BYTES..END]);
}

/// Name:  pack_sk
///
/// Description: Serialize the secret key
///
/// Arguments: - [u8] r:  output serialized secret key
///  - const Polyvec sk: input vector of polynomials (secret key)
fn pack_sk(r: &mut [u8], sk: &mut Polyvec) {
    polyvec_tobytes(r, sk);
}

/// Name:  unpack_sk
///
/// Description: De-serialize the secret key, inverse of pack_sk
///
/// Arguments:   - Polyvec sk: output vector of polynomials (secret key)
///  - const [u8] packedsk: input serialized secret key
fn unpack_sk(sk: &mut Polyvec, packedsk: &[u8]) {
    polyvec_frombytes(sk, packedsk);
}

/// Name:  pack_ciphertext
///
/// Description: Serialize the ciphertext as concatenation of the
///  compressed and serialized vector of polynomials b
///  and the compressed and serialized polynomial v
///
/// Arguments:   [u8] r:  the output serialized ciphertext
///  const poly *pk:  the input vector of polynomials b
///  const [u8] seed: the input polynomial v
fn pack_ciphertext(r: &mut [u8], b: &mut Polyvec, v: Poly) {
    polyvec_compress(r, *b);
    poly_compress(&mut r[KYBER_POLYVEC_COMPRESSED_BYTES..], v);
}

/// Name:  unpack_ciphertext
///
/// Description: De-serialize and decompress ciphertext from a byte array;
///  approximate inverse of pack_ciphertext
///
/// Arguments:   - Polyvec b:   output vector of polynomials b
///  - poly *v:  output polynomial v
///  - const [u8] c:   input serialized ciphertext
fn unpack_ciphertext(b: &mut Polyvec, v: &mut Poly, c: &[u8]) {
    polyvec_decompress(b, c);
    poly_decompress(v, &c[KYBER_POLYVEC_COMPRESSED_BYTES..]);
}

/// Name:  rej_uniform
///
/// Description: Run rejection sampling on uniform random bytes to generate
///  uniform random integers mod q
///
/// Arguments: - i16 *r:  output buffer
///  - usize len:   requested number of 16-bit integers (uniform mod q)
///  - const [u8] buf:  input buffer (assumed to be uniform random bytes)
///  - usize buflen:  length of input buffer in bytes
///
/// Returns number of sampled 16-bit integers (at most len)
fn rej_uniform(
    r: &mut [i16],
    len: usize,
    buf: &[u8],
    buflen: usize,
) -> usize {
    let (mut ctr, mut pos) = (0usize, 0usize);
    let (mut val0, mut val1);

    while ctr < len && pos + 3 <= buflen {
        val0 = ((buf[pos]) as u16 | (buf[pos + 1] as u16) << 8) & 0xFFF;
        val1 = ((buf[pos + 1] >> 4) as u16
            | (buf[pos + 2] as u16) << 4)
            & 0xFFF;
        pos += 3;

        if val0 < KYBER_Q as u16 {
            r[ctr] = val0 as i16;
            ctr += 1;
        }
        if ctr < len && val1 < KYBER_Q as u16 {
            r[ctr] = val1 as i16;
            ctr += 1;
        }
    }
    ctr
}

fn gen_a(a: &mut [Polyvec], b: &[u8]) {
    gen_matrix(a, b, false);
}

fn gen_at(a: &mut [Polyvec], b: &[u8]) {
    gen_matrix(a, b, true);
}

/// Name:  gen_matrix
///
/// Description: Deterministically generate matrix A (or the transpose of A)
///  from a seed. Entries of the matrix are polynomials that look
///  uniformly random. Performs rejection sampling on output of
///  a XOF
///
/// Arguments:   - Polyvec a:   ouptput matrix A
///  - const [u8] seed: input seed
///  - bool transposed: boolean deciding whether A or A^T is generated
fn gen_matrix(a: &mut [Polyvec], seed: &[u8], transposed: bool) {
    let mut ctr;
    // kyberslash-guard: safe — `const` expression evaluated at
    // compile time to size `buf`. No runtime division.  ADR 0003.
    const GEN_MATRIX_NBLOCKS: usize =
        (12 * KYBER_N / 8 * (1 << 12) / KYBER_Q + XOF_BLOCKBYTES)
            / XOF_BLOCKBYTES;
    let mut buf = [0u8; GEN_MATRIX_NBLOCKS * XOF_BLOCKBYTES + 2];
    let mut buflen: usize;
    let mut off: usize;
    let mut state = XofState::new();

    for (i, item) in
        a.iter_mut().enumerate().take(KYBER_SECURITY_PARAMETER)
    {
        for j in 0..KYBER_SECURITY_PARAMETER {
            if transposed {
                xof_absorb(&mut state, seed, i as u8, j as u8);
            } else {
                xof_absorb(&mut state, seed, j as u8, i as u8);
            }
            xof_squeezeblocks(&mut buf, GEN_MATRIX_NBLOCKS, &mut state);
            buflen = GEN_MATRIX_NBLOCKS * XOF_BLOCKBYTES;
            ctr = rej_uniform(
                &mut item.vec[j].coeffs,
                KYBER_N,
                &buf,
                buflen,
            );

            while ctr < KYBER_N {
                off = buflen % 3;
                for k in 0..off {
                    buf[k] = buf[buflen - off + k];
                }
                xof_squeezeblocks(&mut buf[off..], 1, &mut state);
                buflen = off + XOF_BLOCKBYTES;
                ctr += rej_uniform(
                    &mut item.vec[j].coeffs[ctr..],
                    KYBER_N - ctr,
                    &buf,
                    buflen,
                );
            }
        }
    }
}

// Name:  indcpa_keypair
//
// Description: Generates public and private key for the CPA-secure
//  public-key encryption scheme underlying Kyber
//
// Arguments: - [u8] pk: output public key (length KYBER_INDCPA_PUBLIC_KEY_BYTES)
//  - [u8] sk: output private key (length KYBER_INDCPA_SECRET_KEY_BYTES)
pub(crate) fn indcpa_keypair<R>(
    pk: &mut [u8],
    sk: &mut [u8],
    _seed: Option<(&[u8], &[u8])>,
    _rng: &mut R,
) -> Result<(), KyberLibError>
where
    R: CryptoRng + RngCore,
{
    let mut a = [Polyvec::new(); KYBER_SECURITY_PARAMETER];
    let (mut e, mut pkpv, mut skpv) =
        (Polyvec::new(), Polyvec::new(), Polyvec::new());
    let mut nonce = 0u8;
    let mut buf = [0u8; 2 * KYBER_SYM_BYTES];
    let mut randbuf = [0u8; 2 * KYBER_SYM_BYTES];

    if let Some(s) = _seed {
        randbuf[..KYBER_SYM_BYTES].copy_from_slice(s.0);
    } else {
        randombytes(&mut randbuf, KYBER_SYM_BYTES, _rng)?;
    }

    hash_g(&mut buf, &randbuf, KYBER_SYM_BYTES);

    let (publicseed, noiseseed) = buf.split_at(KYBER_SYM_BYTES);
    gen_a(&mut a, publicseed);

    for (i, _item) in
        a.iter().enumerate().take(KYBER_SECURITY_PARAMETER)
    {
        poly_getnoise_eta1(&mut skpv.vec[i], noiseseed, nonce);
        nonce += 1;
    }
    for (i, _item) in
        a.iter().enumerate().take(KYBER_SECURITY_PARAMETER)
    {
        poly_getnoise_eta1(&mut e.vec[i], noiseseed, nonce);
        nonce += 1;
    }

    polyvec_ntt(&mut skpv);
    polyvec_ntt(&mut e);

    // matrix-vector multiplication
    for (i, _item) in
        a.iter().enumerate().take(KYBER_SECURITY_PARAMETER)
    {
        polyvec_basemul_acc_montgomery(&mut pkpv.vec[i], &a[i], &skpv);
        poly_tomont(&mut pkpv.vec[i]);
    }
    polyvec_add(&mut pkpv, &e);
    polyvec_reduce(&mut pkpv);

    pack_sk(sk, &mut skpv);
    pack_pk(pk, &mut pkpv, publicseed);
    Ok(())
}

/// Name:  indcpa_enc
///
/// Description: Encryption function of the CPA-secure
///  public-key encryption scheme underlying Kyber.
///
/// Arguments: - [u8] c:  output ciphertext (length KYBER_INDCPA_BYTES)
///  - const [u8] m:  input message (length KYBER_SYM_BYTES)
///  - const [u8] pk:   input public key (length KYBER_INDCPA_PUBLIC_KEY_BYTES)
///  - const [u8] coin: input random coins used as seed (length KYBER_SYM_BYTES)
///      to deterministically generate all randomness
pub(crate) fn indcpa_enc(
    c: &mut [u8],
    m: &[u8],
    pk: &[u8],
    coins: &[u8],
) {
    let mut at = [Polyvec::new(); KYBER_SECURITY_PARAMETER];
    let (mut sp, mut pkpv, mut ep, mut b) = (
        Polyvec::new(),
        Polyvec::new(),
        Polyvec::new(),
        Polyvec::new(),
    );
    let (mut v, mut k, mut epp) =
        (Poly::new(), Poly::new(), Poly::new());
    let mut seed = [0u8; KYBER_SYM_BYTES];
    let mut nonce = 0u8;

    unpack_pk(&mut pkpv, &mut seed, pk);
    poly_frommsg(&mut k, m);
    gen_at(&mut at, &seed);

    for (i, _item) in
        a.iter().enumerate().take(KYBER_SECURITY_PARAMETER)
    {
        poly_getnoise_eta1(&mut sp.vec[i], coins, nonce);
        nonce += 1;
    }
    for (i, _item) in
        a.iter().enumerate().take(KYBER_SECURITY_PARAMETER)
    {
        poly_getnoise_eta2(&mut ep.vec[i], coins, nonce);
        nonce += 1;
    }
    poly_getnoise_eta2(&mut epp, coins, nonce);

    polyvec_ntt(&mut sp);

    // matrix-vector multiplication
    for (i, _item) in
        a.iter().enumerate().take(KYBER_SECURITY_PARAMETER)
    {
        polyvec_basemul_acc_montgomery(&mut b.vec[i], &at[i], &sp);
    }

    polyvec_basemul_acc_montgomery(&mut v, &pkpv, &sp);
    polyvec_invntt_tomont(&mut b);
    poly_invntt_tomont(&mut v);

    polyvec_add(&mut b, &ep);
    poly_add(&mut v, &epp);
    poly_add(&mut v, &k);
    polyvec_reduce(&mut b);
    poly_reduce(&mut v);

    pack_ciphertext(c, &mut b, v);
}

/// Name:  indcpa_dec
///
/// Description: Decryption function of the CPA-secure
///  public-key encryption scheme underlying Kyber.
///
/// Arguments:   - [u8] m:  output decrypted message (of length KYBER_SYM_BYTES)
///  - const [u8] c:  input ciphertext (of length KYBER_INDCPA_BYTES)
///  - const [u8] sk: input secret key (of length KYBER_INDCPA_SECRET_KEY_BYTES)
pub(crate) fn indcpa_dec(m: &mut [u8], c: &[u8], sk: &[u8]) {
    let (mut b, mut skpv) = (Polyvec::new(), Polyvec::new());
    let (mut v, mut mp) = (Poly::new(), Poly::new());

    unpack_ciphertext(&mut b, &mut v, c);
    unpack_sk(&mut skpv, sk);

    polyvec_ntt(&mut b);
    polyvec_basemul_acc_montgomery(&mut mp, &skpv, &b);
    poly_invntt_tomont(&mut mp);

    poly_sub(&mut mp, &v);
    poly_reduce(&mut mp);

    poly_tomsg(m, mp);
}