#include <Library/UefiLib.h>
#include <Library/BaseLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/BaseMemoryLib.h>
#include <Uefi.h>

#include "thash.h"
#include "address.h"
#include "params.h"
#include "utils.h"
#include "sha2.h"

#if SPX_SHA512
static VOID thash_512(UINT8 *out, const UINT8 *in, UINT32 inblocks,
           const spx_ctx *ctx, UINT32 addr[8]);
#endif

/**
 * Takes an array of inblocks concatenated arrays of SPX_N bytes.
 */
VOID thash(UINT8 *out, const UINT8 *in, UINT32 inblocks,
           const spx_ctx *ctx, UINT32 addr[8])
{
#if SPX_SHA512
    if (inblocks > 1) {
	thash_512(out, in, inblocks, ctx, addr);
        return;
    }
#endif
    UINT8 outbuf[SPX_SHA256_OUTPUT_BYTES];
    SPX_VLA(UINT8, bitmask, inblocks * SPX_N);
    SPX_VLA(UINT8, buf, SPX_N + SPX_SHA256_OUTPUT_BYTES + inblocks*SPX_N);
    UINT8 sha2_state[40];
    UINT32 i;

    CopyMem(buf, ctx->pub_seed, SPX_N);
    CopyMem(buf + SPX_N, addr, SPX_SHA256_ADDR_BYTES);
    mgf1_256(bitmask, inblocks * SPX_N, buf, SPX_N + SPX_SHA256_ADDR_BYTES);

    /* Retrieve precomputed state containing pub_seed */
    CopyMem(sha2_state, ctx->state_seeded, 40 * sizeof(UINT8));

    for (i = 0; i < inblocks * SPX_N; i++) {
        buf[SPX_N + SPX_SHA256_ADDR_BYTES + i] = in[i] ^ bitmask[i];
    }

    sha256_inc_finalize(outbuf, sha2_state, buf + SPX_N,
                        SPX_SHA256_ADDR_BYTES + inblocks*SPX_N);
    CopyMem(out, outbuf, SPX_N);
}

#if SPX_SHA512
static VOID thash_512(UINT8 *out, const UINT8 *in, UINT32 inblocks,
           const spx_ctx *ctx, UINT32 addr[8])
{
    UINT8 outbuf[SPX_SHA512_OUTPUT_BYTES];
    SPX_VLA(UINT8, bitmask, inblocks * SPX_N);
    SPX_VLA(UINT8, buf, SPX_N + SPX_SHA256_ADDR_BYTES + inblocks*SPX_N);
    UINT8 sha2_state[72];
    UINT32 i;

    CopyMem(buf, ctx->pub_seed, SPX_N);
    CopyMem(buf + SPX_N, addr, SPX_SHA256_ADDR_BYTES);
    mgf1_512(bitmask, inblocks * SPX_N, buf, SPX_N + SPX_SHA256_ADDR_BYTES);

    /* Retrieve precomputed state containing pub_seed */
    CopyMem(sha2_state, ctx->state_seeded_512, 72 * sizeof(UINT8));

    for (i = 0; i < inblocks * SPX_N; i++) {
        buf[SPX_N + SPX_SHA256_ADDR_BYTES + i] = in[i] ^ bitmask[i];
    }

    sha512_inc_finalize(outbuf, sha2_state, buf + SPX_N,
                        SPX_SHA256_ADDR_BYTES + inblocks*SPX_N);
    CopyMem(out, outbuf, SPX_N);
}
#endif
