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

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

/**
 * Converts the value of 'in' to 'outlen' bytes in big-endian byte order.
 */
VOID ull_to_bytes(UINT8 *out, UINT32 outlen,
                  UINT64 in)
{
    INTN i;

    /* Iterate over out in decreasing order, for big-endianness. */
    for (i = (INT32)outlen - 1; i >= 0; i--) {
        out[i] = in & 0xff;
        in = in >> 8;
    }
}

VOID u32_to_bytes(UINT8 *out, UINT32 in)
{
    out[0] = (UINT8)(in >> 24);
    out[1] = (UINT8)(in >> 16);
    out[2] = (UINT8)(in >> 8);
    out[3] = (UINT8)in;
}

/**
 * Converts the inlen bytes in 'in' from big-endian byte order to an integer.
 */
UINT64 bytes_to_ull(const UINT8 *in, UINT32 inlen)
{
    UINT64 retval = 0;
    UINT32 i;

    for (i = 0; i < inlen; i++) {
        retval |= ((UINT64)in[i]) << (8*(inlen - 1 - i));
    }
    return retval;
}

/**
 * Computes a root node given a leaf and an auth path.
 * Expects address to be complete other than the tree_height and tree_index.
 */
VOID compute_root(UINT8 *root, const UINT8 *leaf,
                  UINT32 leaf_idx, UINT32 idx_offset,
                  const UINT8 *auth_path, UINT32 tree_height,
                  const spx_ctx *ctx, UINT32 addr[8])
{
    UINT32 i;
    UINT8 buffer[2 * SPX_N];

    /* If leaf_idx is odd (last bit = 1), current path element is a right child
       and auth_path has to go left. Otherwise it is the other way around. */
    if (leaf_idx & 1) {
        CopyMem(buffer + SPX_N, leaf, SPX_N);
        CopyMem(buffer, auth_path, SPX_N);
    }
    else {
        CopyMem(buffer, leaf, SPX_N);
        CopyMem(buffer + SPX_N, auth_path, SPX_N);
    }
    auth_path += SPX_N;

    for (i = 0; i < tree_height - 1; i++) {
        leaf_idx >>= 1;
        idx_offset >>= 1;
        /* Set the address of the node we're creating. */
        set_tree_height(addr, i + 1);
        set_tree_index(addr, leaf_idx + idx_offset);

        /* Pick the right or left neighbor, depending on parity of the node. */
        if (leaf_idx & 1) {
            thash(buffer + SPX_N, buffer, 2, ctx, addr);
            CopyMem(buffer, auth_path, SPX_N);
        }
        else {
            thash(buffer, buffer, 2, ctx, addr);
            CopyMem(buffer + SPX_N, auth_path, SPX_N);
        }
        auth_path += SPX_N;
    }

    /* The last iteration is exceptional; we do not copy an auth_path node. */
    leaf_idx >>= 1;
    idx_offset >>= 1;
    set_tree_height(addr, tree_height);
    set_tree_index(addr, leaf_idx + idx_offset);
    thash(root, buffer, 2, ctx, addr);
}

/**
 * For a given leaf index, computes the authentication path and the resulting
 * root node using Merkle's TreeHash algorithm.
 * Expects the layer and tree parts of the tree_addr to be set, as well as the
 * tree type (i.e. SPX_ADDR_TYPE_HASHTREE or SPX_ADDR_TYPE_FORSTREE).
 * Applies the offset idx_offset to indices before building addresses, so that
 * it is possible to continue counting indices across trees.
 */
VOID treehash(UINT8 *root, UINT8 *auth_path, const spx_ctx* ctx,
              UINT32 leaf_idx, UINT32 idx_offset, UINT32 tree_height,
              VOID (*gen_leaf)(
                 UINT8* /* leaf */,
                 const spx_ctx* /* ctx */,
                 UINT32 /* addr_idx */, const UINT32[8] /* tree_addr */),
              UINT32 tree_addr[8])
{
    SPX_VLA(UINT8, stack, (tree_height+1)*SPX_N);
    SPX_VLA(UINT32, heights, tree_height+1);
    UINT32 offset = 0;
    UINT32 idx;
    UINT32 tree_idx;

    for (idx = 0; idx < (UINT32)(1 << tree_height); idx++) {
        /* Add the next leaf node to the stack. */
        gen_leaf(stack + offset*SPX_N, ctx, idx + idx_offset, tree_addr);
        offset++;
        heights[offset - 1] = 0;

        /* If this is a node we need for the auth path.. */
        if ((leaf_idx ^ 0x1) == idx) {
            CopyMem(auth_path, stack + (offset - 1)*SPX_N, SPX_N);
        }

        /* While the top-most nodes are of equal height.. */
        while (offset >= 2 && heights[offset - 1] == heights[offset - 2]) {
            /* Compute index of the new node, in the next layer. */
            tree_idx = (idx >> (heights[offset - 1] + 1));

            /* Set the address of the node we're creating. */
            set_tree_height(tree_addr, heights[offset - 1] + 1);
            set_tree_index(tree_addr,
                           tree_idx + (idx_offset >> (heights[offset-1] + 1)));
            /* Hash the top-most nodes from the stack together. */
            thash(stack + (offset - 2)*SPX_N,
                  stack + (offset - 2)*SPX_N, 2, ctx, tree_addr);
            offset--;
            /* Note that the top-most node is now one layer higher. */
            heights[offset - 1]++;

            /* If this is a node we need for the auth path.. */
            if (((leaf_idx >> heights[offset - 1]) ^ 0x1) == tree_idx) {
                CopyMem(auth_path + heights[offset - 1]*SPX_N,
                       stack + (offset - 1)*SPX_N, SPX_N);
            }
        }
    }
    CopyMem(root, stack, SPX_N);
}
