/*!

 * Copyrights

 *

 * Portions created or assigned to Cisco Systems, Inc. are

 * Copyright (c) 2018 Cisco Systems, Inc.  All Rights Reserved.

 */

#include "include/concatkdf_int.h"

#include "include/util_int.h"

#ifdef _WIN32

#include <Winsock2.h>

#else

#include <arpa/inet.h>

#endif

#include <openssl/evp.h>

#include <string.h>

#include <cjose/base64.h>

#include <cjose/util.h>

////////////////////////////////////////////////////////////////////////////////

static uint8_t *_apply_uint32(const uint32_t value, uint8_t *buffer)

{

    const uint32_t big_endian_int32 = htonl(value);

    memcpy(buffer, &big_endian_int32, 4);

    return buffer + 4;

}

static uint8_t *_apply_lendata(const uint8_t *data, const size_t len, uint8_t *buffer)

{

    uint8_t *ptr = buffer;

    ptr = _apply_uint32(len, ptr);

    if (0 < len)

    {

        memcpy(ptr, data, len);

        ptr += len;

    }

    return ptr;

}

size_t min_len(size_t a, size_t b) { return (a < b) ? a : b; }

////////////////////////////////////////////////////////////////////////////////

bool cjose_concatkdf_create_otherinfo(

    const char *alg, const size_t keylen, cjose_header_t *hdr, uint8_t **otherinfo, size_t *otherinfoLen, cjose_err *err)

{

    bool result = false;

    uint8_t *apu = NULL, *apv = NULL;

    size_t apuLen = 0, apvLen = 0;

    // err is optional and may be NULL, so only dereference it when provided.

    // cjose_header_get() records an error only for an invalid header/attr; for a

    // valid hdr and the constant APU/APV attrs an absent field just yields NULL.

    if (NULL != err)

    {

        memset(err, 0, sizeof(cjose_err));

    }

    const char *apuStr = cjose_header_get(hdr, CJOSE_HDR_APU, err);

    const char *apvStr = cjose_header_get(hdr, CJOSE_HDR_APV, err);

    if (NULL != err && CJOSE_ERR_NONE != err->code)

    {

        return false;

    }

    apuLen = (NULL != apuStr) ? strlen(apuStr) : 0;

    if (apuStr != NULL && !cjose_base64url_decode(apuStr, apuLen, &apu, &apuLen, err))

    {

        goto concatkdf_create_otherinfo_finish;

    }

    apvLen = (NULL != apvStr) ? strlen(apvStr) : 0;

    if (apvStr != NULL && !cjose_base64url_decode(apvStr, apvLen, &apv, &apvLen, err))

    {

        goto concatkdf_create_otherinfo_finish;

    }

    const size_t algLen = strlen(alg);

    const size_t bufferLen = (4 + algLen) + (4 + apuLen) + (4 + apvLen) + 4;

    uint8_t *buffer = cjose_get_alloc()(bufferLen);

    if (NULL == buffer)

    {

        CJOSE_ERROR(err, CJOSE_ERR_NO_MEMORY);

        goto concatkdf_create_otherinfo_finish;

    }

    uint8_t *ptr = buffer;

    ptr = _apply_lendata((const uint8_t *)alg, algLen, ptr);

    ptr = _apply_lendata(apu, apuLen, ptr);

    ptr = _apply_lendata(apv, apvLen, ptr);

    // final write; the returned (end) pointer is intentionally not stored

    _apply_uint32(keylen, ptr);

    *otherinfoLen = bufferLen;

    *otherinfo = buffer;

    result = true;

concatkdf_create_otherinfo_finish:

    cjose_get_dealloc()(apu);

    cjose_get_dealloc()(apv);

    return result;

}

////////////////////////////////////////////////////////////////////////////////

uint8_t *cjose_concatkdf_derive(const size_t keylen,

                                const uint8_t *ikm,

                                const size_t ikmLen,

                                const uint8_t *otherinfo,

                                const size_t otherinfoLen,

                                cjose_err *err)

{

    uint8_t *derived = NULL;

    uint8_t *buffer = NULL;

    const EVP_MD *dgst = EVP_sha256();

    EVP_MD_CTX *ctx = EVP_MD_CTX_create();

    if (NULL == ctx)

    {

        CJOSE_ERROR(err, CJOSE_ERR_NO_MEMORY);

        goto concatkdf_derive_finish;

    }

    const size_t hashlen = EVP_MD_size(dgst);

    const size_t N = (keylen + hashlen - 1) / hashlen;

    buffer = cjose_get_alloc()(keylen);

    if (NULL == buffer)

    {

        CJOSE_ERROR(err, CJOSE_ERR_NO_MEMORY);

        goto concatkdf_derive_finish;

    }

    size_t offset = 0;

    for (int idx = 1; N >= idx; idx++)

    {

        uint8_t counter[4];

        _apply_uint32(idx, counter);

        uint8_t *hash = cjose_get_alloc()(hashlen * sizeof(uint8_t));

        if (NULL == hash)

        {

            CJOSE_ERROR(err, CJOSE_ERR_NO_MEMORY);

            goto concatkdf_derive_finish;

        }

        if (1 != EVP_DigestInit_ex(ctx, dgst, NULL) || 1 != EVP_DigestUpdate(ctx, counter, sizeof(counter))

            || 1 != EVP_DigestUpdate(ctx, ikm, ikmLen) || 1 != EVP_DigestUpdate(ctx, otherinfo, otherinfoLen)

            || 1 != EVP_DigestFinal_ex(ctx, hash, NULL))

        {

            _cjose_cleanse_dealloc(hash, hashlen);

            CJOSE_ERROR(err, CJOSE_ERR_CRYPTO);

            goto concatkdf_derive_finish;

        }

        // copy this digest block into the derived key; the final block may be

        // shorter than hashlen. offset stays < keylen on every iteration, so the

        // remaining count (keylen - offset) cannot underflow.

        // hash holds derived key material; wipe it before returning to the allocator

        size_t amt = keylen - offset;

        memcpy(buffer + offset, hash, min_len(hashlen, amt));

        _cjose_cleanse_dealloc(hash, hashlen);

        offset += hashlen;

    }

    derived = buffer;

    buffer = NULL;

concatkdf_derive_finish:

    EVP_MD_CTX_destroy(ctx);

    _cjose_cleanse_dealloc(buffer, keylen);

    return derived;

}