/*

 * Copyright 2024-2025 The OpenSSL Project Authors. All Rights Reserved.

 *

 * Licensed under the Apache License 2.0 (the "License").  You may not use

 * this file except in compliance with the License.  You can obtain a copy

 * in the file LICENSE in the source distribution or at

 * https://www.openssl.org/source/license.html

 */

#include <string.h>

#include <openssl/crypto.h>

#include <openssl/evp.h>

#include <openssl/core_dispatch.h>

#include <openssl/core_names.h>

#include <openssl/params.h>

#include <openssl/err.h>

#include <openssl/proverr.h>

#include "crypto/ml_kem.h"

#include "prov/provider_ctx.h"

#include "prov/implementations.h"

#include "prov/securitycheck.h"

#include "prov/providercommon.h"

static OSSL_FUNC_kem_newctx_fn ml_kem_newctx;

static OSSL_FUNC_kem_freectx_fn ml_kem_freectx;

static OSSL_FUNC_kem_encapsulate_init_fn ml_kem_encapsulate_init;

static OSSL_FUNC_kem_encapsulate_fn ml_kem_encapsulate;

static OSSL_FUNC_kem_decapsulate_init_fn ml_kem_decapsulate_init;

static OSSL_FUNC_kem_decapsulate_fn ml_kem_decapsulate;

static OSSL_FUNC_kem_set_ctx_params_fn ml_kem_set_ctx_params;

static OSSL_FUNC_kem_settable_ctx_params_fn ml_kem_settable_ctx_params;

typedef struct {

    ML_KEM_KEY *key;

    uint8_t entropy_buf[ML_KEM_RANDOM_BYTES];

    uint8_t *entropy;

    int op;

} PROV_ML_KEM_CTX;

static void *ml_kem_newctx(void *provctx)

{

    PROV_ML_KEM_CTX *ctx;

    if ((ctx = OPENSSL_malloc(sizeof(*ctx))) == NULL)

        return NULL;

    ctx->key = NULL;

    ctx->entropy = NULL;

    ctx->op = 0;

    return ctx;

}

static void ml_kem_freectx(void *vctx)

{

    PROV_ML_KEM_CTX *ctx = vctx;

    if (ctx->entropy != NULL)

        OPENSSL_cleanse(ctx->entropy, ML_KEM_RANDOM_BYTES);

    OPENSSL_free(ctx);

}

static int ml_kem_init(void *vctx, int op, void *key,

    const OSSL_PARAM params[])

{

    PROV_ML_KEM_CTX *ctx = vctx;

    if (!ossl_prov_is_running())

        return 0;

    ctx->key = key;

    ctx->op = op;

    return ml_kem_set_ctx_params(vctx, params);

}

static int ml_kem_encapsulate_init(void *vctx, void *vkey,

    const OSSL_PARAM params[])

{

    ML_KEM_KEY *key = vkey;

    if (!ossl_ml_kem_have_pubkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        return 0;

    }

    return ml_kem_init(vctx, EVP_PKEY_OP_ENCAPSULATE, key, params);

}

static int ml_kem_decapsulate_init(void *vctx, void *vkey,

    const OSSL_PARAM params[])

{

    ML_KEM_KEY *key = vkey;

    if (!ossl_ml_kem_have_prvkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        return 0;

    }

    return ml_kem_init(vctx, EVP_PKEY_OP_DECAPSULATE, key, params);

}

static int ml_kem_set_ctx_params(void *vctx, const OSSL_PARAM params[])

{

    PROV_ML_KEM_CTX *ctx = vctx;

    const OSSL_PARAM *p;

    if (ctx == NULL)

        return 0;

    if (ctx->op == EVP_PKEY_OP_DECAPSULATE && ctx->entropy != NULL) {

        /* Decapsulation is deterministic */

        OPENSSL_cleanse(ctx->entropy, ML_KEM_RANDOM_BYTES);

        ctx->entropy = NULL;

    }

    if (ossl_param_is_empty(params))

        return 1;

    /* Encapsulation ephemeral input key material "ikmE" */

    if (ctx->op == EVP_PKEY_OP_ENCAPSULATE

        && (p = OSSL_PARAM_locate_const(params, OSSL_KEM_PARAM_IKME)) != NULL) {

        size_t len = ML_KEM_RANDOM_BYTES;

        ctx->entropy = ctx->entropy_buf;

        if (OSSL_PARAM_get_octet_string(p, (void **)&ctx->entropy,

                len, &len)

            && len == ML_KEM_RANDOM_BYTES)

            return 1;

        /* Possibly, but much less likely wrong type */

        ERR_raise(ERR_LIB_PROV, PROV_R_INVALID_SEED_LENGTH);

        ctx->entropy = NULL;

        return 0;

    }

    return 1;

}

static const OSSL_PARAM *ml_kem_settable_ctx_params(ossl_unused void *vctx,

    ossl_unused void *provctx)

{

    static const OSSL_PARAM params[] = {

        OSSL_PARAM_octet_string(OSSL_KEM_PARAM_IKME, NULL, 0),

        OSSL_PARAM_END

    };

    return params;

}

static int ml_kem_encapsulate(void *vctx, unsigned char *ctext, size_t *clen,

    unsigned char *shsec, size_t *slen)

{

    PROV_ML_KEM_CTX *ctx = vctx;

    ML_KEM_KEY *key = ctx->key;

    const ML_KEM_VINFO *v;

    size_t encap_clen;

    size_t encap_slen;

    int ret = 0;

    if (!ossl_ml_kem_have_pubkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        goto end;

    }

    v = ossl_ml_kem_key_vinfo(key);

    encap_clen = v->ctext_bytes;

    encap_slen = ML_KEM_SHARED_SECRET_BYTES;

    if (ctext == NULL) {

        if (clen == NULL && slen == NULL)

            return 0;

        if (clen != NULL)

            *clen = encap_clen;

        if (slen != NULL)

            *slen = encap_slen;

        return 1;

    }

    if (shsec == NULL) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_NULL_OUTPUT_BUFFER,

            "NULL shared-secret buffer");

        goto end;

    }

    if (clen == NULL) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_NULL_LENGTH_POINTER,

            "null ciphertext input/output length pointer");

        goto end;

    } else if (*clen < encap_clen) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_OUTPUT_BUFFER_TOO_SMALL,

            "ciphertext buffer too small");

        goto end;

    } else {

        *clen = encap_clen;

    }

    if (slen == NULL) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_NULL_LENGTH_POINTER,

            "null shared secret input/output length pointer");

        goto end;

    } else if (*slen < encap_slen) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_OUTPUT_BUFFER_TOO_SMALL,

            "shared-secret buffer too small");

        goto end;

    } else {

        *slen = encap_slen;

    }

    if (ctx->entropy != NULL)

        ret = ossl_ml_kem_encap_seed(ctext, encap_clen, shsec, encap_slen,

            ctx->entropy, ML_KEM_RANDOM_BYTES, key);

    else

        ret = ossl_ml_kem_encap_rand(ctext, encap_clen, shsec, encap_slen, key);

end:

    /*

     * One shot entropy, each encapsulate call must either provide a new

     * "ikmE", or else will use a random value.  If a caller sets an explicit

     * ikmE once for testing, and later performs multiple encapsulations

     * without again calling encapsulate_init(), these should not share the

     * original entropy.

     */

    if (ctx->entropy != NULL) {

        OPENSSL_cleanse(ctx->entropy, ML_KEM_RANDOM_BYTES);

        ctx->entropy = NULL;

    }

    return ret;

}

static int ml_kem_decapsulate(void *vctx, uint8_t *shsec, size_t *slen,

    const uint8_t *ctext, size_t clen)

{

    PROV_ML_KEM_CTX *ctx = vctx;

    ML_KEM_KEY *key = ctx->key;

    size_t decap_slen = ML_KEM_SHARED_SECRET_BYTES;

    if (!ossl_ml_kem_have_prvkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        return 0;

    }

    if (shsec == NULL) {

        if (slen == NULL)

            return 0;

        *slen = ML_KEM_SHARED_SECRET_BYTES;

        return 1;

    }

    /* For now tolerate newly-deprecated NULL length pointers. */

    if (slen == NULL) {

        slen = &decap_slen;

    } else if (*slen < decap_slen) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_OUTPUT_BUFFER_TOO_SMALL,

            "shared-secret buffer too small");

        return 0;

    } else {

        *slen = decap_slen;

    }

    /* ML-KEM decap handles incorrect ciphertext lengths internally */

    return ossl_ml_kem_decap(shsec, decap_slen, ctext, clen, key);

}

const OSSL_DISPATCH ossl_ml_kem_asym_kem_functions[] = {

    { OSSL_FUNC_KEM_NEWCTX, (OSSL_FUNC)ml_kem_newctx },

    { OSSL_FUNC_KEM_ENCAPSULATE_INIT, (OSSL_FUNC)ml_kem_encapsulate_init },

    { OSSL_FUNC_KEM_ENCAPSULATE, (OSSL_FUNC)ml_kem_encapsulate },

    { OSSL_FUNC_KEM_DECAPSULATE_INIT, (OSSL_FUNC)ml_kem_decapsulate_init },

    { OSSL_FUNC_KEM_DECAPSULATE, (OSSL_FUNC)ml_kem_decapsulate },

    { OSSL_FUNC_KEM_FREECTX, (OSSL_FUNC)ml_kem_freectx },

    { OSSL_FUNC_KEM_SET_CTX_PARAMS, (OSSL_FUNC)ml_kem_set_ctx_params },

    { OSSL_FUNC_KEM_SETTABLE_CTX_PARAMS, (OSSL_FUNC)ml_kem_settable_ctx_params },

    OSSL_DISPATCH_END

};