/*
 * LITLS -- ChaCha20 IETF variant (RFC 8439)
 * 256-bit key, 96-bit nonce, 32-bit counter.
 */

#include "litls.h"
#include <string.h>

#define ROTL32(v, n) (((v) << (n)) | ((v) >> (32 - (n))))

#define QR(a, b, c, d) \
    a += b;            \
    d ^= a;            \
    d = ROTL32(d, 16); \
    c += d;            \
    b ^= c;            \
    b = ROTL32(b, 12); \
    a += b;            \
    d ^= a;            \
    d = ROTL32(d, 8);  \
    c += d;            \
    b ^= c;            \
    b = ROTL32(b, 7);

static void chacha20_block_core(uint32_t out[16], const uint32_t in[16]) {
    uint32_t x[16];
    memcpy(x, in, 64);

    for (int i = 0; i < 10; i++) {
        /* column rounds */
        QR(x[0], x[4], x[8], x[12])
        QR(x[1], x[5], x[9], x[13])
        QR(x[2], x[6], x[10], x[14])
        QR(x[3], x[7], x[11], x[15])
        /* diagonal rounds */
        QR(x[0], x[5], x[10], x[15])
        QR(x[1], x[6], x[11], x[12])
        QR(x[2], x[7], x[8], x[13])
        QR(x[3], x[4], x[9], x[14])
    }

    for (int i = 0; i < 16; i++)
        out[i] = x[i] + in[i];
}

static void u32_to_le(uint8_t *out, uint32_t v) {
    out[0] = (uint8_t)(v);
    out[1] = (uint8_t)(v >> 8);
    out[2] = (uint8_t)(v >> 16);
    out[3] = (uint8_t)(v >> 24);
}

static uint32_t le_to_u32(const uint8_t *in) {
    return (uint32_t)in[0] | ((uint32_t)in[1] << 8) | ((uint32_t)in[2] << 16) | ((uint32_t)in[3] << 24);
}

static void setup_state(uint32_t state[16], const uint8_t key[32], const uint8_t nonce[12], uint32_t counter) {
    /* "expand 32-byte k" */
    state[0] = 0x61707865;
    state[1] = 0x3320646e;
    state[2] = 0x79622d32;
    state[3] = 0x6b206574;
    /* key */
    for (int i = 0; i < 8; i++)
        state[4 + i] = le_to_u32(key + i * 4);
    /* counter */
    state[12] = counter;
    /* nonce */
    for (int i = 0; i < 3; i++)
        state[13 + i] = le_to_u32(nonce + i * 4);
}

void litls_chacha20_block(const uint8_t key[32], const uint8_t nonce[12], uint32_t counter, uint8_t out[64]) {
    uint32_t state[16], result[16];
    setup_state(state, key, nonce, counter);
    chacha20_block_core(result, state);
    for (int i = 0; i < 16; i++)
        u32_to_le(out + i * 4, result[i]);
}

void litls_chacha20_encrypt(const uint8_t key[32], const uint8_t nonce[12], uint32_t counter, const uint8_t *in,
                            size_t len, uint8_t *out) {
    uint32_t state[16];
    setup_state(state, key, nonce, counter);

    while (len > 0) {
        uint32_t result[16];
        chacha20_block_core(result, state);

        uint8_t keystream[64];
        for (int i = 0; i < 16; i++)
            u32_to_le(keystream + i * 4, result[i]);

        size_t block_len = len < 64 ? len : 64;
        for (size_t i = 0; i < block_len; i++)
            out[i] = in[i] ^ keystream[i];

        state[12]++;
        in += block_len;
        out += block_len;
        len -= block_len;
    }
}