

#include <string.h>
#include "internal/sha3.h"

void SHA3_squeeze(uint64_t A[5][5], unsigned char *out, size_t len, size_t r, int next);

void ossl_sha3_reset(KECCAK1600_CTX *ctx)
{
    memset(ctx->A, 0, sizeof(ctx->A));
    ctx->bufsz = 0;
    ctx->xof_state = XOF_STATE_INIT;
}

int ossl_sha3_init(KECCAK1600_CTX *ctx, unsigned char pad, size_t bitlen)
{
    size_t bsz = SHA3_BLOCKSIZE(bitlen);

    if (bsz <= sizeof(ctx->buf)) {
        ossl_sha3_reset(ctx);
        ctx->block_size = bsz;
        ctx->md_size = bitlen / 8;
        ctx->pad = pad;
        return 1;
    }

    return 0;
}

int ossl_keccak_kmac_init(KECCAK1600_CTX *ctx, unsigned char pad, size_t bitlen)
{
    int ret = ossl_sha3_init(ctx, pad, bitlen);

    if (ret)
        ctx->md_size *= 2;
    return ret;
}

int ossl_sha3_update(KECCAK1600_CTX *ctx, const void *_inp, size_t len)
{
    const unsigned char *inp = _inp;
    size_t bsz = ctx->block_size;
    size_t num, rem;

    if (len == 0)
        return 1;

    if (ctx->xof_state == XOF_STATE_SQUEEZE
        || ctx->xof_state == XOF_STATE_FINAL)
        return 0;

    if ((num = ctx->bufsz) != 0) {      
        rem = bsz - num;

        if (len < rem) {
            memcpy(ctx->buf + num, inp, len);
            ctx->bufsz += len;
            return 1;
        }
        
        memcpy(ctx->buf + num, inp, rem);
        inp += rem, len -= rem;
        (void)SHA3_absorb(ctx->A, ctx->buf, bsz, bsz);
        ctx->bufsz = 0;
        
    }

    if (len >= bsz)
        rem = SHA3_absorb(ctx->A, inp, len, bsz);
    else
        rem = len;

    if (rem) {
        memcpy(ctx->buf, inp + len - rem, rem);
        ctx->bufsz = rem;
    }

    return 1;
}


int ossl_sha3_final(KECCAK1600_CTX *ctx, unsigned char *out, size_t outlen)
{
    size_t bsz = ctx->block_size;
    size_t num = ctx->bufsz;

    if (outlen == 0)
        return 1;
    if (ctx->xof_state == XOF_STATE_SQUEEZE
        || ctx->xof_state == XOF_STATE_FINAL)
        return 0;

    
    memset(ctx->buf + num, 0, bsz - num);
    ctx->buf[num] = ctx->pad;
    ctx->buf[bsz - 1] |= 0x80;

    (void)SHA3_absorb(ctx->A, ctx->buf, bsz, bsz);

    ctx->xof_state = XOF_STATE_FINAL;
    SHA3_squeeze(ctx->A, out, outlen, bsz, 0);
    return 1;
}


int ossl_sha3_squeeze(KECCAK1600_CTX *ctx, unsigned char *out, size_t outlen)
{
    size_t bsz = ctx->block_size;
    size_t num = ctx->bufsz;
    size_t len;
    int next = 1;

    if (outlen == 0)
        return 1;

    if (ctx->xof_state == XOF_STATE_FINAL)
        return 0;

    
    if (ctx->xof_state != XOF_STATE_SQUEEZE) {
        
        memset(ctx->buf + num, 0, bsz - num);
        ctx->buf[num] = ctx->pad;
        ctx->buf[bsz - 1] |= 0x80;
        (void)SHA3_absorb(ctx->A, ctx->buf, bsz, bsz);
        ctx->xof_state = XOF_STATE_SQUEEZE;
        num = ctx->bufsz = 0;
        next = 0;
    }

    
    if (num != 0) {
        if (outlen > ctx->bufsz)
            len = ctx->bufsz;
        else
            len = outlen;
        memcpy(out, ctx->buf + bsz - ctx->bufsz, len);
        out += len;
        outlen -= len;
        ctx->bufsz -= len;
    }
    if (outlen == 0)
        return 1;

    
    if (outlen >= bsz) {
        len = bsz * (outlen / bsz);
        SHA3_squeeze(ctx->A, out, len, bsz, next);
        next = 1;
        out += len;
        outlen -= len;
    }
    if (outlen > 0) {
        
        SHA3_squeeze(ctx->A, ctx->buf, bsz, bsz, next);
        memcpy(out, ctx->buf, outlen);
        
        ctx->bufsz = bsz - outlen;
    }

    return 1;
}
