/* kyber.c - the Kyber key encapsulation mechanism (main part)

 * Copyright (C) 2024 g10 Code GmbH

 *

 * This file was modified for use by Libgcrypt.

 *

 * This file is free software; you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as

 * published by the Free Software Foundation; either version 2.1 of

 * the License, or (at your option) any later version.

 *

 * This file is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public

 * License along with this program; if not, see <https://www.gnu.org/licenses/>.

 * SPDX-License-Identifier: LGPL-2.1-or-later

 *

 * You can also use this file under the same licence of original code.

 * SPDX-License-Identifier: CC0 OR Apache-2.0

 *

 */

/*

  Original code from:

  Repository: https://github.com/pq-crystals/kyber.git

  Branch: standard

  Commit: 11d00ff1f20cfca1f72d819e5a45165c1e0a2816

  Licence:

  Public Domain (https://creativecommons.org/share-your-work/public-domain/cc0/);

  or Apache 2.0 License (https://www.apache.org/licenses/LICENSE-2.0.html).

  Authors:

        Joppe Bos

        Léo Ducas

        Eike Kiltz

        Tancrède Lepoint

        Vadim Lyubashevsky

        John Schanck

        Peter Schwabe

        Gregor Seiler

        Damien Stehlé

  Kyber Home: https://www.pq-crystals.org/kyber/

 */

/*

 * This implementation consists of four files: kyber.h (header),

 * kyber.c (this), kyber-common.c (common part), and kyber-kdep.c

 * (KYBER_K dependent part).

 *

 * It is for inclusion in libgcrypt library.  Also, standalone use of

 * the implementation is possible.  With KYBER_K defined, it can offer

 * the variant of that KYBER_K specified.  Otherwise, three variants

 * are offered.

 *

 * From original code, following modification was made.

 *

 * - C++ style comments are changed to C-style.

 *

 * - No use of KYBER_NAMESPACE and FIPS202_NAMESPACE.  Don't export

 *   internal symbols.

 *

 * - "verify" routine is changed to return 1 on success, and now has

 *   new name "verify1", so that the use of the routine won't need

 *   negation (since negation might result non-constant-time code with

 *   branch by some compiler).

 *

 * - For "xof" routines, definitions of xof_init and xof_close are

 *   added, so that memory will be possible to be cleared after its

 *   use.

 *

 * - Different external API for shake128, having _init and _close.

 *

 * - New implementation of kyber_shake128_absorb, with the shake128

 *   API.

 *

 * - Added an external function: shake256v with variable arguments.

 *

 * - Macro definitions of xof_squeezeblocks, prf, and rkprf are

 *   modified to use the shake128 API and the shake256v function.

 *

 */

#ifdef HAVE_CONFIG_H

#include <config.h>

#endif

#include <stddef.h>

#include <stdint.h>

#include <string.h>

#ifdef _GCRYPT_IN_LIBGCRYPT

#include <stdarg.h>

#include <gpg-error.h>

#include "types.h"

#include "g10lib.h"

#include "gcrypt-int.h"

#include "const-time.h"

#include "kyber.h"

static int crypto_kem_keypair_derand_2(uint8_t *pk, uint8_t *sk,

                                       const uint8_t *coins);

static int crypto_kem_keypair_derand_3(uint8_t *pk, uint8_t *sk,

                                       const uint8_t *coins);

static int crypto_kem_keypair_derand_4(uint8_t *pk, uint8_t *sk,

                                       const uint8_t *coins);

static int crypto_kem_keypair_2(uint8_t *pk, uint8_t *sk);

static int crypto_kem_keypair_3(uint8_t *pk, uint8_t *sk);

static int crypto_kem_keypair_4(uint8_t *pk, uint8_t *sk);

static int crypto_kem_enc_derand_2(uint8_t *ct, uint8_t *ss, const uint8_t *pk,

                                   const uint8_t *coins);

static int crypto_kem_enc_derand_3(uint8_t *ct, uint8_t *ss, const uint8_t *pk,

                                   const uint8_t *coins);

static int crypto_kem_enc_derand_4(uint8_t *ct, uint8_t *ss, const uint8_t *pk,

                                   const uint8_t *coins);

static int crypto_kem_enc_2(uint8_t *ct, uint8_t *ss, const uint8_t *pk);

static int crypto_kem_enc_3(uint8_t *ct, uint8_t *ss, const uint8_t *pk);

static int crypto_kem_enc_4(uint8_t *ct, uint8_t *ss, const uint8_t *pk);

static int crypto_kem_dec_2(uint8_t *ss, const uint8_t *ct, const uint8_t *sk);

static int crypto_kem_dec_3(uint8_t *ss, const uint8_t *ct, const uint8_t *sk);

static int crypto_kem_dec_4(uint8_t *ss, const uint8_t *ct, const uint8_t *sk);

void

kyber_keypair (int algo, uint8_t *pk, uint8_t *sk, const uint8_t *coins)

{

  switch (algo)

    {

    case GCRY_KEM_MLKEM512:

      if (coins)

        crypto_kem_keypair_derand_2 (pk, sk, coins);

      else

        crypto_kem_keypair_2 (pk, sk);

      break;

    case GCRY_KEM_MLKEM768:

    default:

      if (coins)

        crypto_kem_keypair_derand_3 (pk, sk, coins);

      else

        crypto_kem_keypair_3 (pk, sk);

      break;

    case GCRY_KEM_MLKEM1024:

      if (coins)

        crypto_kem_keypair_derand_4 (pk, sk, coins);

      else

        crypto_kem_keypair_4 (pk, sk);

      break;

    }

}

void

kyber_encap (int algo, uint8_t *ct, uint8_t *ss, const uint8_t *pk,

             const uint8_t *coins)

{

  switch (algo)

    {

    case GCRY_KEM_MLKEM512:

      if (coins)

        crypto_kem_enc_derand_2 (ct, ss, pk, coins);

      else

        crypto_kem_enc_2 (ct, ss, pk);

      break;

    case GCRY_KEM_MLKEM768:

    default:

      if (coins)

        crypto_kem_enc_derand_3 (ct, ss, pk, coins);

      else

        crypto_kem_enc_3 (ct, ss, pk);

      break;

    case GCRY_KEM_MLKEM1024:

      if (coins)

        crypto_kem_enc_derand_4 (ct, ss, pk, coins);

      else

        crypto_kem_enc_4 (ct, ss, pk);

      break;

    }

}

void

kyber_decap (int algo, uint8_t *ss, const uint8_t *ct, const uint8_t *sk)

{

  switch (algo)

    {

    case GCRY_KEM_MLKEM512:

      crypto_kem_dec_2 (ss, ct, sk);

      break;

    case GCRY_KEM_MLKEM768:

    default:

      crypto_kem_dec_3 (ss, ct, sk);

      break;

    case GCRY_KEM_MLKEM1024:

      crypto_kem_dec_4 (ss, ct, sk);

      break;

    }

}

static void

randombytes (uint8_t *out, size_t outlen)

{

  _gcry_randomize (out, outlen, GCRY_VERY_STRONG_RANDOM);

}

typedef struct {

  gcry_md_hd_t h;

} keccak_state;

static void

shake128_init (keccak_state *state)

{

  gcry_err_code_t ec;

  ec = _gcry_md_open (&state->h, GCRY_MD_SHAKE128, 0);

  if (ec)

    log_fatal ("internal md_open failed: %d\n", ec);

}

static void

shake128_absorb (keccak_state *state, const uint8_t *in, size_t inlen)

{

  _gcry_md_write (state->h, in, inlen);

}

static void

shake128_finalize (keccak_state *state)

{

  (void)state;

}

static void

shake128_squeeze (keccak_state *state, uint8_t *out, size_t outlen)

{

  _gcry_md_extract (state->h, GCRY_MD_SHAKE128, out, outlen);

}

static void

shake128_close (keccak_state *state)

{

  _gcry_md_close (state->h);

}

#define MAX_ARGS 16

static void

shake256v (uint8_t *out, size_t outlen, ...)

{

  gcry_buffer_t iov[MAX_ARGS];

  va_list ap;

  int i;

  void *p;

  size_t len;

  va_start (ap, outlen);

  for (i = 0; i < MAX_ARGS; i++)

    {

      p = va_arg (ap, void *);

      len = va_arg (ap, size_t);

      if (!p)

        break;

      iov[i].size = 0;

      iov[i].data = p;

      iov[i].off = 0;

      iov[i].len = len;

    }

  va_end (ap);

  _gcry_md_hash_buffers_extract (GCRY_MD_SHAKE256, 0, out, outlen,

                                 iov, i);

}

static void

sha3_256 (uint8_t h[32], const uint8_t *in, size_t inlen)

{

  _gcry_md_hash_buffer (GCRY_MD_SHA3_256, h, in, inlen);

}

static void

sha3_512 (uint8_t h[64], const uint8_t *in, size_t inlen)

{

  _gcry_md_hash_buffer (GCRY_MD_SHA3_512, h, in, inlen);

}

#define verify1 ct_memequal

#define cmov    ct_memmov_cond

#else

#include "kyber.h"

void randombytes (uint8_t *out, size_t outlen);

typedef struct {

  uint64_t s[25];

  unsigned int pos;

} keccak_state;

void shake128_init (keccak_state *state);

void shake128_absorb (keccak_state *state, const uint8_t *in, size_t inlen);

void shake128_finalize (keccak_state *state);

void shake128_squeeze (keccak_state *state, uint8_t *out, size_t outlen);

void shake128_close (keccak_state *state);

void shake256v (uint8_t *out, size_t outlen, ...);

void sha3_256 (uint8_t h[32], const uint8_t *in, size_t inlen);

void sha3_512 (uint8_t h[64], const uint8_t *in, size_t inlen);

/* Return 1 when success, 0 otherwise.  */

unsigned int verify1 (const uint8_t *a, const uint8_t *b, size_t len);

/* Conditional move.  */

void cmov (uint8_t *r, const uint8_t *x, size_t len, uint8_t b);

/* Select V0 when OP_ENABLE == 1, V1 otherwise.  */

int16_t ct_int16_select (int16_t v0, int16_t v1, unsigned long op_enable);

#endif

/*************** kyber/ref/fips202.h */

#define SHAKE128_RATE 168

/*************** kyber/ref/params.h */

#define KYBER_N 256

#define KYBER_Q 3329

#define KYBER_SYMBYTES 32   /* size in bytes of hashes, and seeds */

#define KYBER_SSBYTES  32   /* size in bytes of shared key */

#define KYBER_POLYBYTES          384

#define KYBER_ETA2 2

#define KYBER_INDCPA_MSGBYTES       (KYBER_SYMBYTES)

/* KYBER_K dependent values (part 1) */

#define KYBER_ETA1_2   3

#define KYBER_ETA1_3_4 2

#define KYBER_POLYCOMPRESSEDBYTES_2_3 128

#define KYBER_POLYCOMPRESSEDBYTES_4   160

/*************** kyber/ref/poly.h */

/*

 * Elements of R_q = Z_q[X]/(X^n + 1). Represents polynomial

 * coeffs[0] + X*coeffs[1] + X^2*coeffs[2] + ... + X^{n-1}*coeffs[n-1]

 */

typedef struct{

  int16_t coeffs[KYBER_N];

} poly;

#if !defined(KYBER_K) || KYBER_K == 2 || KYBER_K == 3

static void poly_compress_128(uint8_t r[KYBER_POLYCOMPRESSEDBYTES_2_3], const poly *a);

static void poly_decompress_128(poly *r, const uint8_t a[KYBER_POLYCOMPRESSEDBYTES_2_3]);

#endif

#if !defined(KYBER_K) || KYBER_K == 4

static void poly_compress_160(uint8_t r[KYBER_POLYCOMPRESSEDBYTES_4], const poly *a);

static void poly_decompress_160(poly *r, const uint8_t a[KYBER_POLYCOMPRESSEDBYTES_4]);

#endif

static void poly_tobytes(uint8_t r[KYBER_POLYBYTES], const poly *a);

static void poly_frombytes(poly *r, const uint8_t a[KYBER_POLYBYTES]);

static void poly_frommsg(poly *r, const uint8_t msg[KYBER_INDCPA_MSGBYTES]);

static void poly_tomsg(uint8_t msg[KYBER_INDCPA_MSGBYTES], const poly *r);

#if !defined(KYBER_K) || KYBER_K == 2

static void poly_getnoise_eta1_2(poly *r, const uint8_t seed[KYBER_SYMBYTES], uint8_t nonce);

#endif

#if !defined(KYBER_K) || KYBER_K == 3 || KYBER_K == 4

static void poly_getnoise_eta1_3_4(poly *r, const uint8_t seed[KYBER_SYMBYTES], uint8_t nonce);

#endif

static void poly_getnoise_eta2(poly *r, const uint8_t seed[KYBER_SYMBYTES], uint8_t nonce);

static void poly_ntt(poly *r);

static void poly_invntt_tomont(poly *r);

static void poly_basemul_montgomery(poly *r, const poly *a, const poly *b);

static void poly_tomont(poly *r);

static void poly_reduce(poly *r);

static void poly_add(poly *r, const poly *a, const poly *b);

static void poly_sub(poly *r, const poly *a, const poly *b);

/*************** kyber/ref/ntt.h */

static const int16_t zetas[128];

static void ntt(int16_t poly[256]);

static void invntt(int16_t poly[256]);

static void basemul(int16_t r[2], const int16_t a[2], const int16_t b[2], int16_t zeta);

/*************** kyber/ref/reduce.h */

#define MONT -1044 /* 2^16 mod q */

#define QINV -3327 /* q^-1 mod 2^16 */

static int16_t montgomery_reduce(int32_t a);

static int16_t barrett_reduce(int16_t a);

/*************** kyber/ref/symmetric.h */

typedef keccak_state xof_state;

static void kyber_shake128_absorb (keccak_state *state,

                                   const uint8_t seed[KYBER_SYMBYTES],

                                   uint8_t x, uint8_t y)

{

  shake128_absorb (state, seed, KYBER_SYMBYTES);

  shake128_absorb (state, &x, 1);

  shake128_absorb (state, &y, 1);

  shake128_finalize (state);

}

#define XOF_BLOCKBYTES SHAKE128_RATE

#define hash_h(OUT, IN, INBYTES) sha3_256(OUT, IN, INBYTES)

#define hash_g(OUT, IN, INBYTES) sha3_512(OUT, IN, INBYTES)

#define xof_init(STATE) shake128_init(STATE)

#define xof_close(STATE) shake128_close(STATE)

#define xof_absorb(STATE, SEED, X, Y) kyber_shake128_absorb(STATE, SEED, X, Y)

#define xof_squeezeblocks(OUT, OUTBLOCKS, STATE) shake128_squeeze(STATE, OUT, SHAKE128_RATE * OUTBLOCKS)

#define prf(OUT, OUTBYTES, KEY, NONCE) \

  shake256v(OUT, OUTBYTES, (void *)(KEY), (size_t)KYBER_SYMBYTES, \

         (void *)&(NONCE), (size_t)1, \

         NULL, (size_t)0)

#define rkprf(OUT, KEY, INPUT) \

  shake256v(OUT, KYBER_SSBYTES, (void *)(KEY), (size_t)KYBER_SYMBYTES, \

        (void *)(INPUT), (size_t)KYBER_CIPHERTEXTBYTES, \

        NULL, (size_t)0)

#include "kyber-common.c"

#define VARIANT2(name) name ## _2

#define VARIANT3(name) name ## _3

#define VARIANT4(name) name ## _4

/* KYBER_K dependent values (part 2) */

#define KYBER_POLYVECBYTES      (KYBER_K * KYBER_POLYBYTES)

#define KYBER_INDCPA_PUBLICKEYBYTES (KYBER_POLYVECBYTES + KYBER_SYMBYTES)

#define KYBER_INDCPA_SECRETKEYBYTES (KYBER_POLYVECBYTES)

#define KYBER_INDCPA_BYTES          (KYBER_POLYVECCOMPRESSEDBYTES + KYBER_POLYCOMPRESSEDBYTES)

#define KYBER_PUBLICKEYBYTES  (KYBER_INDCPA_PUBLICKEYBYTES)

/* 32 bytes of additional space to save H(pk) */

#define KYBER_SECRETKEYBYTES  (KYBER_INDCPA_SECRETKEYBYTES + KYBER_INDCPA_PUBLICKEYBYTES + 2*KYBER_SYMBYTES)

#define KYBER_CIPHERTEXTBYTES (KYBER_INDCPA_BYTES)

#ifdef KYBER_K

# if KYBER_K == 2

#  define KYBER_POLYCOMPRESSEDBYTES    128

#  define KYBER_POLYVECCOMPRESSEDBYTES (KYBER_K * 320)

#  define poly_compress poly_compress_128

#  define poly_decompress poly_decompress_128

#  define poly_getnoise_eta1 poly_getnoise_eta1_2

# elif KYBER_K == 3

#  define KYBER_POLYCOMPRESSEDBYTES    128

#  define KYBER_POLYVECCOMPRESSEDBYTES (KYBER_K * 320)

#  define poly_compress poly_compress_128

#  define poly_decompress poly_decompress_128

#  define poly_getnoise_eta1 poly_getnoise_eta1_3_4

# elif KYBER_K == 4

#  define KYBER_POLYCOMPRESSEDBYTES    160

#  define KYBER_POLYVECCOMPRESSEDBYTES (KYBER_K * 352)

#  define poly_compress poly_compress_160

#  define poly_decompress poly_decompress_160

#  define poly_getnoise_eta1 poly_getnoise_eta1_3_4

# endif

# include "kyber-kdep.c"

# else

# define KYBER_K 2

# define KYBER_POLYCOMPRESSEDBYTES    128

# define KYBER_POLYVECCOMPRESSEDBYTES (KYBER_K * 320)

# define poly_compress poly_compress_128

# define poly_decompress poly_decompress_128

# define poly_getnoise_eta1 poly_getnoise_eta1_2

# define crypto_kem_keypair_derand VARIANT2(crypto_kem_keypair_derand)

# define crypto_kem_enc_derand VARIANT2(crypto_kem_enc_derand)

# define crypto_kem_keypair VARIANT2(crypto_kem_keypair)

# define crypto_kem_enc VARIANT2(crypto_kem_enc)

# define crypto_kem_dec VARIANT2(crypto_kem_dec)

# define polyvec VARIANT2(polyvec)

# define polyvec_compress VARIANT2(polyvec_compress)

# define polyvec_decompress VARIANT2(polyvec_decompress)

# define polyvec_tobytes VARIANT2(polyvec_tobytes)

# define polyvec_frombytes VARIANT2(polyvec_frombytes)

# define polyvec_ntt VARIANT2(polyvec_ntt)

# define polyvec_invntt_tomont VARIANT2(polyvec_invntt_tomont)

# define polyvec_basemul_acc_montgomery VARIANT2(polyvec_basemul_acc_montgomery)

# define polyvec_reduce VARIANT2(polyvec_reduce)

# define polyvec_add VARIANT2(polyvec_add)

# define pack_pk VARIANT2(pack_pk)

# define unpack_pk VARIANT2(unpack_pk)

# define pack_sk VARIANT2(pack_sk)

# define unpack_sk VARIANT2(unpack_sk)

# define pack_ciphertext VARIANT2(pack_ciphertext)

# define unpack_ciphertext VARIANT2(unpack_ciphertext)

# define gen_matrix VARIANT2(gen_matrix)

# define indcpa_keypair_derand VARIANT2(indcpa_keypair_derand)

# define indcpa_enc VARIANT2(indcpa_enc)

# define indcpa_dec VARIANT2(indcpa_dec)

# include "kyber-kdep.c"

# define KYBER_K 3

# define KYBER_POLYCOMPRESSEDBYTES    128

# define KYBER_POLYVECCOMPRESSEDBYTES (KYBER_K * 320)

# define poly_compress poly_compress_128

# define poly_decompress poly_decompress_128

# define poly_getnoise_eta1 poly_getnoise_eta1_3_4

# define crypto_kem_keypair_derand VARIANT3(crypto_kem_keypair_derand)

# define crypto_kem_enc_derand VARIANT3(crypto_kem_enc_derand)

# define crypto_kem_keypair VARIANT3(crypto_kem_keypair)

# define crypto_kem_enc VARIANT3(crypto_kem_enc)

# define crypto_kem_dec VARIANT3(crypto_kem_dec)

# define polyvec VARIANT3(polyvec)

# define polyvec_compress VARIANT3(polyvec_compress)

# define polyvec_decompress VARIANT3(polyvec_decompress)

# define polyvec_tobytes VARIANT3(polyvec_tobytes)

# define polyvec_frombytes VARIANT3(polyvec_frombytes)

# define polyvec_ntt VARIANT3(polyvec_ntt)

# define polyvec_invntt_tomont VARIANT3(polyvec_invntt_tomont)

# define polyvec_basemul_acc_montgomery VARIANT3(polyvec_basemul_acc_montgomery)

# define polyvec_reduce VARIANT3(polyvec_reduce)

# define polyvec_add VARIANT3(polyvec_add)

# define pack_pk VARIANT3(pack_pk)

# define unpack_pk VARIANT3(unpack_pk)

# define pack_sk VARIANT3(pack_sk)

# define unpack_sk VARIANT3(unpack_sk)

# define pack_ciphertext VARIANT3(pack_ciphertext)

# define unpack_ciphertext VARIANT3(unpack_ciphertext)

# define gen_matrix VARIANT3(gen_matrix)

# define indcpa_keypair_derand VARIANT3(indcpa_keypair_derand)

# define indcpa_enc VARIANT3(indcpa_enc)

# define indcpa_dec VARIANT3(indcpa_dec)

# include "kyber-kdep.c"

# define KYBER_K 4

# define KYBER_POLYCOMPRESSEDBYTES    160

# define KYBER_POLYVECCOMPRESSEDBYTES (KYBER_K * 352)

# define poly_compress poly_compress_160

# define poly_decompress poly_decompress_160

# define poly_getnoise_eta1 poly_getnoise_eta1_3_4

# define crypto_kem_keypair_derand VARIANT4(crypto_kem_keypair_derand)

# define crypto_kem_enc_derand VARIANT4(crypto_kem_enc_derand)

# define crypto_kem_keypair VARIANT4(crypto_kem_keypair)

# define crypto_kem_enc VARIANT4(crypto_kem_enc)

# define crypto_kem_dec VARIANT4(crypto_kem_dec)

# define polyvec VARIANT4(polyvec)

# define polyvec_compress VARIANT4(polyvec_compress)

# define polyvec_decompress VARIANT4(polyvec_decompress)

# define polyvec_tobytes VARIANT4(polyvec_tobytes)

# define polyvec_frombytes VARIANT4(polyvec_frombytes)

# define polyvec_ntt VARIANT4(polyvec_ntt)

# define polyvec_invntt_tomont VARIANT4(polyvec_invntt_tomont)

# define polyvec_basemul_acc_montgomery VARIANT4(polyvec_basemul_acc_montgomery)

# define polyvec_reduce VARIANT4(polyvec_reduce)

# define polyvec_add VARIANT4(polyvec_add)

# define pack_pk VARIANT4(pack_pk)

# define unpack_pk VARIANT4(unpack_pk)

# define pack_sk VARIANT4(pack_sk)

# define unpack_sk VARIANT4(unpack_sk)

# define pack_ciphertext VARIANT4(pack_ciphertext)

# define unpack_ciphertext VARIANT4(unpack_ciphertext)

# define gen_matrix VARIANT4(gen_matrix)

# define indcpa_keypair_derand VARIANT4(indcpa_keypair_derand)

# define indcpa_enc VARIANT4(indcpa_enc)

# define indcpa_dec VARIANT4(indcpa_dec)

# include "kyber-kdep.c"

#endif