/*

 * Copyright 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 <assert.h>

#include <limits.h>

#include <openssl/cms.h>

#include <openssl/core_names.h>

#include <openssl/err.h>

#include <openssl/decoder.h>

#include "internal/sizes.h"

#include "crypto/asn1.h"

#include "crypto/evp.h"

#include "cms_local.h"

static int kem_cms_decrypt(CMS_RecipientInfo *ri)

{

    uint32_t *kekLength;

    X509_ALGOR *wrap;

    EVP_PKEY_CTX *pctx;

    EVP_CIPHER_CTX *kekctx;

    uint32_t cipher_length;

    char name[OSSL_MAX_NAME_SIZE];

    EVP_CIPHER *kekcipher = NULL;

    int rv = 0;

    if (!ossl_cms_RecipientInfo_kemri_get0_alg(ri, &kekLength, &wrap))

        goto err;

    pctx = CMS_RecipientInfo_get0_pkey_ctx(ri);

    if (pctx == NULL)

        goto err;

    kekctx = CMS_RecipientInfo_kemri_get0_ctx(ri);

    if (kekctx == NULL)

        goto err;

    OBJ_obj2txt(name, sizeof(name), wrap->algorithm, 0);

    kekcipher = EVP_CIPHER_fetch(pctx->libctx, name, pctx->propquery);

    if (kekcipher == NULL || EVP_CIPHER_get_mode(kekcipher) != EVP_CIPH_WRAP_MODE)

        goto err;

    if (!EVP_EncryptInit_ex(kekctx, kekcipher, NULL, NULL, NULL))

        goto err;

    if (EVP_CIPHER_asn1_to_param(kekctx, wrap->parameter) <= 0)

        goto err;

    cipher_length = EVP_CIPHER_CTX_get_key_length(kekctx);

    if (cipher_length != *kekLength) {

        ERR_raise(ERR_LIB_CMS, CMS_R_INVALID_KEY_LENGTH);

        goto err;

    }

    rv = 1;

err:

    EVP_CIPHER_free(kekcipher);

    return rv;

}

static int kem_cms_encrypt(CMS_RecipientInfo *ri)

{

    uint32_t *kekLength;

    X509_ALGOR *wrap;

    X509_ALGOR *kdf;

    EVP_PKEY_CTX *pctx;

    EVP_PKEY *pkey;

    int security_bits;

    const ASN1_OBJECT *kdf_obj = NULL;

    unsigned char kemri_x509_algor[OSSL_MAX_ALGORITHM_ID_SIZE];

    OSSL_PARAM params[2];

    X509_ALGOR *x509_algor = NULL;

    EVP_CIPHER_CTX *kekctx;

    int wrap_nid;

    int rv = 0;

    if (!ossl_cms_RecipientInfo_kemri_get0_alg(ri, &kekLength, &wrap))

        goto err;

    kdf = CMS_RecipientInfo_kemri_get0_kdf_alg(ri);

    if (kdf == NULL)

        goto err;

    pctx = CMS_RecipientInfo_get0_pkey_ctx(ri);

    if (pctx == NULL)

        goto err;

    pkey = EVP_PKEY_CTX_get0_pkey(pctx);

    if (pkey == NULL)

        goto err;

    security_bits = EVP_PKEY_get_security_bits(pkey);

    if (security_bits == 0)

        goto err;

    X509_ALGOR_get0(&kdf_obj, NULL, NULL, kdf);

    if (kdf_obj == NULL || OBJ_obj2nid(kdf_obj) == NID_undef) {

        /*

         * If the KDF OID hasn't already been set, then query the provider

         * for a default KDF.

         */

        params[0] = OSSL_PARAM_construct_octet_string(OSSL_PKEY_PARAM_CMS_KEMRI_KDF_ALGORITHM,

            kemri_x509_algor, sizeof(kemri_x509_algor));

        params[1] = OSSL_PARAM_construct_end();

        if (!EVP_PKEY_get_params(pkey, params))

            goto err;

        if (OSSL_PARAM_modified(&params[0])) {

            const unsigned char *p = kemri_x509_algor;

            x509_algor = d2i_X509_ALGOR(NULL, &p, (long)params[0].return_size);

            if (x509_algor == NULL)

                goto err;

            if (!X509_ALGOR_copy(kdf, x509_algor))

                goto err;

        } else {

            if (!X509_ALGOR_set0(kdf, OBJ_nid2obj(NID_HKDF_SHA256), V_ASN1_UNDEF, NULL))

                return 0;

        }

    }

    /* Get wrap NID */

    kekctx = CMS_RecipientInfo_kemri_get0_ctx(ri);

    if (kekctx == NULL)

        goto err;

    *kekLength = EVP_CIPHER_CTX_get_key_length(kekctx);

    wrap_nid = EVP_CIPHER_CTX_get_type(kekctx);

    /* Package wrap algorithm in an AlgorithmIdentifier */

    ASN1_OBJECT_free(wrap->algorithm);

    ASN1_TYPE_free(wrap->parameter);

    wrap->algorithm = OBJ_nid2obj(wrap_nid);

    wrap->parameter = ASN1_TYPE_new();

    if (wrap->parameter == NULL)

        goto err;

    if (EVP_CIPHER_param_to_asn1(kekctx, wrap->parameter) <= 0) {

        ASN1_TYPE_free(wrap->parameter);

        wrap->parameter = NULL;

        goto err;

    }

    if (ASN1_TYPE_get(wrap->parameter) == NID_undef) {

        ASN1_TYPE_free(wrap->parameter);

        wrap->parameter = NULL;

    }

    rv = 1;

err:

    X509_ALGOR_free(x509_algor);

    return rv;

}

int ossl_cms_kem_envelope(CMS_RecipientInfo *ri, int decrypt)

{

    assert(decrypt == 0 || decrypt == 1);

    if (decrypt == 1)

        return kem_cms_decrypt(ri);

    if (decrypt == 0)

        return kem_cms_encrypt(ri);

    ERR_raise(ERR_LIB_CMS, CMS_R_NOT_SUPPORTED_FOR_THIS_KEY_TYPE);

    return 0;

}