/*

 * 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 <openssl/core_dispatch.h>

#include <openssl/core_names.h>

#include <openssl/crypto.h>

#include <openssl/err.h>

#include <openssl/evp.h>

#include <openssl/params.h>

#include <openssl/proverr.h>

#include <openssl/rand.h>

#include "prov/implementations.h"

#include "prov/mlx_kem.h"

#include "prov/provider_ctx.h"

#include "prov/providercommon.h"

static OSSL_FUNC_kem_newctx_fn mlx_kem_newctx;

static OSSL_FUNC_kem_freectx_fn mlx_kem_freectx;

static OSSL_FUNC_kem_encapsulate_init_fn mlx_kem_encapsulate_init;

static OSSL_FUNC_kem_encapsulate_fn mlx_kem_encapsulate;

static OSSL_FUNC_kem_decapsulate_init_fn mlx_kem_decapsulate_init;

static OSSL_FUNC_kem_decapsulate_fn mlx_kem_decapsulate;

static OSSL_FUNC_kem_set_ctx_params_fn mlx_kem_set_ctx_params;

static OSSL_FUNC_kem_settable_ctx_params_fn mlx_kem_settable_ctx_params;

typedef struct {

    OSSL_LIB_CTX *libctx;

    MLX_KEY *key;

    int op;

} PROV_MLX_KEM_CTX;

static void *mlx_kem_newctx(void *provctx)

{

    PROV_MLX_KEM_CTX *ctx;

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

        return NULL;

    ctx->libctx = PROV_LIBCTX_OF(provctx);

    ctx->key = NULL;

    ctx->op = 0;

    return ctx;

}

static void mlx_kem_freectx(void *vctx)

{

    OPENSSL_free(vctx);

}

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

                        ossl_unused const OSSL_PARAM params[])

{

    PROV_MLX_KEM_CTX *ctx = vctx;

    if (!ossl_prov_is_running())

        return 0;

    ctx->key = key;

    ctx->op = op;

    return 1;

}

static int

mlx_kem_encapsulate_init(void *vctx, void *vkey, const OSSL_PARAM params[])

{

    MLX_KEY *key = vkey;

    if (!mlx_kem_have_pubkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        return 0;

    }

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

}

static int

mlx_kem_decapsulate_init(void *vctx, void *vkey, const OSSL_PARAM params[])

{

    MLX_KEY *key = vkey;

    if (!mlx_kem_have_prvkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        return 0;

    }

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

}

static const OSSL_PARAM *mlx_kem_settable_ctx_params(ossl_unused void *vctx,

                                                     ossl_unused void *provctx)

{

    static const OSSL_PARAM params[] = { OSSL_PARAM_END };

    return params;

}

static int

mlx_kem_set_ctx_params(void *vctx, const OSSL_PARAM params[])

{

    return 1;

}

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

                               unsigned char *shsec, size_t *slen)

{

    MLX_KEY *key = ((PROV_MLX_KEM_CTX *) vctx)->key;

    EVP_PKEY_CTX *ctx = NULL;

    EVP_PKEY *xkey = NULL;

    size_t encap_clen;

    size_t encap_slen;

    uint8_t *cbuf;

    uint8_t *sbuf;

    int ml_kem_slot = key->xinfo->ml_kem_slot;

    int ret = 0;

    if (!mlx_kem_have_pubkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        goto end;

    }

    encap_clen = key->minfo->ctext_bytes + key->xinfo->pubkey_bytes;

    encap_slen = ML_KEM_SHARED_SECRET_BYTES + key->xinfo->shsec_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 output buffer");

        return 0;

    }

    if (clen == NULL) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_NULL_LENGTH_POINTER,

                       "null ciphertext input/output length pointer");

        return 0;

    } else if (*clen < encap_clen) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_OUTPUT_BUFFER_TOO_SMALL,

                       "ciphertext buffer too small");

        return 0;

    } 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");

        return 0;

    } else if (*slen < encap_slen) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_OUTPUT_BUFFER_TOO_SMALL,

                       "shared-secret buffer too small");

        return 0;

    } else {

        *slen = encap_slen;

    }

    /* ML-KEM encapsulation */

    encap_clen = key->minfo->ctext_bytes;

    encap_slen = ML_KEM_SHARED_SECRET_BYTES;

    cbuf = ctext + ml_kem_slot * key->xinfo->pubkey_bytes;

    sbuf = shsec + ml_kem_slot * key->xinfo->shsec_bytes;

    ctx = EVP_PKEY_CTX_new_from_pkey(key->libctx, key->mkey, key->propq);

    if (ctx == NULL

        || EVP_PKEY_encapsulate_init(ctx, NULL) <= 0

        || EVP_PKEY_encapsulate(ctx, cbuf, &encap_clen, sbuf, &encap_slen) <= 0)

        goto end;

    if (encap_clen != key->minfo->ctext_bytes) {

        ERR_raise_data(ERR_LIB_PROV, ERR_R_INTERNAL_ERROR,

                       "unexpected %s ciphertext output size: %lu",

                       key->minfo->algorithm_name, (unsigned long) encap_clen);

        goto end;

    }

    if (encap_slen != ML_KEM_SHARED_SECRET_BYTES) {

        ERR_raise_data(ERR_LIB_PROV, ERR_R_INTERNAL_ERROR,

                       "unexpected %s shared secret output size: %lu",

                       key->minfo->algorithm_name, (unsigned long) encap_slen);

        goto end;

    }

    EVP_PKEY_CTX_free(ctx);

    /*-

     * ECDHE encapsulation

     *

     * Generate own ephemeral private key and add its public key to ctext.

     *

     * Note, we could support a settable parameter that sets an extant ECDH

     * keypair as the keys to use in encap, making it possible to reuse the

     * same (TLS client) ECDHE keypair for both the classical EC keyshare and a

     * corresponding ECDHE + ML-KEM keypair.  But the TLS layer would then need

     * know that this is a hybrid, and that it can partly reuse the same keys

     * as another group for which a keyshare will be sent.  Deferred until we

     * support generating multiple keyshares, there's a workable keyshare

     * prediction specification, and the optimisation is justified.

     */

    cbuf = ctext + (1 - ml_kem_slot) * key->minfo->ctext_bytes;

    encap_clen = key->xinfo->pubkey_bytes;

    ctx = EVP_PKEY_CTX_new_from_pkey(key->libctx, key->xkey, key->propq);

    if (ctx == NULL

        || EVP_PKEY_keygen_init(ctx) <= 0

        || EVP_PKEY_keygen(ctx, &xkey) <= 0

        || EVP_PKEY_get_octet_string_param(xkey, OSSL_PKEY_PARAM_ENCODED_PUBLIC_KEY,

                                           cbuf, encap_clen, &encap_clen) <= 0)

        goto end;

    if (encap_clen != key->xinfo->pubkey_bytes) {

        ERR_raise_data(ERR_LIB_PROV, ERR_R_INTERNAL_ERROR,

                       "unexpected %s public key output size: %lu",

                       key->xinfo->algorithm_name, (unsigned long) encap_clen);

        goto end;

    }

    EVP_PKEY_CTX_free(ctx);

    /* Derive the ECDH shared secret */

    encap_slen = key->xinfo->shsec_bytes;

    sbuf = shsec + (1 - ml_kem_slot) * ML_KEM_SHARED_SECRET_BYTES;

    ctx = EVP_PKEY_CTX_new_from_pkey(key->libctx, xkey, key->propq);

    if (ctx == NULL

        || EVP_PKEY_derive_init(ctx) <= 0

        || EVP_PKEY_derive_set_peer(ctx, key->xkey) <= 0

        || EVP_PKEY_derive(ctx, sbuf, &encap_slen) <= 0)

        goto end;

    if (encap_slen != key->xinfo->shsec_bytes) {

        ERR_raise_data(ERR_LIB_PROV, ERR_R_INTERNAL_ERROR,

                       "unexpected %s shared secret output size: %lu",

                       key->xinfo->algorithm_name, (unsigned long) encap_slen);

        goto end;

    }

    ret = 1;

 end:

    EVP_PKEY_free(xkey);

    EVP_PKEY_CTX_free(ctx);

    return ret;

}

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

                               const uint8_t *ctext, size_t clen)

{

    MLX_KEY *key = ((PROV_MLX_KEM_CTX *) vctx)->key;

    EVP_PKEY_CTX *ctx = NULL;

    EVP_PKEY *xkey = NULL;

    const uint8_t *cbuf;

    uint8_t *sbuf;

    size_t decap_slen = ML_KEM_SHARED_SECRET_BYTES + key->xinfo->shsec_bytes;

    size_t decap_clen = key->minfo->ctext_bytes + key->xinfo->pubkey_bytes;

    int ml_kem_slot = key->xinfo->ml_kem_slot;

    int ret = 0;

    if (!mlx_kem_have_prvkey(key)) {

        ERR_raise(ERR_LIB_PROV, PROV_R_MISSING_KEY);

        return 0;

    }

    if (shsec == NULL) {

        if (slen == NULL)

            return 0;

        *slen = decap_slen;

        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;

    }

    if (clen != decap_clen) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_WRONG_CIPHERTEXT_SIZE,

                       "wrong decapsulation input ciphertext size: %lu",

                       (unsigned long) clen);

        return 0;

    }

    /* ML-KEM decapsulation */

    decap_clen = key->minfo->ctext_bytes;

    decap_slen = ML_KEM_SHARED_SECRET_BYTES;

    cbuf = ctext + ml_kem_slot * key->xinfo->pubkey_bytes;

    sbuf = shsec + ml_kem_slot * key->xinfo->shsec_bytes;

    ctx = EVP_PKEY_CTX_new_from_pkey(key->libctx, key->mkey, key->propq);

    if (ctx == NULL

        || EVP_PKEY_decapsulate_init(ctx, NULL) <= 0

        || EVP_PKEY_decapsulate(ctx, sbuf, &decap_slen, cbuf, decap_clen) <= 0)

        goto end;

    if (decap_slen != ML_KEM_SHARED_SECRET_BYTES) {

        ERR_raise_data(ERR_LIB_PROV, ERR_R_INTERNAL_ERROR,

                       "unexpected %s shared secret output size: %lu",

                       key->minfo->algorithm_name, (unsigned long) decap_slen);

        goto end;

    }

    EVP_PKEY_CTX_free(ctx);

    /* ECDH decapsulation */

    decap_clen = key->xinfo->pubkey_bytes;

    decap_slen = key->xinfo->shsec_bytes;

    cbuf = ctext + (1 - ml_kem_slot) * key->minfo->ctext_bytes;

    sbuf = shsec + (1 - ml_kem_slot) * ML_KEM_SHARED_SECRET_BYTES;

    ctx = EVP_PKEY_CTX_new_from_pkey(key->libctx, key->xkey, key->propq);

    if (ctx == NULL

        || (xkey = EVP_PKEY_new()) == NULL

        || EVP_PKEY_copy_parameters(xkey, key->xkey) <= 0

        || EVP_PKEY_set1_encoded_public_key(xkey, cbuf, decap_clen) <= 0

        || EVP_PKEY_derive_init(ctx) <= 0

        || EVP_PKEY_derive_set_peer(ctx, xkey) <= 0

        || EVP_PKEY_derive(ctx, sbuf, &decap_slen) <= 0)

        goto end;

    if (decap_slen != key->xinfo->shsec_bytes) {

        ERR_raise_data(ERR_LIB_PROV, ERR_R_INTERNAL_ERROR,

                       "unexpected %s shared secret output size: %lu",

                       key->xinfo->algorithm_name, (unsigned long) decap_slen);

        goto end;

    }

    ret = 1;

 end:

    EVP_PKEY_CTX_free(ctx);

    EVP_PKEY_free(xkey);

    return ret;

}

const OSSL_DISPATCH ossl_mlx_kem_asym_kem_functions[] = {

    { OSSL_FUNC_KEM_NEWCTX, (OSSL_FUNC) mlx_kem_newctx },

    { OSSL_FUNC_KEM_ENCAPSULATE_INIT, (OSSL_FUNC) mlx_kem_encapsulate_init },

    { OSSL_FUNC_KEM_ENCAPSULATE, (OSSL_FUNC) mlx_kem_encapsulate },

    { OSSL_FUNC_KEM_DECAPSULATE_INIT, (OSSL_FUNC) mlx_kem_decapsulate_init },

    { OSSL_FUNC_KEM_DECAPSULATE, (OSSL_FUNC) mlx_kem_decapsulate },

    { OSSL_FUNC_KEM_FREECTX, (OSSL_FUNC) mlx_kem_freectx },

    { OSSL_FUNC_KEM_SET_CTX_PARAMS, (OSSL_FUNC) mlx_kem_set_ctx_params },

    { OSSL_FUNC_KEM_SETTABLE_CTX_PARAMS, (OSSL_FUNC) mlx_kem_settable_ctx_params },

    OSSL_DISPATCH_END

};