/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#ifdef FREEBL_NO_DEPEND
#include "stubs.h"
#endif

#include <stdbool.h>

#include "blapi.h"
#include "secerr.h"
#include "secitem.h"

#include "kyber.h"
/* Vendored libcrux ML-KEM, from the combined ML-KEM + ML-DSA extraction in
 * lib/freebl/libcrux/ (shares its SHA-3/core with ml_dsa.c). Compiled into the
 * 'crux' static library; see lib/freebl/libcrux/README.md and the freebl build
 * files for the isolated include scope. */
#include "libcrux_mlkem512.h"
#include "libcrux_mlkem512_portable.h"
#include "libcrux_mlkem768.h"
#include "libcrux_mlkem768_portable.h"
#include "libcrux_mlkem1024.h"
#include "libcrux_mlkem1024_portable.h"

/* Local names for the libcrux entry point types that eurydice does not give a
 * stable name to: the trailing hash on Eurydice_arr_* and tuple_* is
 * regenerated by every extraction, so a re-vendor only has to update this block
 * rather than each use below. The keys have stable upstream names and are used
 * directly. */
typedef Eurydice_arr_c7 MlKemKeypairCoins;   /* uint8_t[64] */
typedef Eurydice_arr_ec MlKemEncCoins;       /* uint8_t[32] */
typedef Eurydice_arr_ec MlKemSharedSecret;   /* uint8_t[32] */
typedef Eurydice_arr_d2 MlKem512Ciphertext;  /* uint8_t[768] */
typedef Eurydice_arr_2b MlKem768Ciphertext;  /* uint8_t[1088] */
typedef Eurydice_arr_d1 MlKem1024Ciphertext; /* uint8_t[1568] */
typedef tuple_ab MlKem512Encapsulation;      /* { MlKem512Ciphertext fst; MlKemSharedSecret snd; } */
typedef tuple_f4 MlKem768Encapsulation;      /* { MlKem768Ciphertext fst; MlKemSharedSecret snd; } */
typedef tuple_25 MlKem1024Encapsulation;     /* { MlKem1024Ciphertext fst; MlKemSharedSecret snd; } */

/* Consistency check between the vendored libcrux array types and kyber.h. The
 * key/ciphertext/coin lengths are the sizes of the Eurydice_arr_* structs the
 * libcrux entry points take and return; if a re-vendor changes a size these fail
 * loudly. */
PR_STATIC_ASSERT(MLKEM512_PUBLIC_KEY_BYTES == sizeof(((libcrux_ml_kem_mlkem512_MlKem512PublicKey *)0)->data));
PR_STATIC_ASSERT(MLKEM512_PRIVATE_KEY_BYTES == sizeof(((libcrux_ml_kem_mlkem512_MlKem512PrivateKey *)0)->data));
PR_STATIC_ASSERT(MLKEM512_CIPHERTEXT_BYTES == sizeof(((MlKem512Ciphertext *)0)->data));
PR_STATIC_ASSERT(KYBER768_PUBLIC_KEY_BYTES == sizeof(((libcrux_ml_kem_mlkem768_MlKem768PublicKey *)0)->data));
PR_STATIC_ASSERT(KYBER768_PRIVATE_KEY_BYTES == sizeof(((libcrux_ml_kem_mlkem768_MlKem768PrivateKey *)0)->data));
PR_STATIC_ASSERT(KYBER768_CIPHERTEXT_BYTES == sizeof(((MlKem768Ciphertext *)0)->data));
PR_STATIC_ASSERT(KYBER_SHARED_SECRET_BYTES == sizeof(((MlKemSharedSecret *)0)->data));
PR_STATIC_ASSERT(KYBER_KEYPAIR_COIN_BYTES == sizeof(((MlKemKeypairCoins *)0)->data));
PR_STATIC_ASSERT(KYBER_ENC_COIN_BYTES == sizeof(((MlKemEncCoins *)0)->data));
PR_STATIC_ASSERT(MLKEM1024_PUBLIC_KEY_BYTES == sizeof(((libcrux_ml_kem_mlkem1024_MlKem1024PublicKey *)0)->data));
PR_STATIC_ASSERT(MLKEM1024_PRIVATE_KEY_BYTES == sizeof(((libcrux_ml_kem_mlkem1024_MlKem1024PrivateKey *)0)->data));
PR_STATIC_ASSERT(MLKEM1024_CIPHERTEXT_BYTES == sizeof(((MlKem1024Ciphertext *)0)->data));

static bool
valid_params(KyberParams params)
{
    switch (params) {
        case params_ml_kem512:
        case params_ml_kem512_test_mode:
        case params_ml_kem768:
        case params_ml_kem768_test_mode:
        case params_ml_kem1024:
        case params_ml_kem1024_test_mode:
            return true;
        default:
            return false;
    }
}

static bool
valid_pubkey(KyberParams params, const SECItem *pubkey)
{
    switch (params) {
        case params_ml_kem512:
        case params_ml_kem512_test_mode:
            return pubkey && pubkey->len == MLKEM512_PUBLIC_KEY_BYTES;
        case params_ml_kem768:
        case params_ml_kem768_test_mode:
            return pubkey && pubkey->len == KYBER768_PUBLIC_KEY_BYTES;
        case params_ml_kem1024:
        case params_ml_kem1024_test_mode:
            return pubkey && pubkey->len == MLKEM1024_PUBLIC_KEY_BYTES;
        default:
            return false;
    }
}

static bool
valid_privkey(KyberParams params, const SECItem *privkey)
{
    switch (params) {
        case params_ml_kem512:
        case params_ml_kem512_test_mode:
            return privkey && privkey->len == MLKEM512_PRIVATE_KEY_BYTES;
        case params_ml_kem768:
        case params_ml_kem768_test_mode:
            return privkey && privkey->len == KYBER768_PRIVATE_KEY_BYTES;
        case params_ml_kem1024:
        case params_ml_kem1024_test_mode:
            return privkey && privkey->len == MLKEM1024_PRIVATE_KEY_BYTES;
        default:
            return false;
    }
}

static bool
valid_ciphertext(KyberParams params, const SECItem *ciphertext)
{
    switch (params) {
        case params_ml_kem512:
        case params_ml_kem512_test_mode:
            return ciphertext && ciphertext->len == MLKEM512_CIPHERTEXT_BYTES;
        case params_ml_kem768:
        case params_ml_kem768_test_mode:
            return ciphertext && ciphertext->len == KYBER768_CIPHERTEXT_BYTES;
        case params_ml_kem1024:
        case params_ml_kem1024_test_mode:
            return ciphertext && ciphertext->len == MLKEM1024_CIPHERTEXT_BYTES;
        default:
            return false;
    }
}

static bool
valid_secret(KyberParams params, const SECItem *secret)
{
    switch (params) {
        case params_ml_kem512:
        case params_ml_kem512_test_mode:
        case params_ml_kem768:
        case params_ml_kem768_test_mode:
        case params_ml_kem1024:
        case params_ml_kem1024_test_mode:
            return secret && secret->len == KYBER_SHARED_SECRET_BYTES;
        default:
            return false;
    }
}

static bool
valid_keypair_seed(KyberParams params, const SECItem *seed)
{
    switch (params) {
        case params_ml_kem512:
        case params_ml_kem512_test_mode:
        case params_ml_kem768:
        case params_ml_kem768_test_mode:
        case params_ml_kem1024:
        case params_ml_kem1024_test_mode:
            return !seed || seed->len == KYBER_KEYPAIR_COIN_BYTES;
        default:
            return false;
    }
}

static bool
valid_enc_seed(KyberParams params, const SECItem *seed)
{
    switch (params) {
        /* Only the test-mode parameter sets accept caller-supplied
         * encapsulation randomness; everywhere else it must come from the
         * RNG. */
        case params_ml_kem512:
        case params_ml_kem768:
        case params_ml_kem1024:
            return !seed;
        case params_ml_kem512_test_mode:
        case params_ml_kem768_test_mode:
        case params_ml_kem1024_test_mode:
            return !seed || seed->len == KYBER_ENC_COIN_BYTES;
        default:
            return false;
    }
}

SECStatus
Kyber_NewKey(KyberParams params, const SECItem *keypair_seed, SECItem *privkey, SECItem *pubkey)
{
    if (!valid_params(params)) {
        PORT_SetError(SEC_ERROR_INVALID_ALGORITHM);
        return SECFailure;
    }

    if (!(valid_keypair_seed(params, keypair_seed) && valid_privkey(params, privkey) && valid_pubkey(params, pubkey))) {
        PORT_SetError(SEC_ERROR_INVALID_ARGS);
        return SECFailure;
    }

    uint8_t randbuf[KYBER_KEYPAIR_COIN_BYTES];
    uint8_t *coins;
    if (keypair_seed) {
        coins = keypair_seed->data;
    } else {
        if (RNG_GenerateGlobalRandomBytes(randbuf, sizeof randbuf) != SECSuccess) {
            PORT_SetError(SEC_ERROR_NEED_RANDOM);
            return SECFailure;
        }
        coins = randbuf;
    }
    NSS_CLASSIFY(coins, KYBER_KEYPAIR_COIN_BYTES);
    if (params == params_ml_kem512 || params == params_ml_kem512_test_mode) {
        MlKemKeypairCoins coinsArr;
        memcpy(coinsArr.data, coins, KYBER_KEYPAIR_COIN_BYTES);
        libcrux_ml_kem_mlkem512_MlKem512KeyPair keys = libcrux_ml_kem_mlkem512_portable_generate_key_pair(coinsArr);
        memcpy(pubkey->data, keys.pk.data, MLKEM512_PUBLIC_KEY_BYTES);
        memcpy(privkey->data, keys.sk.data, MLKEM512_PRIVATE_KEY_BYTES);
    } else if (params == params_ml_kem768 || params == params_ml_kem768_test_mode) {
        MlKemKeypairCoins coinsArr;
        memcpy(coinsArr.data, coins, KYBER_KEYPAIR_COIN_BYTES);
        libcrux_ml_kem_mlkem768_MlKem768KeyPair keys = libcrux_ml_kem_mlkem768_portable_generate_key_pair(coinsArr);
        memcpy(pubkey->data, keys.pk.data, KYBER768_PUBLIC_KEY_BYTES);
        memcpy(privkey->data, keys.sk.data, KYBER768_PRIVATE_KEY_BYTES);
    } else if (params == params_ml_kem1024 || params == params_ml_kem1024_test_mode) {
        MlKemKeypairCoins coinsArr;
        memcpy(coinsArr.data, coins, KYBER_KEYPAIR_COIN_BYTES);
        libcrux_ml_kem_mlkem1024_MlKem1024KeyPair keys = libcrux_ml_kem_mlkem1024_portable_generate_key_pair(coinsArr);
        memcpy(pubkey->data, keys.pk.data, MLKEM1024_PUBLIC_KEY_BYTES);
        memcpy(privkey->data, keys.sk.data, MLKEM1024_PRIVATE_KEY_BYTES);
    } else {
        /* round-3 Kyber removed; only ML-KEM is supported */
        PORT_SetError(SEC_ERROR_INVALID_ALGORITHM);
        return SECFailure;
    }
    NSS_DECLASSIFY(pubkey->data, pubkey->len);
    return SECSuccess;
}

SECStatus
Kyber_Encapsulate(KyberParams params, const SECItem *enc_seed, const SECItem *pubkey, SECItem *ciphertext, SECItem *secret)
{
    if (!valid_params(params)) {
        PORT_SetError(SEC_ERROR_INVALID_ALGORITHM);
        return SECFailure;
    }

    if (!(valid_enc_seed(params, enc_seed) && valid_pubkey(params, pubkey) && valid_ciphertext(params, ciphertext) && valid_secret(params, secret))) {
        PORT_SetError(SEC_ERROR_INVALID_ARGS);
        return SECFailure;
    }

    uint8_t randbuf[KYBER_ENC_COIN_BYTES];
    uint8_t *coins;
    if (enc_seed) {
        coins = enc_seed->data;
    } else {
        if (RNG_GenerateGlobalRandomBytes(randbuf, sizeof randbuf) != SECSuccess) {
            PORT_SetError(SEC_ERROR_NEED_RANDOM);
            return SECFailure;
        }
        coins = randbuf;
    }
    NSS_CLASSIFY(coins, KYBER_ENC_COIN_BYTES);
    if (params == params_ml_kem512 || params == params_ml_kem512_test_mode) {
        libcrux_ml_kem_mlkem512_MlKem512PublicKey pk_value;
        memcpy(pk_value.data, pubkey->data, MLKEM512_PUBLIC_KEY_BYTES);

        bool valid_pk = libcrux_ml_kem_mlkem512_portable_validate_public_key(&pk_value);
        if (!valid_pk) {
            PORT_SetError(SEC_ERROR_INVALID_ARGS);
            return SECFailure;
        }

        MlKemEncCoins coinsArr;
        memcpy(coinsArr.data, coins, KYBER_ENC_COIN_BYTES);
        MlKem512Encapsulation encap = libcrux_ml_kem_mlkem512_portable_encapsulate(&pk_value, coinsArr);
        memcpy(ciphertext->data, encap.fst.data, MLKEM512_CIPHERTEXT_BYTES);
        memcpy(secret->data, encap.snd.data, KYBER_SHARED_SECRET_BYTES);
    } else if (params == params_ml_kem768 || params == params_ml_kem768_test_mode) {
        libcrux_ml_kem_mlkem768_MlKem768PublicKey pk_value;
        memcpy(pk_value.data, pubkey->data, KYBER768_PUBLIC_KEY_BYTES);

        bool valid_pk = libcrux_ml_kem_mlkem768_portable_validate_public_key(&pk_value);
        if (!valid_pk) {
            PORT_SetError(SEC_ERROR_INVALID_ARGS);
            return SECFailure;
        }

        MlKemEncCoins coinsArr;
        memcpy(coinsArr.data, coins, KYBER_ENC_COIN_BYTES);
        MlKem768Encapsulation encap = libcrux_ml_kem_mlkem768_portable_encapsulate(&pk_value, coinsArr);
        memcpy(ciphertext->data, encap.fst.data, KYBER768_CIPHERTEXT_BYTES);
        memcpy(secret->data, encap.snd.data, KYBER_SHARED_SECRET_BYTES);
    } else if (params == params_ml_kem1024 || params == params_ml_kem1024_test_mode) {
        libcrux_ml_kem_mlkem1024_MlKem1024PublicKey pk_value;
        memcpy(pk_value.data, pubkey->data, MLKEM1024_PUBLIC_KEY_BYTES);

        bool valid_pk = libcrux_ml_kem_mlkem1024_portable_validate_public_key(&pk_value);
        if (!valid_pk) {
            PORT_SetError(SEC_ERROR_INVALID_ARGS);
            return SECFailure;
        }

        MlKemEncCoins coinsArr;
        memcpy(coinsArr.data, coins, KYBER_ENC_COIN_BYTES);
        MlKem1024Encapsulation encap = libcrux_ml_kem_mlkem1024_portable_encapsulate(&pk_value, coinsArr);
        memcpy(ciphertext->data, encap.fst.data, MLKEM1024_CIPHERTEXT_BYTES);
        memcpy(secret->data, encap.snd.data, KYBER_SHARED_SECRET_BYTES);
    } else {
        /* round-3 Kyber removed; only ML-KEM is supported */
        PORT_SetError(SEC_ERROR_INVALID_ALGORITHM);
        return SECFailure;
    }

    return SECSuccess;
}

SECStatus
Kyber_Decapsulate(KyberParams params, const SECItem *privkey, const SECItem *ciphertext, SECItem *secret)
{
    if (!valid_params(params)) {
        PORT_SetError(SEC_ERROR_INVALID_ALGORITHM);
        return SECFailure;
    }

    if (!(valid_privkey(params, privkey) && valid_ciphertext(params, ciphertext) && valid_secret(params, secret))) {
        PORT_SetError(SEC_ERROR_INVALID_ARGS);
        return SECFailure;
    }

    if (params == params_ml_kem512 || params == params_ml_kem512_test_mode) {
        libcrux_ml_kem_mlkem512_MlKem512PrivateKey private_key;
        memcpy(private_key.data, privkey->data, MLKEM512_PRIVATE_KEY_BYTES);

        MlKem512Ciphertext cipher_text;
        memcpy(cipher_text.data, ciphertext->data, MLKEM512_CIPHERTEXT_BYTES);

        bool valid = libcrux_ml_kem_mlkem512_portable_validate_private_key(&private_key, &cipher_text);
        if (!valid) {
            PORT_SetError(SEC_ERROR_INVALID_ARGS);
            return SECFailure;
        }

        MlKemSharedSecret ss = libcrux_ml_kem_mlkem512_portable_decapsulate(&private_key, &cipher_text);
        memcpy(secret->data, ss.data, KYBER_SHARED_SECRET_BYTES);
    } else if (params == params_ml_kem768 || params == params_ml_kem768_test_mode) {
        libcrux_ml_kem_mlkem768_MlKem768PrivateKey private_key;
        memcpy(private_key.data, privkey->data, KYBER768_PRIVATE_KEY_BYTES);

        MlKem768Ciphertext cipher_text;
        memcpy(cipher_text.data, ciphertext->data, KYBER768_CIPHERTEXT_BYTES);

        bool valid = libcrux_ml_kem_mlkem768_portable_validate_private_key(&private_key, &cipher_text);
        if (!valid) {
            PORT_SetError(SEC_ERROR_INVALID_ARGS);
            return SECFailure;
        }

        MlKemSharedSecret ss = libcrux_ml_kem_mlkem768_portable_decapsulate(&private_key, &cipher_text);
        memcpy(secret->data, ss.data, KYBER_SHARED_SECRET_BYTES);
    } else if (params == params_ml_kem1024 || params == params_ml_kem1024_test_mode) {
        libcrux_ml_kem_mlkem1024_MlKem1024PrivateKey private_key;
        memcpy(private_key.data, privkey->data, MLKEM1024_PRIVATE_KEY_BYTES);

        MlKem1024Ciphertext cipher_text;
        memcpy(cipher_text.data, ciphertext->data, MLKEM1024_CIPHERTEXT_BYTES);

        bool valid = libcrux_ml_kem_mlkem1024_portable_validate_private_key(&private_key, &cipher_text);
        if (!valid) {
            PORT_SetError(SEC_ERROR_INVALID_ARGS);
            return SECFailure;
        }

        MlKemSharedSecret ss = libcrux_ml_kem_mlkem1024_portable_decapsulate(&private_key, &cipher_text);
        memcpy(secret->data, ss.data, KYBER_SHARED_SECRET_BYTES);
    } else {
        /* round-3 Kyber removed; only ML-KEM is supported */
        PORT_SetError(SEC_ERROR_INVALID_ALGORITHM);
        return SECFailure;
    }

    return SECSuccess;
}