/*
 * Copyright Supranational LLC
 * Licensed under the Apache License, Version 2.0, see LICENSE for details.
 * SPDX-License-Identifier: Apache-2.0
 */

#include "consts.h"
#include "sha256.h"

static const vec384 BLS12_381_RRRR = {  /* RR^2 */
    TO_LIMB_T(0xed48ac6bd94ca1e0), TO_LIMB_T(0x315f831e03a7adf8),
    TO_LIMB_T(0x9a53352a615e29dd), TO_LIMB_T(0x34c04e5e921e1761),
    TO_LIMB_T(0x2512d43565724728), TO_LIMB_T(0x0aa6346091755d4d)
};

#ifdef expand_message_xmd
void expand_message_xmd(unsigned char *bytes, size_t len_in_bytes,
                        const unsigned char *aug, size_t aug_len,
                        const unsigned char *msg, size_t msg_len,
                        const unsigned char *DST, size_t DST_len);
#else
static void sha256_init_Zpad(SHA256_CTX *ctx)
{
    ctx->h[0] = 0xda5698beU;
    ctx->h[1] = 0x17b9b469U;
    ctx->h[2] = 0x62335799U;
    ctx->h[3] = 0x779fbecaU;
    ctx->h[4] = 0x8ce5d491U;
    ctx->h[5] = 0xc0d26243U;
    ctx->h[6] = 0xbafef9eaU;
    ctx->h[7] = 0x1837a9d8U;
    ctx->N = 64;
    vec_zero(ctx->buf, sizeof(ctx->buf));
    ctx->off = 0;
}

static void vec_xor(void *restrict ret, const void *restrict a,
                                        const void *restrict b, size_t num)
{
    limb_t *rp = (limb_t *)ret;
    const limb_t *ap = (const limb_t *)a;
    const limb_t *bp = (const limb_t *)b;
    size_t i;

    num /= sizeof(limb_t);

    for (i = 0; i < num; i++)
        rp[i] = ap[i] ^ bp[i];
}

static void expand_message_xmd(unsigned char *bytes, size_t len_in_bytes,
                               const unsigned char *aug, size_t aug_len,
                               const unsigned char *msg, size_t msg_len,
                               const unsigned char *DST, size_t DST_len)
{
    union { limb_t align; unsigned char c[32]; } b_0;
    union { limb_t align; unsigned char c[33+256+31]; } b_i;
    unsigned char *p;
    size_t i, b_i_bits, b_i_blocks;
    SHA256_CTX ctx;

    /*
     * compose template for 'strxor(b_0, b_(i-1)) || I2OSP(i, 1) || DST_prime'
     */
    if (DST_len > 255) {
        sha256_init(&ctx);
        sha256_update(&ctx, "H2C-OVERSIZE-DST-", 17);
        sha256_update(&ctx, DST, DST_len);
        sha256_final(b_0.c, &ctx);
        DST = b_0.c, DST_len = 32;
    }
    b_i_blocks = ((33 + DST_len + 1 + 9) + 63) & -64;
    vec_zero(b_i.c + b_i_blocks - 64, 64);

    p = b_i.c + 33;
    for (i = 0; i < DST_len; i++)
        p[i] = DST[i];
    p[i++] = (unsigned char)DST_len;
    p[i++] = 0x80;
    p[i+6] = p[i+5] = p[i+4] = p[i+3] = p[i+2] = p[i+1] = p[i+0] = 0;
    b_i_bits = (33 + DST_len + 1) * 8;
    p = b_i.c + b_i_blocks;
    p[-2] = (unsigned char)(b_i_bits >> 8);
    p[-1] = (unsigned char)(b_i_bits);

    sha256_init_Zpad(&ctx);                         /* Z_pad | */
    sha256_update(&ctx, aug, aug_len);              /* | aug | */
    sha256_update(&ctx, msg, msg_len);              /* | msg | */
    /* | I2OSP(len_in_bytes, 2) || I2OSP(0, 1) || DST_prime    */
    b_i.c[30] = (unsigned char)(len_in_bytes >> 8);
    b_i.c[31] = (unsigned char)(len_in_bytes);
    b_i.c[32] = 0;
    sha256_update(&ctx, b_i.c + 30, 3 + DST_len + 1);
    sha256_final(b_0.c, &ctx);

    sha256_init_h(ctx.h);
    vec_copy(b_i.c, b_0.c, 32);
    ++b_i.c[32];
    sha256_block_data_order(ctx.h, b_i.c, b_i_blocks / 64);
    sha256_emit(bytes, ctx.h);

    len_in_bytes += 31; /* ell = ceil(len_in_bytes / b_in_bytes), with */
    len_in_bytes /= 32; /* caller being responsible for accordingly large
                         * buffer. hash_to_field passes one with length
                         * divisible by 64, remember? which works... */
    while (--len_in_bytes) {
        sha256_init_h(ctx.h);
        vec_xor(b_i.c, b_0.c, bytes, 32);
        bytes += 32;
        ++b_i.c[32];
        sha256_block_data_order(ctx.h, b_i.c, b_i_blocks / 64);
        sha256_emit(bytes, ctx.h);
    }
}
#endif

/*
 * |nelems| is 'count * m' from spec
 */
static void hash_to_field(vec384 elems[], size_t nelems,
                          const unsigned char *aug, size_t aug_len,
                          const unsigned char *msg, size_t msg_len,
                          const unsigned char *DST, size_t DST_len)
{
    size_t L = sizeof(vec384) + 128/8;  /* ceil((ceil(log2(p)) + k) / 8) */
    size_t len_in_bytes = L * nelems;   /* divisible by 64, hurray!      */
#if !defined(__STDC_VERSION__) || __STDC_VERSION__<199901 \
                               || defined(__STDC_NO_VLA__)
    limb_t *pseudo_random = alloca(len_in_bytes);
#else
    limb_t pseudo_random[len_in_bytes/sizeof(limb_t)];
#endif
    unsigned char *bytes;
    vec768 elem;

    aug_len = aug!=NULL ? aug_len : 0;
    DST_len = DST!=NULL ? DST_len : 0;

    expand_message_xmd((unsigned char *)pseudo_random, len_in_bytes,
                       aug, aug_len, msg, msg_len, DST, DST_len);

    vec_zero(elem, sizeof(elem));
    bytes = (unsigned char *)pseudo_random;
    while (nelems--) {
        limbs_from_be_bytes(elem, bytes, L);
        bytes += L;
        /*
         * L-bytes block % P, output is in Montgomery domain...
         */
        redc_mont_384(elems[0], elem, BLS12_381_P, p0);
        mul_mont_384(elems[0], elems[0], BLS12_381_RRRR, BLS12_381_P, p0);
        elems++;
    }
}

void blst_expand_message_xmd(unsigned char *bytes, size_t len_in_bytes,
                             const unsigned char *msg, size_t msg_len,
                             const unsigned char *DST, size_t DST_len)
{
    size_t buf_len = (len_in_bytes+31) & ((size_t)0-32);
    unsigned char *buf_ptr = bytes;

    if (buf_len > 255*32)
        return;

    if (buf_len != len_in_bytes)
        buf_ptr = alloca(buf_len);

    expand_message_xmd(buf_ptr, len_in_bytes, NULL, 0, msg, msg_len,
                                              DST, DST_len);
    if (buf_ptr != bytes) {
        unsigned char *ptr = buf_ptr;
        while (len_in_bytes--)
            *bytes++ = *ptr++;
        vec_zero(buf_ptr, buf_len);
    }
}