/*
 * LITLS -- AES-GCM AEAD (NIST SP 800-38D)
 * Composed from BearSSL's AES-CTR + GHASH (ctmul32).
 */

#include "bearssl/inner.h"
#include "litls.h"

/*
 * Internal GCM implementation.
 * encrypt: 1 = encrypt, 0 = decrypt.
 * Returns 0 on success, -1 on failure (decrypt tag mismatch).
 */
static int aes_gcm_crypt(const uint8_t *key, size_t key_len, const uint8_t *nonce, const uint8_t *aad, size_t aad_len,
                         const uint8_t *in, size_t in_len, uint8_t *out, uint8_t *tag, const uint8_t *expected_tag,
                         int encrypt) {
    br_aes_ct_ctr_keys ctx;
    uint8_t h[16];  /* GHASH subkey H = AES_K(0^128) */
    uint8_t j0[16]; /* J0 = nonce || 0x00000001 */
    uint8_t tag_mask[16];
    uint8_t ghash_acc[16];
    uint8_t len_block[16];

    /* Initialize AES-CTR context */
    br_aes_ct_ctr_init(&ctx, key, key_len);

    /*
     * Derive GHASH subkey: H = AES_K(0^128)
     * Use CTR mode with IV=0 and counter=0 to encrypt a zero block.
     */
    memset(h, 0, 16);
    {
        uint8_t zero_iv[12];
        memset(zero_iv, 0, 12);
        /* CTR with counter=0 encrypts: AES(IV||counter) XOR data.
         * With data=0, result = AES(0||0) = AES(0^128). */
        br_aes_ct_ctr_run(&ctx, zero_iv, 0, h, 16);
    }

    /*
     * Build J0 = nonce || 0x00000001 (for 96-bit nonce)
     */
    memcpy(j0, nonce, 12);
    j0[12] = 0x00;
    j0[13] = 0x00;
    j0[14] = 0x00;
    j0[15] = 0x01;

    /*
     * Generate tag mask: E_K(J0) = AES-CTR with nonce and counter=1.
     * BearSSL's CTR uses IV(12 bytes) + counter(4 bytes big-endian).
     * Counter=1 with our nonce gives J0.
     */
    memset(tag_mask, 0, 16);
    br_aes_ct_ctr_run(&ctx, nonce, 1, tag_mask, 16);

    /*
     * Encrypt or decrypt data using AES-CTR starting at counter=2.
     */
    if (in_len > 0) {
        memcpy(out, in, in_len);
        br_aes_ct_ctr_run(&ctx, nonce, 2, out, in_len);
    }

    /*
     * For decryption, we need to verify the tag before returning.
     * The GHASH is computed over the ciphertext (which is 'in' for decrypt).
     */
    const uint8_t *ct = encrypt ? out : in;

    /*
     * Compute GHASH over: AAD || pad || CT || pad || len(AAD)*8 || len(CT)*8
     * br_ghash_ctmul32 handles partial blocks by zero-padding internally,
     * so we must only feed full 16-byte aligned data per section. Feed
     * aligned portion, then manually pad the last partial block.
     */
    memset(ghash_acc, 0, 16);

    /* Feed AAD (aligned blocks, then padded partial block) */
    if (aad_len > 0) {
        size_t aad_full = aad_len & ~(size_t)15;
        if (aad_full > 0)
            br_ghash_ctmul32(ghash_acc, h, aad, aad_full);
        if (aad_len > aad_full) {
            uint8_t pad_block[16];
            memset(pad_block, 0, 16);
            memcpy(pad_block, aad + aad_full, aad_len - aad_full);
            br_ghash_ctmul32(ghash_acc, h, pad_block, 16);
        }
    }

    /* Feed ciphertext (aligned blocks, then padded partial block) */
    if (in_len > 0) {
        size_t ct_full = in_len & ~(size_t)15;
        if (ct_full > 0)
            br_ghash_ctmul32(ghash_acc, h, ct, ct_full);
        if (in_len > ct_full) {
            uint8_t pad_block[16];
            memset(pad_block, 0, 16);
            memcpy(pad_block, ct + ct_full, in_len - ct_full);
            br_ghash_ctmul32(ghash_acc, h, pad_block, 16);
        }
    }

    /* Feed length block: len(AAD)*8 || len(CT)*8, both as big-endian 64-bit */
    br_enc64be(len_block, (uint64_t)aad_len << 3);
    br_enc64be(len_block + 8, (uint64_t)in_len << 3);
    br_ghash_ctmul32(ghash_acc, h, len_block, 16);

    /* Tag = GHASH XOR E_K(J0) */
    for (int i = 0; i < 16; i++)
        ghash_acc[i] ^= tag_mask[i];

    if (encrypt) {
        memcpy(tag, ghash_acc, 16);
    } else {
        /* Verify tag */
        if (litls_secure_memcmp(ghash_acc, expected_tag, 16) != 0) {
            /* Wipe output on failure */
            memset(out, 0, in_len);
            return -1;
        }
    }

    return 0;
}

int litls_aes128_gcm_encrypt(const uint8_t key[16], const uint8_t nonce[12], const uint8_t *aad, size_t aad_len,
                             const uint8_t *plaintext, size_t pt_len, uint8_t *ciphertext, uint8_t tag[16]) {
    return aes_gcm_crypt(key, 16, nonce, aad, aad_len, plaintext, pt_len, ciphertext, tag, NULL, 1);
}

int litls_aes128_gcm_decrypt(const uint8_t key[16], const uint8_t nonce[12], const uint8_t *aad, size_t aad_len,
                             const uint8_t *ciphertext, size_t ct_len, const uint8_t tag[16], uint8_t *plaintext) {
    return aes_gcm_crypt(key, 16, nonce, aad, aad_len, ciphertext, ct_len, plaintext, NULL, tag, 0);
}

int litls_aes256_gcm_encrypt(const uint8_t key[32], const uint8_t nonce[12], const uint8_t *aad, size_t aad_len,
                             const uint8_t *plaintext, size_t pt_len, uint8_t *ciphertext, uint8_t tag[16]) {
    return aes_gcm_crypt(key, 32, nonce, aad, aad_len, plaintext, pt_len, ciphertext, tag, NULL, 1);
}

int litls_aes256_gcm_decrypt(const uint8_t key[32], const uint8_t nonce[12], const uint8_t *aad, size_t aad_len,
                             const uint8_t *ciphertext, size_t ct_len, const uint8_t tag[16], uint8_t *plaintext) {
    return aes_gcm_crypt(key, 32, nonce, aad, aad_len, ciphertext, ct_len, plaintext, NULL, tag, 0);
}