



#include <stdlib.h>
#include <string.h>
#include <openssl/core_dispatch.h>
#include <openssl/core_names.h>
#include <openssl/params.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/proverr.h>

#include "prov/implementations.h"
#include "prov/provider_ctx.h"
#include "prov/provider_util.h"
#include "prov/providercommon.h"
#include "internal/cryptlib.h" 


static OSSL_FUNC_mac_newctx_fn kmac128_new;
static OSSL_FUNC_mac_newctx_fn kmac256_new;
static OSSL_FUNC_mac_dupctx_fn kmac_dup;
static OSSL_FUNC_mac_freectx_fn kmac_free;
static OSSL_FUNC_mac_gettable_ctx_params_fn kmac_gettable_ctx_params;
static OSSL_FUNC_mac_get_ctx_params_fn kmac_get_ctx_params;
static OSSL_FUNC_mac_settable_ctx_params_fn kmac_settable_ctx_params;
static OSSL_FUNC_mac_set_ctx_params_fn kmac_set_ctx_params;
static OSSL_FUNC_mac_init_fn kmac_init;
static OSSL_FUNC_mac_update_fn kmac_update;
static OSSL_FUNC_mac_final_fn kmac_final;

#define KMAC_MAX_BLOCKSIZE ((1600 - 128 * 2) / 8) 


#define KMAC_MAX_OUTPUT_LEN (0xFFFFFF / 8)
#define KMAC_MAX_ENCODED_HEADER_LEN (1 + 3)


#define KMAC_MAX_CUSTOM 512


#define KMAC_MAX_CUSTOM_ENCODED (KMAC_MAX_CUSTOM + KMAC_MAX_ENCODED_HEADER_LEN)


#define KMAC_MAX_KEY 512
#define KMAC_MIN_KEY 4


#define KMAC_MAX_KEY_ENCODED (KMAC_MAX_BLOCKSIZE * 4)


static const unsigned char kmac_string[] = {
    0x01, 0x20, 0x4B, 0x4D, 0x41, 0x43
};

#define KMAC_FLAG_XOF_MODE          1

struct kmac_data_st {
    void  *provctx;
    EVP_MD_CTX *ctx;
    PROV_DIGEST digest;
    size_t out_len;
    size_t key_len;
    size_t custom_len;
    
    int xof_mode;
    
    unsigned char key[KMAC_MAX_KEY_ENCODED];
    unsigned char custom[KMAC_MAX_CUSTOM_ENCODED];
};

static int encode_string(unsigned char *out, size_t out_max_len, size_t *out_len,
                         const unsigned char *in, size_t in_len);
static int right_encode(unsigned char *out, size_t out_max_len, size_t *out_len,
                        size_t bits);
static int bytepad(unsigned char *out, size_t *out_len,
                   const unsigned char *in1, size_t in1_len,
                   const unsigned char *in2, size_t in2_len,
                   size_t w);
static int kmac_bytepad_encode_key(unsigned char *out, size_t out_max_len,
                                   size_t *out_len,
                                   const unsigned char *in, size_t in_len,
                                   size_t w);

static void kmac_free(void *vmacctx)
{
    struct kmac_data_st *kctx = vmacctx;

    if (kctx != NULL) {
        EVP_MD_CTX_free(kctx->ctx);
        ossl_prov_digest_reset(&kctx->digest);
        OPENSSL_cleanse(kctx->key, kctx->key_len);
        OPENSSL_cleanse(kctx->custom, kctx->custom_len);
        OPENSSL_free(kctx);
    }
}


static struct kmac_data_st *kmac_new(void *provctx)
{
    struct kmac_data_st *kctx;

    if (!ossl_prov_is_running())
        return NULL;

    if ((kctx = OPENSSL_zalloc(sizeof(*kctx))) == NULL
            || (kctx->ctx = EVP_MD_CTX_new()) == NULL) {
        kmac_free(kctx);
        return NULL;
    }
    kctx->provctx = provctx;
    return kctx;
}

static void *kmac_fetch_new(void *provctx, const OSSL_PARAM *params)
{
    struct kmac_data_st *kctx = kmac_new(provctx);

    if (kctx == NULL)
        return 0;
    if (!ossl_prov_digest_load_from_params(&kctx->digest, params,
                                      PROV_LIBCTX_OF(provctx))) {
        kmac_free(kctx);
        return 0;
    }

    kctx->out_len = EVP_MD_get_size(ossl_prov_digest_md(&kctx->digest));
    return kctx;
}

static void *kmac128_new(void *provctx)
{
    static const OSSL_PARAM kmac128_params[] = {
        OSSL_PARAM_utf8_string("digest", OSSL_DIGEST_NAME_KECCAK_KMAC128,
                               sizeof(OSSL_DIGEST_NAME_KECCAK_KMAC128)),
        OSSL_PARAM_END
    };
    return kmac_fetch_new(provctx, kmac128_params);
}

static void *kmac256_new(void *provctx)
{
    static const OSSL_PARAM kmac256_params[] = {
        OSSL_PARAM_utf8_string("digest", OSSL_DIGEST_NAME_KECCAK_KMAC256,
                               sizeof(OSSL_DIGEST_NAME_KECCAK_KMAC256)),
        OSSL_PARAM_END
    };
    return kmac_fetch_new(provctx, kmac256_params);
}

static void *kmac_dup(void *vsrc)
{
    struct kmac_data_st *src = vsrc;
    struct kmac_data_st *dst;

    if (!ossl_prov_is_running())
        return NULL;

    dst = kmac_new(src->provctx);
    if (dst == NULL)
        return NULL;

    if (!EVP_MD_CTX_copy(dst->ctx, src->ctx)
        || !ossl_prov_digest_copy(&dst->digest, &src->digest)) {
        kmac_free(dst);
        return NULL;
    }

    dst->out_len = src->out_len;
    dst->key_len = src->key_len;
    dst->custom_len = src->custom_len;
    dst->xof_mode = src->xof_mode;
    memcpy(dst->key, src->key, src->key_len);
    memcpy(dst->custom, src->custom, dst->custom_len);

    return dst;
}

static int kmac_setkey(struct kmac_data_st *kctx, const unsigned char *key,
                       size_t keylen)
{
    const EVP_MD *digest = ossl_prov_digest_md(&kctx->digest);
    int w = EVP_MD_get_block_size(digest);

    if (keylen < KMAC_MIN_KEY || keylen > KMAC_MAX_KEY) {
        ERR_raise(ERR_LIB_PROV, PROV_R_INVALID_KEY_LENGTH);
        return 0;
    }
    if (w <= 0) {
        ERR_raise(ERR_LIB_PROV, PROV_R_INVALID_DIGEST_LENGTH);
        return 0;
    }
    if (!kmac_bytepad_encode_key(kctx->key, sizeof(kctx->key), &kctx->key_len,
                                 key, keylen, (size_t)w))
        return 0;
    return 1;
}


static int kmac_init(void *vmacctx, const unsigned char *key,
                     size_t keylen, const OSSL_PARAM params[])
{
    struct kmac_data_st *kctx = vmacctx;
    EVP_MD_CTX *ctx = kctx->ctx;
    unsigned char *out;
    size_t out_len, block_len;
    int res, t;

    if (!ossl_prov_is_running() || !kmac_set_ctx_params(kctx, params))
        return 0;

    if (key != NULL) {
        if (!kmac_setkey(kctx, key, keylen))
            return 0;
    } else if (kctx->key_len == 0) {
        
        ERR_raise(ERR_LIB_PROV, PROV_R_NO_KEY_SET);
        return 0;
    }
    if (!EVP_DigestInit_ex(kctx->ctx, ossl_prov_digest_md(&kctx->digest),
                           NULL))
        return 0;

    t = EVP_MD_get_block_size(ossl_prov_digest_md(&kctx->digest));
    if (t <= 0) {
        ERR_raise(ERR_LIB_PROV, PROV_R_INVALID_DIGEST_LENGTH);
        return 0;
    }
    block_len = t;

    
    if (kctx->custom_len == 0) {
        const OSSL_PARAM cparams[] = {
            OSSL_PARAM_octet_string(OSSL_MAC_PARAM_CUSTOM, "", 0),
            OSSL_PARAM_END
        };
        (void)kmac_set_ctx_params(kctx, cparams);
    }

    if (!bytepad(NULL, &out_len, kmac_string, sizeof(kmac_string),
                 kctx->custom, kctx->custom_len, block_len)) {
        ERR_raise(ERR_LIB_PROV, ERR_R_INTERNAL_ERROR);
        return 0;
    }
    out = OPENSSL_malloc(out_len);
    if (out == NULL)
        return 0;
    res = bytepad(out, NULL, kmac_string, sizeof(kmac_string),
                  kctx->custom, kctx->custom_len, block_len)
          && EVP_DigestUpdate(ctx, out, out_len)
          && EVP_DigestUpdate(ctx, kctx->key, kctx->key_len);
    OPENSSL_free(out);
    return res;
}

static int kmac_update(void *vmacctx, const unsigned char *data,
                       size_t datalen)
{
    struct kmac_data_st *kctx = vmacctx;

    return EVP_DigestUpdate(kctx->ctx, data, datalen);
}

static int kmac_final(void *vmacctx, unsigned char *out, size_t *outl,
                      size_t outsize)
{
    struct kmac_data_st *kctx = vmacctx;
    EVP_MD_CTX *ctx = kctx->ctx;
    size_t lbits, len;
    unsigned char encoded_outlen[KMAC_MAX_ENCODED_HEADER_LEN];
    int ok;

    if (!ossl_prov_is_running())
        return 0;

    
    lbits = (kctx->xof_mode ? 0 : (kctx->out_len * 8));

    ok = right_encode(encoded_outlen, sizeof(encoded_outlen), &len, lbits)
        && EVP_DigestUpdate(ctx, encoded_outlen, len)
        && EVP_DigestFinalXOF(ctx, out, kctx->out_len);
    *outl = kctx->out_len;
    return ok;
}

static const OSSL_PARAM known_gettable_ctx_params[] = {
    OSSL_PARAM_size_t(OSSL_MAC_PARAM_SIZE, NULL),
    OSSL_PARAM_size_t(OSSL_MAC_PARAM_BLOCK_SIZE, NULL),
    OSSL_PARAM_END
};
static const OSSL_PARAM *kmac_gettable_ctx_params(ossl_unused void *ctx,
                                                  ossl_unused void *provctx)
{
    return known_gettable_ctx_params;
}

static int kmac_get_ctx_params(void *vmacctx, OSSL_PARAM params[])
{
    struct kmac_data_st *kctx = vmacctx;
    OSSL_PARAM *p;
    int sz;

    if ((p = OSSL_PARAM_locate(params, OSSL_MAC_PARAM_SIZE)) != NULL
            && !OSSL_PARAM_set_size_t(p, kctx->out_len))
        return 0;

    if ((p = OSSL_PARAM_locate(params, OSSL_MAC_PARAM_BLOCK_SIZE)) != NULL) {
        sz = EVP_MD_block_size(ossl_prov_digest_md(&kctx->digest));
        if (!OSSL_PARAM_set_int(p, sz))
            return 0;
    }

    return 1;
}

static const OSSL_PARAM known_settable_ctx_params[] = {
    OSSL_PARAM_int(OSSL_MAC_PARAM_XOF, NULL),
    OSSL_PARAM_size_t(OSSL_MAC_PARAM_SIZE, NULL),
    OSSL_PARAM_octet_string(OSSL_MAC_PARAM_KEY, NULL, 0),
    OSSL_PARAM_octet_string(OSSL_MAC_PARAM_CUSTOM, NULL, 0),
    OSSL_PARAM_END
};
static const OSSL_PARAM *kmac_settable_ctx_params(ossl_unused void *ctx,
                                                  ossl_unused void *provctx)
{
    return known_settable_ctx_params;
}


static int kmac_set_ctx_params(void *vmacctx, const OSSL_PARAM *params)
{
    struct kmac_data_st *kctx = vmacctx;
    const OSSL_PARAM *p;

    if (params == NULL)
        return 1;

    if ((p = OSSL_PARAM_locate_const(params, OSSL_MAC_PARAM_XOF)) != NULL
        && !OSSL_PARAM_get_int(p, &kctx->xof_mode))
        return 0;
    if ((p = OSSL_PARAM_locate_const(params, OSSL_MAC_PARAM_SIZE)) != NULL) {
        size_t sz = 0;

        if (!OSSL_PARAM_get_size_t(p, &sz))
            return 0;
        if (sz > KMAC_MAX_OUTPUT_LEN) {
            ERR_raise(ERR_LIB_PROV, PROV_R_INVALID_OUTPUT_LENGTH);
            return 0;
        }
        kctx->out_len = sz;
    }
    if ((p = OSSL_PARAM_locate_const(params, OSSL_MAC_PARAM_KEY)) != NULL
            && !kmac_setkey(kctx, p->data, p->data_size))
        return 0;
    if ((p = OSSL_PARAM_locate_const(params, OSSL_MAC_PARAM_CUSTOM))
        != NULL) {
        if (p->data_size > KMAC_MAX_CUSTOM) {
            ERR_raise(ERR_LIB_PROV, PROV_R_INVALID_CUSTOM_LENGTH);
            return 0;
        }
        if (!encode_string(kctx->custom, sizeof(kctx->custom), &kctx->custom_len,
                           p->data, p->data_size))
            return 0;
    }
    return 1;
}




static unsigned int get_encode_size(size_t bits)
{
    unsigned int cnt = 0, sz = sizeof(size_t);

    while (bits && (cnt < sz)) {
        ++cnt;
        bits >>= 8;
    }
    
    if (cnt == 0)
        cnt = 1;
    return cnt;
}


static int right_encode(unsigned char *out, size_t out_max_len, size_t *out_len,
                        size_t bits)
{
    unsigned int len = get_encode_size(bits);
    int i;

    if (len >= out_max_len) {
        ERR_raise(ERR_LIB_PROV, PROV_R_LENGTH_TOO_LARGE);
        return 0;
    }

    
    for (i = len - 1; i >= 0; --i) {
        out[i] = (unsigned char)(bits & 0xFF);
        bits >>= 8;
    }
    
    out[len] = (unsigned char)len;

    
    *out_len = len + 1;
    return 1;
}


static int encode_string(unsigned char *out, size_t out_max_len, size_t *out_len,
                         const unsigned char *in, size_t in_len)
{
    if (in == NULL) {
        *out_len = 0;
    } else {
        size_t i, bits, len, sz;

        bits = 8 * in_len;
        len = get_encode_size(bits);
        sz = 1 + len + in_len;

        if (sz > out_max_len) {
            ERR_raise(ERR_LIB_PROV, PROV_R_LENGTH_TOO_LARGE);
            return 0;
        }

        out[0] = (unsigned char)len;
        for (i = len; i > 0; --i) {
            out[i] = (bits & 0xFF);
            bits >>= 8;
        }
        memcpy(out + len + 1, in, in_len);
        *out_len = sz;
    }
    return 1;
}


static int bytepad(unsigned char *out, size_t *out_len,
                   const unsigned char *in1, size_t in1_len,
                   const unsigned char *in2, size_t in2_len, size_t w)
{
    int len;
    unsigned char *p = out;
    int sz = w;

    if (out == NULL) {
        if (out_len == NULL) {
            ERR_raise(ERR_LIB_PROV, ERR_R_PASSED_NULL_PARAMETER);
            return 0;
        }
        sz = 2 + in1_len + (in2 != NULL ? in2_len : 0);
        *out_len = (sz + w - 1) / w * w;
        return 1;
    }

    if (!ossl_assert(w <= 255))
        return 0;

    
    *p++ = 1;
    *p++ = (unsigned char)w;
    
    memcpy(p, in1, in1_len);
    p += in1_len;
    
    if (in2 != NULL && in2_len > 0) {
        memcpy(p, in2, in2_len);
        p += in2_len;
    }
    
    len = p - out;
    sz = (len + w - 1) / w * w;
    
    if (sz != len)
        memset(p, 0, sz - len);
    if (out_len != NULL)
        *out_len = sz;
    return 1;
}


static int kmac_bytepad_encode_key(unsigned char *out, size_t out_max_len,
                                   size_t *out_len,
                                   const unsigned char *in, size_t in_len,
                                   size_t w)
{
    unsigned char tmp[KMAC_MAX_KEY + KMAC_MAX_ENCODED_HEADER_LEN];
    size_t tmp_len;

    if (!encode_string(tmp, sizeof(tmp), &tmp_len, in, in_len))
        return 0;
    if (!bytepad(NULL, out_len, tmp, tmp_len, NULL, 0, w))
        return 0;
    if (!ossl_assert(*out_len <= out_max_len))
        return 0;
    return bytepad(out, NULL, tmp, tmp_len, NULL, 0, w);
}

const OSSL_DISPATCH ossl_kmac128_functions[] = {
    { OSSL_FUNC_MAC_NEWCTX, (void (*)(void))kmac128_new },
    { OSSL_FUNC_MAC_DUPCTX, (void (*)(void))kmac_dup },
    { OSSL_FUNC_MAC_FREECTX, (void (*)(void))kmac_free },
    { OSSL_FUNC_MAC_INIT, (void (*)(void))kmac_init },
    { OSSL_FUNC_MAC_UPDATE, (void (*)(void))kmac_update },
    { OSSL_FUNC_MAC_FINAL, (void (*)(void))kmac_final },
    { OSSL_FUNC_MAC_GETTABLE_CTX_PARAMS,
      (void (*)(void))kmac_gettable_ctx_params },
    { OSSL_FUNC_MAC_GET_CTX_PARAMS, (void (*)(void))kmac_get_ctx_params },
    { OSSL_FUNC_MAC_SETTABLE_CTX_PARAMS,
      (void (*)(void))kmac_settable_ctx_params },
    { OSSL_FUNC_MAC_SET_CTX_PARAMS, (void (*)(void))kmac_set_ctx_params },
    OSSL_DISPATCH_END
};

const OSSL_DISPATCH ossl_kmac256_functions[] = {
    { OSSL_FUNC_MAC_NEWCTX, (void (*)(void))kmac256_new },
    { OSSL_FUNC_MAC_DUPCTX, (void (*)(void))kmac_dup },
    { OSSL_FUNC_MAC_FREECTX, (void (*)(void))kmac_free },
    { OSSL_FUNC_MAC_INIT, (void (*)(void))kmac_init },
    { OSSL_FUNC_MAC_UPDATE, (void (*)(void))kmac_update },
    { OSSL_FUNC_MAC_FINAL, (void (*)(void))kmac_final },
    { OSSL_FUNC_MAC_GETTABLE_CTX_PARAMS,
      (void (*)(void))kmac_gettable_ctx_params },
    { OSSL_FUNC_MAC_GET_CTX_PARAMS, (void (*)(void))kmac_get_ctx_params },
    { OSSL_FUNC_MAC_SETTABLE_CTX_PARAMS,
      (void (*)(void))kmac_settable_ctx_params },
    { OSSL_FUNC_MAC_SET_CTX_PARAMS, (void (*)(void))kmac_set_ctx_params },
    OSSL_DISPATCH_END
};
