/**
 * @file se050_hkdf_blake2s.c
 * @brief HKDF Implementation using HMAC-BLAKE2s (RFC 5861)
 */

#include "se050_hkdf_blake2s.h"
#include "se050_hmac_blake2s.h"
#include <string.h>

#define HMAC_BLAKE2S_DIGEST_SIZE 32

int se050_hkdf_extract(uint8_t prk[32],
                       const uint8_t *salt, size_t saltlen,
                       const uint8_t *ikm, size_t ikmlen)
{
    uint8_t default_salt[HMAC_BLAKE2S_DIGEST_SIZE];

    if (!prk || !ikm) {
        return -1;
    }

    if (!salt || saltlen == 0) {
        memset(default_salt, 0, HMAC_BLAKE2S_DIGEST_SIZE);
        salt = default_salt;
        saltlen = HMAC_BLAKE2S_DIGEST_SIZE;
    }

    return se050_hmac_blake2s(prk, salt, saltlen, ikm, ikmlen);
}

int se050_hkdf_expand(uint8_t *okm, size_t okmlen,
                      const uint8_t prk[32],
                      const uint8_t *info, size_t infolen)
{
    uint8_t t[HMAC_BLAKE2S_DIGEST_SIZE];
    size_t n;
    int ret;

    if (!okm || !prk || okmlen == 0 || okmlen > HKDF_BLAKE2S_MAX_OUTPUT) {
        return -1;
    }

    n = (okmlen + HMAC_BLAKE2S_DIGEST_SIZE - 1) / HMAC_BLAKE2S_DIGEST_SIZE;

    memset(t, 0, sizeof(t));

    for (size_t i = 1; i <= n; i++) {
        size_t data_len = (i == 1) ? 0 : HMAC_BLAKE2S_DIGEST_SIZE;
        
        ret = se050_hmac_blake2s(t, prk, HMAC_BLAKE2S_DIGEST_SIZE,
                                 t, data_len);
        if (ret != 0) {
            return ret;
        }

        if (info && infolen > 0) {
            ret = se050_hmac_blake2s_variable(t, HMAC_BLAKE2S_DIGEST_SIZE,
                                              prk, HMAC_BLAKE2S_DIGEST_SIZE,
                                              info, infolen);
            if (ret != 0) {
                return ret;
            }
        }

        size_t block_len = (i < n) ? HMAC_BLAKE2S_DIGEST_SIZE :
                           (okmlen - (i - 1) * HMAC_BLAKE2S_DIGEST_SIZE);
        memcpy(okm + (i - 1) * HMAC_BLAKE2S_DIGEST_SIZE, t, block_len);
    }

    memset(t, 0, sizeof(t));
    return 0;
}

int se050_hkdf(uint8_t *okm, size_t okmlen,
               const uint8_t *salt, size_t saltlen,
               const uint8_t *ikm, size_t ikmlen,
               const uint8_t *info, size_t infolen)
{
    uint8_t prk[HMAC_BLAKE2S_DIGEST_SIZE];
    int ret;

    ret = se050_hkdf_extract(prk, salt, saltlen, ikm, ikmlen);
    if (ret != 0) {
        return ret;
    }

    ret = se050_hkdf_expand(okm, okmlen, prk, info, infolen);
    memset(prk, 0, sizeof(prk));
    return ret;
}