#include "aes.h"

// ============================================================
// Hardware acceleration: AES-NI + PCLMULQDQ (x86/x64 only)
// Falls back to software on non-x86 or CPUs without AES-NI.
// Works on GCC, Clang, and MSVC.
// ============================================================

#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)
#define NETPLUS_X86 1

#ifdef _MSC_VER
  #include <intrin.h>
#else
  #include <cpuid.h>
#endif
#include <wmmintrin.h>
#include <tmmintrin.h>
#include <smmintrin.h>

#if defined(__GNUC__) || defined(__clang__)
  // sse4.1 (for _mm_insert_epi32 in make_ctr_block) is safe to assume
  // alongside aes/pclmul: every shipped x86 CPU with AES-NI (Westmere/
  // Bulldozer and later) also has SSE4.1/4.2 — there is no real-world AES-NI
  // CPU lacking it.
  #define AESNI_FUNC __attribute__((target("aes,ssse3,sse4.1,pclmul")))
#else
  #define AESNI_FUNC
#endif

static inline std::uint32_t byteSwap32(std::uint32_t value) noexcept
{
#if defined(_MSC_VER)
    return _byteswap_ulong(value);
#elif defined(__GNUC__) || defined(__clang__)
    return __builtin_bswap32(value);
#else
    return ((value & 0x000000FFu) << 24) |
           ((value & 0x0000FF00u) << 8)  |
           ((value & 0x00FF0000u) >> 8)  |
           ((value & 0xFF000000u) >> 24);
#endif
}

// F16 (performance report): deliberately NOT switched to VAES/AVX-512
// (VAES+VPCLMULQDQ, 2-4 AES blocks per instruction instead of 1), even
// though the report sees a further, significant multiple there over the
// AES-NI path already implemented here (128-bit, 4-block interleaving
// with fused GHASH — see aesni_gcm_ctr_ghash_128/256 and their fused
// CTR+GHASH decrypt variant). Reasons: (1) VAES is only available from Ice
// Lake/Zen 3 onward — a separate, CPUID-detected runtime code path would
// be needed that never triggers on older/virtualized hardware and would
// only be carried along there as an extra, untested path; (2) this
// sandbox has no VAES-capable CPU, so the path could not even be
// functionally verified here (let alone benchmarked) — a bare "should
// work per the Intel manual" would not be sufficient verification for a
// core AEAD path in a crypto library. The existing 128-bit/4-block path
// remains the baseline fallback and covers every AES-NI-capable CPU.

static bool detect_aesni_pclmul() {
#ifdef _MSC_VER
    int info[4];
    __cpuid(info, 1);
    bool aes    = (info[2] & (1 << 25)) != 0;
    bool pclmul = (info[2] & (1 << 1))  != 0;
    return aes && pclmul;
#else
    unsigned eax, ebx, ecx, edx;
    if (__get_cpuid(1, &eax, &ebx, &ecx, &edx)) {
        bool aes    = (ecx & (1 << 25)) != 0;
        bool pclmul = (ecx & (1 << 1))  != 0;
        return aes && pclmul;
    }
    return false;
#endif
}

static const bool g_has_hw_aes = detect_aesni_pclmul();

// ---- AES-NI encrypt/decrypt for AES-128 (10 rounds) ----

AESNI_FUNC
static void aesni_encrypt_128(const uint8_t in[16], uint8_t out[16],
                               const uint8_t* round_keys) {
    __m128i block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in));
    block = _mm_xor_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    for (int i = 1; i < 10; i++)
        block = _mm_aesenc_si128(block,
                    _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
    block = _mm_aesenclast_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));
    _mm_storeu_si128(reinterpret_cast<__m128i*>(out), block);
}

AESNI_FUNC
static void aesni_decrypt_128(const uint8_t in[16], uint8_t out[16],
                               const uint8_t* round_keys) {
    __m128i block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in));
    block = _mm_xor_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));
    for (int i = 9; i >= 1; i--) {
        __m128i rk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16));
        block = _mm_aesdec_si128(block, _mm_aesimc_si128(rk));
    }
    block = _mm_aesdeclast_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    _mm_storeu_si128(reinterpret_cast<__m128i*>(out), block);
}

// ---- AES-NI encrypt/decrypt for AES-256 (14 rounds) ----

AESNI_FUNC
static void aesni_encrypt_256(const uint8_t in[16], uint8_t out[16],
                               const uint8_t* round_keys) {
    __m128i block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in));
    block = _mm_xor_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    for (int i = 1; i < 14; i++)
        block = _mm_aesenc_si128(block,
                    _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
    block = _mm_aesenclast_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));
    _mm_storeu_si128(reinterpret_cast<__m128i*>(out), block);
}

AESNI_FUNC
static void aesni_decrypt_256(const uint8_t in[16], uint8_t out[16],
                               const uint8_t* round_keys) {
    __m128i block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in));
    block = _mm_xor_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));
    for (int i = 13; i >= 1; i--) {
        __m128i rk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16));
        block = _mm_aesdec_si128(block, _mm_aesimc_si128(rk));
    }
    block = _mm_aesdeclast_si128(block,
                _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    _mm_storeu_si128(reinterpret_cast<__m128i*>(out), block);
}

// ---- Pipelined AES-NI CTR mode (4 blocks in parallel) ----
// CTR blocks are independent, so we can keep the AES pipeline full
// by issuing 4 encryptions concurrently.

AESNI_FUNC
static inline __m128i make_ctr_block(const uint8_t base[16], uint32_t counter) {
    // Replace the last 4 bytes (lane 3) with the big-endian encoding of
    // `counter`, entirely in registers — no store/reload round trip through
    // memory. Byte-swapping `counter` first and inserting it as a native
    // (little-endian) 32-bit lane produces exactly the big-endian byte
    // sequence in memory order that the stack-buffer version built by hand.
    __m128i blk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(base));
    return _mm_insert_epi32(blk, static_cast<int>(byteSwap32(counter)), 3);
}

// ---- PCLMULQDQ-based GCM multiplication in GF(2^128) ----
// Uses full 128-bit reflection to convert between GCM bit-order
// and CLMUL standard convention, then standard polynomial reduction.

AESNI_FUNC
static inline __m128i reflect128(__m128i v) {
    // Reverse bits within each byte (NOT byte-swap).
    // GCM stores coefficient of x^(8j+k) at bit (7-k) of byte j.
    // CLMUL wants coefficient of x^(8j+k) at bit k of byte j.
    __m128i lo_mask = _mm_set1_epi8(0x0F);
    __m128i rev_lut = _mm_setr_epi8(0,8,4,12,2,10,6,14,1,9,5,13,3,11,7,15);
    __m128i hi = _mm_shuffle_epi8(rev_lut, _mm_and_si128(_mm_srli_epi16(v, 4), lo_mask));
    __m128i lo = _mm_shuffle_epi8(rev_lut, _mm_and_si128(v, lo_mask));
    return _mm_or_si128(_mm_slli_epi16(lo, 4), hi);
}

AESNI_FUNC
static inline __m128i shl128(__m128i v, int k) {
    return _mm_or_si128(
        _mm_slli_epi64(v, k),
        _mm_slli_si128(_mm_srli_epi64(v, 64 - k), 8)
    );
}

AESNI_FUNC
static void gcm_mul_clmul(std::array<uint8_t,16>& X, const uint8_t H[16]) {
    __m128i x = reflect128(_mm_loadu_si128(reinterpret_cast<const __m128i*>(X.data())));
    __m128i h = reflect128(_mm_loadu_si128(reinterpret_cast<const __m128i*>(H)));

    // Karatsuba CLMUL: 256-bit product [hi128:lo128]
    __m128i ll  = _mm_clmulepi64_si128(x, h, 0x00);
    __m128i hh  = _mm_clmulepi64_si128(x, h, 0x11);
    __m128i mid = _mm_xor_si128(_mm_clmulepi64_si128(x, h, 0x01),
                                _mm_clmulepi64_si128(x, h, 0x10));
    __m128i lo128 = _mm_xor_si128(ll, _mm_slli_si128(mid, 8));
    __m128i hi128 = _mm_xor_si128(hh, _mm_srli_si128(mid, 8));

    // Reduce [hi128:lo128] mod P(x) = x^128 + x^7 + x^2 + x + 1
    // fold(hi128) = hi128 ^ (hi128<<1) ^ (hi128<<2) ^ (hi128<<7)
    __m128i s1 = shl128(hi128, 1);
    __m128i s2 = shl128(hi128, 2);
    __m128i s7 = shl128(hi128, 7);
    __m128i fold = _mm_xor_si128(hi128, _mm_xor_si128(s1, _mm_xor_si128(s2, s7)));

    // Overflow bits (128..134) from shifts of hi128
    alignas(16) uint64_t hipair[2];
    _mm_store_si128(reinterpret_cast<__m128i*>(hipair), hi128);
    uint64_t hi_hi = hipair[1];
    uint64_t ov = (hi_hi >> 63) ^ (hi_hi >> 62) ^ (hi_hi >> 57);
    // Second-level fold of overflow (max 14 bits)
    uint64_t ov2 = ov ^ (ov << 1) ^ (ov << 2) ^ (ov << 7);

    __m128i result = _mm_xor_si128(lo128, fold);
    result = _mm_xor_si128(result, _mm_set_epi64x(0, static_cast<long long>(ov2)));

    _mm_storeu_si128(reinterpret_cast<__m128i*>(X.data()), reflect128(result));
}

// ---- Inline __m128i GHASH multiply for interleaved use ----
// Operates directly on __m128i in GCM byte order (no array conversion).
AESNI_FUNC
static inline __m128i gcm_mul_m128i(__m128i X, __m128i H_reflected) {
    __m128i x = reflect128(X);

    __m128i ll  = _mm_clmulepi64_si128(x, H_reflected, 0x00);
    __m128i hh  = _mm_clmulepi64_si128(x, H_reflected, 0x11);
    __m128i mid = _mm_xor_si128(_mm_clmulepi64_si128(x, H_reflected, 0x01),
                                _mm_clmulepi64_si128(x, H_reflected, 0x10));
    __m128i lo128 = _mm_xor_si128(ll, _mm_slli_si128(mid, 8));
    __m128i hi128 = _mm_xor_si128(hh, _mm_srli_si128(mid, 8));

    __m128i s1 = shl128(hi128, 1);
    __m128i s2 = shl128(hi128, 2);
    __m128i s7 = shl128(hi128, 7);
    __m128i fold = _mm_xor_si128(hi128, _mm_xor_si128(s1, _mm_xor_si128(s2, s7)));

    alignas(16) uint64_t hipair[2];
    _mm_store_si128(reinterpret_cast<__m128i*>(hipair), hi128);
    uint64_t hi_hi = hipair[1];
    uint64_t ov = (hi_hi >> 63) ^ (hi_hi >> 62) ^ (hi_hi >> 57);
    uint64_t ov2 = ov ^ (ov << 1) ^ (ov << 2) ^ (ov << 7);

    __m128i result = _mm_xor_si128(lo128, fold);
    result = _mm_xor_si128(result, _mm_set_epi64x(0, static_cast<long long>(ov2)));

    return reflect128(result);
}

// ---- GHASH update: Y = gcm_mul(Y ^ block, H) ----
AESNI_FUNC
static inline __m128i ghash_update(__m128i Y, __m128i block, __m128i H_reflected) {
    return gcm_mul_m128i(_mm_xor_si128(Y, block), H_reflected);
}

// ---- Interleaved AES-CTR + GHASH for AES-128 GCM ----
// Performs CTR encryption and GHASH authentication in a single pass.
// While AES rounds execute on the next batch of counters, we GHASH
// the ciphertext blocks from the current batch.
AESNI_FUNC
static void aesni_gcm_ctr_ghash_128(
    const uint8_t* in, uint8_t* out, size_t len,
    const uint8_t J0[16],
    const uint8_t* round_keys,
    const uint8_t* aad, size_t aad_len,
    const uint8_t H[16],
    uint8_t tag16[16]
) {
    __m128i H_ref = reflect128(_mm_loadu_si128(reinterpret_cast<const __m128i*>(H)));
    __m128i Y = _mm_setzero_si128(); // GHASH accumulator

    // Process AAD first
    {
        size_t aoff = 0;
        while (aoff + 16 <= aad_len) {
            __m128i a = _mm_loadu_si128(reinterpret_cast<const __m128i*>(aad + aoff));
            Y = ghash_update(Y, a, H_ref);
            aoff += 16;
        }
        if (aoff < aad_len) {
            alignas(16) uint8_t last[16] = {};
            std::memcpy(last, aad + aoff, aad_len - aoff);
            Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(last)), H_ref);
        }
    }

    uint32_t ctr = (uint32_t(J0[12]) << 24) | (uint32_t(J0[13]) << 16) |
                   (uint32_t(J0[14]) << 8)  |  uint32_t(J0[15]);
    size_t off = 0;

    // Main loop: 4 blocks at a time with interleaved GHASH
    while (off + 64 <= len) {
        // Generate 4 counter blocks
        __m128i c0 = make_ctr_block(J0, ++ctr);
        __m128i c1 = make_ctr_block(J0, ++ctr);
        __m128i c2 = make_ctr_block(J0, ++ctr);
        __m128i c3 = make_ctr_block(J0, ++ctr);

        // Round 0: AddRoundKey
        __m128i rk0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys));
        c0 = _mm_xor_si128(c0, rk0);
        c1 = _mm_xor_si128(c1, rk0);
        c2 = _mm_xor_si128(c2, rk0);
        c3 = _mm_xor_si128(c3, rk0);

        // Rounds 1-9: interleave AES rounds with GHASH of previous ciphertext
        // We do GHASH on previously produced ct blocks during AES round latency
        for (int i = 1; i < 10; i++) {
            __m128i rk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16));
            c0 = _mm_aesenc_si128(c0, rk);
            c1 = _mm_aesenc_si128(c1, rk);
            c2 = _mm_aesenc_si128(c2, rk);
            c3 = _mm_aesenc_si128(c3, rk);
        }

        // Final round
        __m128i rkLast = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160));
        c0 = _mm_aesenclast_si128(c0, rkLast);
        c1 = _mm_aesenclast_si128(c1, rkLast);
        c2 = _mm_aesenclast_si128(c2, rkLast);
        c3 = _mm_aesenclast_si128(c3, rkLast);

        // XOR with plaintext to produce ciphertext
        __m128i p0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        __m128i p1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 16));
        __m128i p2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 32));
        __m128i p3 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 48));

        __m128i ct0 = _mm_xor_si128(c0, p0);
        __m128i ct1 = _mm_xor_si128(c1, p1);
        __m128i ct2 = _mm_xor_si128(c2, p2);
        __m128i ct3 = _mm_xor_si128(c3, p3);

        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off),      ct0);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 16), ct1);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 32), ct2);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 48), ct3);

        // GHASH the 4 ciphertext blocks we just produced
        Y = ghash_update(Y, ct0, H_ref);
        Y = ghash_update(Y, ct1, H_ref);
        Y = ghash_update(Y, ct2, H_ref);
        Y = ghash_update(Y, ct3, H_ref);

        off += 64;
    }

    // Handle remaining full blocks
    while (off + 16 <= len) {
        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 10; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));

        __m128i p = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        __m128i ct_blk = _mm_xor_si128(c, p);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off), ct_blk);
        Y = ghash_update(Y, ct_blk, H_ref);
        off += 16;
    }

    // Final partial block
    if (off < len) {
        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 10; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));

        alignas(16) uint8_t ks[16];
        _mm_storeu_si128(reinterpret_cast<__m128i*>(ks), c);
        alignas(16) uint8_t partial[16] = {};
        size_t rem = len - off;
        for (size_t i = 0; i < rem; i++) {
            out[off + i] = in[off + i] ^ ks[i];
            partial[i] = out[off + i];
        }
        Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(partial)), H_ref);
    }

    // Finalize GHASH: append length block
    alignas(16) uint8_t lenblk[16] = {};
    uint64_t aad_bits = uint64_t(aad_len) * 8;
    uint64_t ct_bits  = uint64_t(len) * 8;
    lenblk[0] = uint8_t(aad_bits >> 56); lenblk[1] = uint8_t(aad_bits >> 48);
    lenblk[2] = uint8_t(aad_bits >> 40); lenblk[3] = uint8_t(aad_bits >> 32);
    lenblk[4] = uint8_t(aad_bits >> 24); lenblk[5] = uint8_t(aad_bits >> 16);
    lenblk[6] = uint8_t(aad_bits >> 8);  lenblk[7] = uint8_t(aad_bits);
    lenblk[8]  = uint8_t(ct_bits >> 56); lenblk[9]  = uint8_t(ct_bits >> 48);
    lenblk[10] = uint8_t(ct_bits >> 40); lenblk[11] = uint8_t(ct_bits >> 32);
    lenblk[12] = uint8_t(ct_bits >> 24); lenblk[13] = uint8_t(ct_bits >> 16);
    lenblk[14] = uint8_t(ct_bits >> 8);  lenblk[15] = uint8_t(ct_bits);

    Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(lenblk)), H_ref);

    // Encrypt J0 for tag
    __m128i j0_block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(J0));
    j0_block = _mm_xor_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    for (int i = 1; i < 10; i++)
        j0_block = _mm_aesenc_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
    j0_block = _mm_aesenclast_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));

    __m128i tag = _mm_xor_si128(j0_block, Y);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(tag16), tag);
}

// ---- Interleaved AES-CTR-decrypt + GHASH for AES-128 GCM (decrypt side) ----
// Unlike the encrypt direction, the ciphertext GHASH needs is already fully
// available up front (it's the input, `in`) rather than being produced
// block-by-block — so GHASH-of-ciphertext and CTR-decrypt of that same
// block have no data dependency on each other and can be computed from the
// same loaded ciphertext register, instead of the old two-pass shape (one
// full GHASH pass over `ct`, then a full separate CTR pass). Writes
// plaintext to `out` and the *computed* tag to `tag16` — the caller (see
// aes128::aes_gcm_decrypt) compares it against the wire tag itself.
AESNI_FUNC
static void aesni_gcm_ctr_ghash_decrypt_128(
    const uint8_t* in, uint8_t* out, size_t len,
    const uint8_t J0[16],
    const uint8_t* round_keys,
    const uint8_t* aad, size_t aad_len,
    const uint8_t H[16],
    uint8_t tag16[16]
) {
    __m128i H_ref = reflect128(_mm_loadu_si128(reinterpret_cast<const __m128i*>(H)));
    __m128i Y = _mm_setzero_si128(); // GHASH accumulator

    // Process AAD first
    {
        size_t aoff = 0;
        while (aoff + 16 <= aad_len) {
            __m128i a = _mm_loadu_si128(reinterpret_cast<const __m128i*>(aad + aoff));
            Y = ghash_update(Y, a, H_ref);
            aoff += 16;
        }
        if (aoff < aad_len) {
            alignas(16) uint8_t last[16] = {};
            std::memcpy(last, aad + aoff, aad_len - aoff);
            Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(last)), H_ref);
        }
    }

    uint32_t ctr = (uint32_t(J0[12]) << 24) | (uint32_t(J0[13]) << 16) |
                   (uint32_t(J0[14]) << 8)  |  uint32_t(J0[15]);
    size_t off = 0;

    // Main loop: 4 blocks at a time. Ciphertext is loaded once per block and
    // used for both GHASH (immediately) and the CTR XOR (once the keystream
    // is ready) — no dependency chain between the two, unlike the encrypt
    // side.
    while (off + 64 <= len) {
        __m128i ct0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        __m128i ct1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 16));
        __m128i ct2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 32));
        __m128i ct3 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 48));

        Y = ghash_update(Y, ct0, H_ref);
        Y = ghash_update(Y, ct1, H_ref);
        Y = ghash_update(Y, ct2, H_ref);
        Y = ghash_update(Y, ct3, H_ref);

        __m128i c0 = make_ctr_block(J0, ++ctr);
        __m128i c1 = make_ctr_block(J0, ++ctr);
        __m128i c2 = make_ctr_block(J0, ++ctr);
        __m128i c3 = make_ctr_block(J0, ++ctr);

        __m128i rk0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys));
        c0 = _mm_xor_si128(c0, rk0);
        c1 = _mm_xor_si128(c1, rk0);
        c2 = _mm_xor_si128(c2, rk0);
        c3 = _mm_xor_si128(c3, rk0);

        for (int i = 1; i < 10; i++) {
            __m128i rk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16));
            c0 = _mm_aesenc_si128(c0, rk);
            c1 = _mm_aesenc_si128(c1, rk);
            c2 = _mm_aesenc_si128(c2, rk);
            c3 = _mm_aesenc_si128(c3, rk);
        }

        __m128i rkLast = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160));
        c0 = _mm_aesenclast_si128(c0, rkLast);
        c1 = _mm_aesenclast_si128(c1, rkLast);
        c2 = _mm_aesenclast_si128(c2, rkLast);
        c3 = _mm_aesenclast_si128(c3, rkLast);

        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off),      _mm_xor_si128(c0, ct0));
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 16), _mm_xor_si128(c1, ct1));
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 32), _mm_xor_si128(c2, ct2));
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 48), _mm_xor_si128(c3, ct3));

        off += 64;
    }

    // Remaining full blocks
    while (off + 16 <= len) {
        __m128i ct_blk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        Y = ghash_update(Y, ct_blk, H_ref);

        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 10; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));

        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off), _mm_xor_si128(c, ct_blk));
        off += 16;
    }

    // Final partial block
    if (off < len) {
        alignas(16) uint8_t partial[16] = {};
        size_t rem = len - off;
        std::memcpy(partial, in + off, rem);
        Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(partial)), H_ref);

        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 10; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));

        alignas(16) uint8_t ks[16];
        _mm_storeu_si128(reinterpret_cast<__m128i*>(ks), c);
        for (size_t i = 0; i < rem; i++)
            out[off + i] = in[off + i] ^ ks[i];
    }

    // Finalize GHASH: append length block
    alignas(16) uint8_t lenblk[16] = {};
    uint64_t aad_bits = uint64_t(aad_len) * 8;
    uint64_t ct_bits  = uint64_t(len) * 8;
    lenblk[0] = uint8_t(aad_bits >> 56); lenblk[1] = uint8_t(aad_bits >> 48);
    lenblk[2] = uint8_t(aad_bits >> 40); lenblk[3] = uint8_t(aad_bits >> 32);
    lenblk[4] = uint8_t(aad_bits >> 24); lenblk[5] = uint8_t(aad_bits >> 16);
    lenblk[6] = uint8_t(aad_bits >> 8);  lenblk[7] = uint8_t(aad_bits);
    lenblk[8]  = uint8_t(ct_bits >> 56); lenblk[9]  = uint8_t(ct_bits >> 48);
    lenblk[10] = uint8_t(ct_bits >> 40); lenblk[11] = uint8_t(ct_bits >> 32);
    lenblk[12] = uint8_t(ct_bits >> 24); lenblk[13] = uint8_t(ct_bits >> 16);
    lenblk[14] = uint8_t(ct_bits >> 8);  lenblk[15] = uint8_t(ct_bits);

    Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(lenblk)), H_ref);

    // Encrypt J0 for tag
    __m128i j0_block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(J0));
    j0_block = _mm_xor_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    for (int i = 1; i < 10; i++)
        j0_block = _mm_aesenc_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
    j0_block = _mm_aesenclast_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 160)));

    __m128i tag = _mm_xor_si128(j0_block, Y);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(tag16), tag);
}

// ---- Interleaved AES-CTR + GHASH for AES-256 GCM ----
AESNI_FUNC
static void aesni_gcm_ctr_ghash_256(
    const uint8_t* in, uint8_t* out, size_t len,
    const uint8_t J0[16],
    const uint8_t* round_keys,
    const uint8_t* aad, size_t aad_len,
    const uint8_t H[16],
    uint8_t tag16[16]
) {
    __m128i H_ref = reflect128(_mm_loadu_si128(reinterpret_cast<const __m128i*>(H)));
    __m128i Y = _mm_setzero_si128();

    // Process AAD
    {
        size_t aoff = 0;
        while (aoff + 16 <= aad_len) {
            __m128i a = _mm_loadu_si128(reinterpret_cast<const __m128i*>(aad + aoff));
            Y = ghash_update(Y, a, H_ref);
            aoff += 16;
        }
        if (aoff < aad_len) {
            alignas(16) uint8_t last[16] = {};
            std::memcpy(last, aad + aoff, aad_len - aoff);
            Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(last)), H_ref);
        }
    }

    uint32_t ctr = (uint32_t(J0[12]) << 24) | (uint32_t(J0[13]) << 16) |
                   (uint32_t(J0[14]) << 8)  |  uint32_t(J0[15]);
    size_t off = 0;

    // Main loop: 4 blocks at a time
    while (off + 64 <= len) {
        __m128i c0 = make_ctr_block(J0, ++ctr);
        __m128i c1 = make_ctr_block(J0, ++ctr);
        __m128i c2 = make_ctr_block(J0, ++ctr);
        __m128i c3 = make_ctr_block(J0, ++ctr);

        __m128i rk0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys));
        c0 = _mm_xor_si128(c0, rk0);
        c1 = _mm_xor_si128(c1, rk0);
        c2 = _mm_xor_si128(c2, rk0);
        c3 = _mm_xor_si128(c3, rk0);

        for (int i = 1; i < 14; i++) {
            __m128i rk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16));
            c0 = _mm_aesenc_si128(c0, rk);
            c1 = _mm_aesenc_si128(c1, rk);
            c2 = _mm_aesenc_si128(c2, rk);
            c3 = _mm_aesenc_si128(c3, rk);
        }

        __m128i rkLast = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224));
        c0 = _mm_aesenclast_si128(c0, rkLast);
        c1 = _mm_aesenclast_si128(c1, rkLast);
        c2 = _mm_aesenclast_si128(c2, rkLast);
        c3 = _mm_aesenclast_si128(c3, rkLast);

        __m128i p0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        __m128i p1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 16));
        __m128i p2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 32));
        __m128i p3 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 48));

        __m128i ct0 = _mm_xor_si128(c0, p0);
        __m128i ct1 = _mm_xor_si128(c1, p1);
        __m128i ct2 = _mm_xor_si128(c2, p2);
        __m128i ct3 = _mm_xor_si128(c3, p3);

        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off),      ct0);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 16), ct1);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 32), ct2);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 48), ct3);

        Y = ghash_update(Y, ct0, H_ref);
        Y = ghash_update(Y, ct1, H_ref);
        Y = ghash_update(Y, ct2, H_ref);
        Y = ghash_update(Y, ct3, H_ref);

        off += 64;
    }

    while (off + 16 <= len) {
        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 14; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));

        __m128i p = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        __m128i ct_blk = _mm_xor_si128(c, p);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off), ct_blk);
        Y = ghash_update(Y, ct_blk, H_ref);
        off += 16;
    }

    if (off < len) {
        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 14; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));

        alignas(16) uint8_t ks[16];
        _mm_storeu_si128(reinterpret_cast<__m128i*>(ks), c);
        alignas(16) uint8_t partial[16] = {};
        size_t rem = len - off;
        for (size_t i = 0; i < rem; i++) {
            out[off + i] = in[off + i] ^ ks[i];
            partial[i] = out[off + i];
        }
        Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(partial)), H_ref);
    }

    // Finalize GHASH
    alignas(16) uint8_t lenblk[16] = {};
    uint64_t aad_bits = uint64_t(aad_len) * 8;
    uint64_t ct_bits  = uint64_t(len) * 8;
    lenblk[0] = uint8_t(aad_bits >> 56); lenblk[1] = uint8_t(aad_bits >> 48);
    lenblk[2] = uint8_t(aad_bits >> 40); lenblk[3] = uint8_t(aad_bits >> 32);
    lenblk[4] = uint8_t(aad_bits >> 24); lenblk[5] = uint8_t(aad_bits >> 16);
    lenblk[6] = uint8_t(aad_bits >> 8);  lenblk[7] = uint8_t(aad_bits);
    lenblk[8]  = uint8_t(ct_bits >> 56); lenblk[9]  = uint8_t(ct_bits >> 48);
    lenblk[10] = uint8_t(ct_bits >> 40); lenblk[11] = uint8_t(ct_bits >> 32);
    lenblk[12] = uint8_t(ct_bits >> 24); lenblk[13] = uint8_t(ct_bits >> 16);
    lenblk[14] = uint8_t(ct_bits >> 8);  lenblk[15] = uint8_t(ct_bits);

    Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(lenblk)), H_ref);

    // Encrypt J0 for tag
    __m128i j0_block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(J0));
    j0_block = _mm_xor_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    for (int i = 1; i < 14; i++)
        j0_block = _mm_aesenc_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
    j0_block = _mm_aesenclast_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));

    __m128i tag = _mm_xor_si128(j0_block, Y);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(tag16), tag);
}

// ---- Interleaved AES-CTR-decrypt + GHASH for AES-256 GCM (decrypt side) ----
// See aesni_gcm_ctr_ghash_decrypt_128's comment — same fusion, 14 rounds.
AESNI_FUNC
static void aesni_gcm_ctr_ghash_decrypt_256(
    const uint8_t* in, uint8_t* out, size_t len,
    const uint8_t J0[16],
    const uint8_t* round_keys,
    const uint8_t* aad, size_t aad_len,
    const uint8_t H[16],
    uint8_t tag16[16]
) {
    __m128i H_ref = reflect128(_mm_loadu_si128(reinterpret_cast<const __m128i*>(H)));
    __m128i Y = _mm_setzero_si128();

    {
        size_t aoff = 0;
        while (aoff + 16 <= aad_len) {
            __m128i a = _mm_loadu_si128(reinterpret_cast<const __m128i*>(aad + aoff));
            Y = ghash_update(Y, a, H_ref);
            aoff += 16;
        }
        if (aoff < aad_len) {
            alignas(16) uint8_t last[16] = {};
            std::memcpy(last, aad + aoff, aad_len - aoff);
            Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(last)), H_ref);
        }
    }

    uint32_t ctr = (uint32_t(J0[12]) << 24) | (uint32_t(J0[13]) << 16) |
                   (uint32_t(J0[14]) << 8)  |  uint32_t(J0[15]);
    size_t off = 0;

    while (off + 64 <= len) {
        __m128i ct0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        __m128i ct1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 16));
        __m128i ct2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 32));
        __m128i ct3 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off + 48));

        Y = ghash_update(Y, ct0, H_ref);
        Y = ghash_update(Y, ct1, H_ref);
        Y = ghash_update(Y, ct2, H_ref);
        Y = ghash_update(Y, ct3, H_ref);

        __m128i c0 = make_ctr_block(J0, ++ctr);
        __m128i c1 = make_ctr_block(J0, ++ctr);
        __m128i c2 = make_ctr_block(J0, ++ctr);
        __m128i c3 = make_ctr_block(J0, ++ctr);

        __m128i rk0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys));
        c0 = _mm_xor_si128(c0, rk0);
        c1 = _mm_xor_si128(c1, rk0);
        c2 = _mm_xor_si128(c2, rk0);
        c3 = _mm_xor_si128(c3, rk0);

        for (int i = 1; i < 14; i++) {
            __m128i rk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16));
            c0 = _mm_aesenc_si128(c0, rk);
            c1 = _mm_aesenc_si128(c1, rk);
            c2 = _mm_aesenc_si128(c2, rk);
            c3 = _mm_aesenc_si128(c3, rk);
        }

        __m128i rkLast = _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224));
        c0 = _mm_aesenclast_si128(c0, rkLast);
        c1 = _mm_aesenclast_si128(c1, rkLast);
        c2 = _mm_aesenclast_si128(c2, rkLast);
        c3 = _mm_aesenclast_si128(c3, rkLast);

        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off),      _mm_xor_si128(c0, ct0));
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 16), _mm_xor_si128(c1, ct1));
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 32), _mm_xor_si128(c2, ct2));
        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off + 48), _mm_xor_si128(c3, ct3));

        off += 64;
    }

    while (off + 16 <= len) {
        __m128i ct_blk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(in + off));
        Y = ghash_update(Y, ct_blk, H_ref);

        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 14; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));

        _mm_storeu_si128(reinterpret_cast<__m128i*>(out + off), _mm_xor_si128(c, ct_blk));
        off += 16;
    }

    if (off < len) {
        alignas(16) uint8_t partial[16] = {};
        size_t rem = len - off;
        std::memcpy(partial, in + off, rem);
        Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(partial)), H_ref);

        __m128i c = make_ctr_block(J0, ++ctr);
        c = _mm_xor_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
        for (int i = 1; i < 14; i++)
            c = _mm_aesenc_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
        c = _mm_aesenclast_si128(c, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));

        alignas(16) uint8_t ks[16];
        _mm_storeu_si128(reinterpret_cast<__m128i*>(ks), c);
        for (size_t i = 0; i < rem; i++)
            out[off + i] = in[off + i] ^ ks[i];
    }

    // Finalize GHASH
    alignas(16) uint8_t lenblk[16] = {};
    uint64_t aad_bits = uint64_t(aad_len) * 8;
    uint64_t ct_bits  = uint64_t(len) * 8;
    lenblk[0] = uint8_t(aad_bits >> 56); lenblk[1] = uint8_t(aad_bits >> 48);
    lenblk[2] = uint8_t(aad_bits >> 40); lenblk[3] = uint8_t(aad_bits >> 32);
    lenblk[4] = uint8_t(aad_bits >> 24); lenblk[5] = uint8_t(aad_bits >> 16);
    lenblk[6] = uint8_t(aad_bits >> 8);  lenblk[7] = uint8_t(aad_bits);
    lenblk[8]  = uint8_t(ct_bits >> 56); lenblk[9]  = uint8_t(ct_bits >> 48);
    lenblk[10] = uint8_t(ct_bits >> 40); lenblk[11] = uint8_t(ct_bits >> 32);
    lenblk[12] = uint8_t(ct_bits >> 24); lenblk[13] = uint8_t(ct_bits >> 16);
    lenblk[14] = uint8_t(ct_bits >> 8);  lenblk[15] = uint8_t(ct_bits);

    Y = ghash_update(Y, _mm_loadu_si128(reinterpret_cast<const __m128i*>(lenblk)), H_ref);

    __m128i j0_block = _mm_loadu_si128(reinterpret_cast<const __m128i*>(J0));
    j0_block = _mm_xor_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys)));
    for (int i = 1; i < 14; i++)
        j0_block = _mm_aesenc_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + i * 16)));
    j0_block = _mm_aesenclast_si128(j0_block, _mm_loadu_si128(reinterpret_cast<const __m128i*>(round_keys + 224)));

    __m128i tag = _mm_xor_si128(j0_block, Y);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(tag16), tag);
}

#endif // NETPLUS_X86

namespace netplus {

// ============================================================
// AES-128 Implementation
// ============================================================

// Constant-time GF(2^8) multiply (AES field poly 0x11B)
uint8_t aes128::gf_mul(uint8_t a, uint8_t b) {
    uint8_t p = 0;
    for (int i = 0; i < 8; i++) {
        uint8_t mask = static_cast<uint8_t>(0 - (b & 1));
        p ^= (a & mask);
        uint8_t hi = static_cast<uint8_t>(a & 0x80);
        a <<= 1;
        a ^= static_cast<uint8_t>(0x1B & static_cast<uint8_t>(0 - (hi >> 7)));
        b >>= 1;
    }
    return p;
}

uint8_t aes128::gf_inv(uint8_t a) {
    // constant-time inversion by exponentiation: a^(254)
    // 0 maps to 0
    uint8_t a2  = gf_mul(a,a);
    uint8_t a4  = gf_mul(a2,a2);
    uint8_t a8  = gf_mul(a4,a4);
    uint8_t a16 = gf_mul(a8,a8);
    uint8_t a32 = gf_mul(a16,a16);
    uint8_t a64 = gf_mul(a32,a32);
    uint8_t a128= gf_mul(a64,a64);

    uint8_t t = gf_mul(a128,a64);
    t = gf_mul(t,a32);
    t = gf_mul(t,a16);
    t = gf_mul(t,a8);
    t = gf_mul(t,a4);
    t = gf_mul(t,a2);

    return t;
}

// Constant-time AES S-box computed (no lookup tables)
uint8_t aes128::ct_sbox(uint8_t x) {
    uint8_t inv = gf_inv(x);

    // affine transform
    uint8_t y = inv;
    y ^= rotl8(inv,1);
    y ^= rotl8(inv,2);
    y ^= rotl8(inv,3);
    y ^= rotl8(inv,4);
    y ^= 0x63;
    return y;
}

uint8_t aes128::ct_inv_sbox(uint8_t x) {
    // inverse affine then inversion
    uint8_t y = x;
    y ^= 0x63;
    uint8_t t = y;
    y = static_cast<uint8_t>(
        rotl8(t,1) ^ rotl8(t,3) ^ rotl8(t,6)
    );
    return gf_inv(y);
}

// AES helper multiply for MixColumns (still constant time)
uint8_t aes128::gmul(uint8_t a, uint8_t b) {
    return gf_mul(a,b);
}

// Key schedule helpers
uint32_t aes128::rotWord(uint32_t w) {
    return (w << 8) | (w >> 24);
}

uint32_t aes128::subWord_ct(uint32_t w) {
    uint8_t b0 = ct_sbox(uint8_t(w >> 24));
    uint8_t b1 = ct_sbox(uint8_t(w >> 16));
    uint8_t b2 = ct_sbox(uint8_t(w >>  8));
    uint8_t b3 = ct_sbox(uint8_t(w));
    return (uint32_t(b0) << 24) |
           (uint32_t(b1) << 16) |
           (uint32_t(b2) <<  8) |
           (uint32_t(b3));
}

void aes128::keyExpansion(const key128& key) {
    if (key.size() != 16)
        throw std::invalid_argument("aes128::keyExpansion: key must be 16 bytes");

    static const uint8_t RCON[10] = {
        0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80,0x1B,0x36
    };

    m_roundKeyBytes.assign(176, 0);
    std::memcpy(m_roundKeyBytes.data(), key.data(), 16);

    // Copy initial key
    m_roundKeyBytes.resize(176); // 11 rounds * 16 bytes
    std::memcpy(m_roundKeyBytes.data(), key.data(), 16);

    uint32_t temp;
    int bytesGenerated = 16;
    int rconIter = 0;

    while (bytesGenerated < 176) {
        temp =
            (uint32_t(m_roundKeyBytes[bytesGenerated - 4]) << 24) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 3]) << 16) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 2]) <<  8) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 1]));

        if (bytesGenerated % 16 == 0) {
            temp = subWord_ct(rotWord(temp));
            temp ^= (uint32_t(RCON[rconIter]) << 24);
            rconIter++;
        }

        uint32_t prev =
            (uint32_t(m_roundKeyBytes[bytesGenerated - 16]) << 24) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 15]) << 16) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 14]) <<  8) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 13]));

        temp ^= prev;

        m_roundKeyBytes[bytesGenerated + 0] = uint8_t(temp >> 24);
        m_roundKeyBytes[bytesGenerated + 1] = uint8_t(temp >> 16);
        m_roundKeyBytes[bytesGenerated + 2] = uint8_t(temp >>  8);
        m_roundKeyBytes[bytesGenerated + 3] = uint8_t(temp);

        bytesGenerated += 4;
    }
}

// AES round ops (constant-time S-box)
void aes128::addRoundKey(uint8_t* s, const uint8_t* rk) {
    for (int i = 0; i < 16; i++) s[i] ^= rk[i];
}

void aes128::subBytes(uint8_t* s) {
    for (int i = 0; i < 16; i++) s[i] = ct_sbox(s[i]);
}

void aes128::invSubBytes(uint8_t* s) {
    for (int i = 0; i < 16; i++) s[i] = ct_inv_sbox(s[i]);
}

// state is column-major
void aes128::shiftRows(uint8_t* s) {
    uint8_t t[16];
    std::memcpy(t, s, 16);

    s[0]  = t[0];  s[4]  = t[4];  s[8]  = t[8];  s[12] = t[12];
    s[1]  = t[5];  s[5]  = t[9];  s[9]  = t[13]; s[13] = t[1];
    s[2]  = t[10]; s[6]  = t[14]; s[10] = t[2];  s[14] = t[6];
    s[3]  = t[15]; s[7]  = t[3];  s[11] = t[7];  s[15] = t[11];
}

void aes128::invShiftRows(uint8_t* s) {
    uint8_t t[16];
    std::memcpy(t, s, 16);

    s[0]  = t[0];  s[4]  = t[4];  s[8]  = t[8];  s[12] = t[12];
    s[1]  = t[13]; s[5]  = t[1];  s[9]  = t[5];  s[13] = t[9];
    s[2]  = t[10]; s[6]  = t[14]; s[10] = t[2];  s[14] = t[6];
    s[3]  = t[7];  s[7]  = t[11]; s[11] = t[15]; s[15] = t[3];
}

void aes128::mixColumns(uint8_t* s) {
    for (int c = 0; c < 4; c++) {
        uint8_t* col = &s[4*c];
        uint8_t a0=col[0], a1=col[1], a2=col[2], a3=col[3];

        uint8_t r0 = xtime(a0) ^ (xtime(a1)^a1) ^ a2 ^ a3;
        uint8_t r1 = a0 ^ xtime(a1) ^ (xtime(a2)^a2) ^ a3;
        uint8_t r2 = a0 ^ a1 ^ xtime(a2) ^ (xtime(a3)^a3);
        uint8_t r3 = (xtime(a0)^a0) ^ a1 ^ a2 ^ xtime(a3);

        col[0]=r0; col[1]=r1; col[2]=r2; col[3]=r3;
    }
}

void aes128::invMixColumns(uint8_t* s) {
    for (int c = 0; c < 4; c++) {
        uint8_t* col = &s[4*c];
        uint8_t a0=col[0], a1=col[1], a2=col[2], a3=col[3];

        col[0] = gmul(0x0E,a0)^gmul(0x0B,a1)^gmul(0x0D,a2)^gmul(0x09,a3);
        col[1] = gmul(0x09,a0)^gmul(0x0E,a1)^gmul(0x0B,a2)^gmul(0x0D,a3);
        col[2] = gmul(0x0D,a0)^gmul(0x09,a1)^gmul(0x0E,a2)^gmul(0x0B,a3);
        col[3] = gmul(0x0B,a0)^gmul(0x0D,a1)^gmul(0x09,a2)^gmul(0x0E,a3);
    }
}

// Constructor
aes128::aes128(const key128& key) {
    keyExpansion(key);
    uint8_t zero[16] = {0};
    encrypt_block(zero, m_H);
}

// AES block encrypt/decrypt
void aes128::encrypt_block(const uint8_t in[16], uint8_t out[16]) const {
#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        aesni_encrypt_128(in, out, m_roundKeyBytes.data());
        return;
    }
#endif
    uint8_t s[16];
    std::memcpy(s, in, 16);

    addRoundKey(s, rk(0));
    for (int r = 1; r < 10; r++) {
        subBytes(s);
        shiftRows(s);
        mixColumns(s);
        addRoundKey(s, rk(r));
    }
    subBytes(s);
    shiftRows(s);
    addRoundKey(s, rk(10));

    std::memcpy(out, s, 16);
}

void aes128::decrypt_block(const uint8_t in[16], uint8_t out[16]) const {
#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        aesni_decrypt_128(in, out, m_roundKeyBytes.data());
        return;
    }
#endif
    uint8_t s[16];
    std::memcpy(s, in, 16);

    addRoundKey(s, rk(10));
    for (int r = 9; r >= 1; r--) {
        invShiftRows(s);
        invSubBytes(s);
        addRoundKey(s, rk(r));
        invMixColumns(s);
    }
    invShiftRows(s);
    invSubBytes(s);
    addRoundKey(s, rk(0));

    std::memcpy(out, s, 16);
}

// Vector API
std::vector<uint8_t> aes128::encrypt(const std::vector<uint8_t>& pt) {
    if (pt.size() != 16) throw std::invalid_argument("encrypt: block must be 16 bytes");
    uint8_t out[16];
    encrypt_block(pt.data(), out);
    return std::vector<uint8_t>(out, out+16);
}

std::vector<uint8_t> aes128::decrypt(const std::vector<uint8_t>& ct) {
    if (ct.size() != 16) throw std::invalid_argument("decrypt: block must be 16 bytes");
    uint8_t out[16];
    decrypt_block(ct.data(), out);
    return std::vector<uint8_t>(out, out+16);
}

// CBC mode
std::vector<uint8_t> aes128::encryptCBC(const std::vector<uint8_t>& pt,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("encryptCBC: iv must be 16 bytes");
    if (pt.empty() || (pt.size() % 16) != 0) throw std::invalid_argument("encryptCBC: pt multiple of 16");

    std::vector<uint8_t> out(pt.size());
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < pt.size(); off += 16) {
        uint8_t blk[16];
        for (int i = 0; i < 16; i++) blk[i] = pt[off+i] ^ prev[i];
        encrypt_block(blk, &out[off]);
        std::memcpy(prev, &out[off], 16);
    }
    return out;
}

std::vector<uint8_t> aes128::decryptCBC(const std::vector<uint8_t>& ct,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("decryptCBC: iv must be 16 bytes");
    if (ct.empty() || (ct.size() % 16) != 0) throw std::invalid_argument("decryptCBC: ct multiple of 16");

    std::vector<uint8_t> out(ct.size());
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < ct.size(); off += 16) {
        uint8_t dec[16];
        decrypt_block(&ct[off], dec);
        for (int i = 0; i < 16; i++) out[off+i] = dec[i] ^ prev[i];
        std::memcpy(prev, &ct[off], 16);
    }
    return out;
}

// CTS mode (Cipher Text Stealing) - RFC 3962 for Kerberos
// CBC-CS3 variant: last two ciphertext blocks are swapped.
// Input must be >= 16 bytes. Output size == input size (no padding).
std::vector<uint8_t> aes128::encryptCTS(const std::vector<uint8_t>& pt,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("encryptCTS: iv must be 16 bytes");
    if (pt.size() < 16)  throw std::invalid_argument("encryptCTS: plaintext must be >= 16 bytes");

    size_t n = pt.size();

    // Special case: exactly one block - same as ECB
    if (n == 16) {
        uint8_t blk[16], out[16];
        for (int i = 0; i < 16; i++) blk[i] = pt[i] ^ iv[i];
        encrypt_block(blk, out);
        return std::vector<uint8_t>(out, out + 16);
    }

    // Pad plaintext to multiple of 16 with zeros for CBC processing
    size_t nBlocks = (n + 15) / 16;
    std::vector<uint8_t> padded = pt;
    padded.resize(nBlocks * 16, 0);

    // Standard CBC encrypt
    std::vector<uint8_t> cbcOut(nBlocks * 16);
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < nBlocks * 16; off += 16) {
        uint8_t blk[16];
        for (int i = 0; i < 16; i++) blk[i] = padded[off + i] ^ prev[i];
        encrypt_block(blk, &cbcOut[off]);
        std::memcpy(prev, &cbcOut[off], 16);
    }

    // CTS swap: output = CBC[0..n-3] || CBC[n-1] (full) || CBC[n-2] (truncated)
    std::vector<uint8_t> result(n);
    size_t penultimateOff = (nBlocks - 2) * 16;
    size_t lastOff = (nBlocks - 1) * 16;
    size_t lastPlainLen = n - penultimateOff - 16; // bytes in last partial block

    // Copy all complete blocks except last two
    if (nBlocks > 2)
        std::memcpy(result.data(), cbcOut.data(), penultimateOff);

    // Swap: last full ciphertext block, then truncated penultimate
    std::memcpy(result.data() + penultimateOff, &cbcOut[lastOff], 16);
    std::memcpy(result.data() + penultimateOff + 16, &cbcOut[penultimateOff], lastPlainLen);

    return result;
}

std::vector<uint8_t> aes128::decryptCTS(const std::vector<uint8_t>& ct,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("decryptCTS: iv must be 16 bytes");
    if (ct.size() < 16)  throw std::invalid_argument("decryptCTS: ciphertext must be >= 16 bytes");

    size_t n = ct.size();

    // Special case: exactly one block
    if (n == 16) {
        uint8_t dec[16];
        decrypt_block(ct.data(), dec);
        std::vector<uint8_t> out(16);
        for (int i = 0; i < 16; i++) out[i] = dec[i] ^ iv[i];
        return out;
    }

    size_t nBlocks = (n + 15) / 16;
    size_t penultimateOff = (nBlocks - 2) * 16;
    size_t lastPlainLen = n - penultimateOff - 16;

    // Un-swap CTS: ct has [blocks 0..n-3] || Cn-1 (full 16) || Cn-2 (truncated lastPlainLen)
    // Actual CBC order: Cn-2 is the full block, Cn-1 is the padded truncated one
    uint8_t Cn_1[16]; // the swapped full block (at penultimateOff in ct)
    std::memcpy(Cn_1, ct.data() + penultimateOff, 16);

    uint8_t Cn_2_trunc[16] = {}; // the truncated block (at penultimateOff+16 in ct)
    std::memcpy(Cn_2_trunc, ct.data() + penultimateOff + 16, lastPlainLen);

    // Decrypt Cn_1 (the swapped block) to get intermediate
    uint8_t decCn1[16];
    decrypt_block(Cn_1, decCn1);

    // Reconstruct the full Cn_2 block: truncated bytes from ct, rest from decrypted Cn_1
    uint8_t Cn_2_full[16];
    std::memcpy(Cn_2_full, Cn_2_trunc, lastPlainLen);
    std::memcpy(Cn_2_full + lastPlainLen, decCn1 + lastPlainLen, 16 - lastPlainLen);

    // Now decrypt all blocks in proper CBC order
    std::vector<uint8_t> result(n);

    // Decrypt blocks 0..n-3 normally with CBC
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < penultimateOff; off += 16) {
        uint8_t dec[16];
        decrypt_block(ct.data() + off, dec);
        for (int i = 0; i < 16; i++) result[off + i] = dec[i] ^ prev[i];
        std::memcpy(prev, ct.data() + off, 16);
    }

    // Decrypt reconstructed Cn_2_full (penultimate in CBC order)
    {
        uint8_t dec[16];
        decrypt_block(Cn_2_full, dec);
        for (int i = 0; i < 16; i++) result[penultimateOff + i] = dec[i] ^ prev[i];
        std::memcpy(prev, Cn_2_full, 16);
    }

    // Last partial block: XOR decrypted Cn_1 with prev (Cn_2_full)
    for (size_t i = 0; i < lastPlainLen; i++)
        result[penultimateOff + 16 + i] = decCn1[i] ^ prev[i];

    return result;
}

// GCM operations
void aes128::gcm_mul(std::array<uint8_t,16>& X) const {
#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        gcm_mul_clmul(X, m_H);
        return;
    }
#endif
    uint64_t Xhi = load_be64(X.data());
    uint64_t Xlo = load_be64(X.data() + 8);
    
    uint64_t Hhi = load_be64(m_H);
    uint64_t Hlo = load_be64(m_H + 8);
    
    uint64_t Zhi = 0, Zlo = 0;
    
    // Process each bit of X from MSB to LSB
    for (int i = 0; i < 128; i++) {
        // Get bit i of X (MSB first)
        uint64_t xi;
        if (i < 64)
            xi = (Xhi >> (63 - i)) & 1;
        else
            xi = (Xlo >> (63 - (i - 64))) & 1;
        
        // If bit is set, XOR H into Z
        if (xi) {
            Zhi ^= Hhi;
            Zlo ^= Hlo;
        }
        
        // Multiply H by x (shift right, reduce if needed)
        uint64_t lsb = Hlo & 1;
        Hlo = (Hlo >> 1) | (Hhi << 63);
        Hhi = Hhi >> 1;
        
        // If LSB was 1, reduce by XORing with R (0xE1 << 56)
        if (lsb) {
            Hhi ^= 0xE100000000000000ULL;
        }
    }
    
    store_be64(X.data(), Zhi);
    store_be64(X.data() + 8, Zlo);
}

void aes128::ghash(std::array<uint8_t,16>& out,
                   const uint8_t* aad, size_t aad_len,
                   const uint8_t* ct,  size_t ct_len) const {
    std::array<uint8_t,16> Y{};
    Y.fill(0);

    auto update = [&](const uint8_t* data, size_t len) {
        while (len >= 16) {
            for (int i = 0; i < 16; i++) Y[i] ^= data[i];
            gcm_mul(Y);
            data += 16;
            len  -= 16;
        }
        if (len > 0) {
            std::array<uint8_t,16> last{};
            last.fill(0);
            std::memcpy(last.data(), data, len);
            for (int i = 0; i < 16; i++) Y[i] ^= last[i];
            gcm_mul(Y);
        }
    };

    update(aad, aad_len);
    update(ct,  ct_len);

    uint8_t lenblk[16] = {0};
    store_be64(lenblk,     uint64_t(aad_len) * 8);
    store_be64(lenblk + 8, uint64_t(ct_len)  * 8);

    for (int i = 0; i < 16; i++) Y[i] ^= lenblk[i];
    gcm_mul(Y);

    out = Y;
}

bool aes128::aes_gcm_encrypt(
    const uint8_t iv12[12],
    const uint8_t* aad, size_t aad_len,
    const uint8_t* pt,  size_t pt_len,
    uint8_t* ct,
    uint8_t tag16[16]
) const {
    uint8_t J0[16] = {0};
    std::memcpy(J0, iv12, 12);
    J0[15] = 1;

#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        aesni_gcm_ctr_ghash_128(pt, ct, pt_len, J0, m_roundKeyBytes.data(),
                                 aad, aad_len, m_H, tag16);
        return true;
    }
#endif
    {
        uint8_t ctr[16];
        std::memcpy(ctr, J0, 16);

        size_t off = 0;
        while (off < pt_len) {
            inc32(ctr);
            uint8_t S[16];
            encrypt_block(ctr, S);

            size_t take = std::min<size_t>(16, pt_len - off);
            for (size_t i = 0; i < take; i++)
                ct[off+i] = pt[off+i] ^ S[i];
            off += take;
        }
    }

    std::array<uint8_t,16> Sgh{};
    Sgh.fill(0);
    ghash(Sgh, aad, aad_len, ct, pt_len);

    uint8_t EJ0[16];
    encrypt_block(J0, EJ0);

    for (int i = 0; i < 16; i++)
        tag16[i] = EJ0[i] ^ Sgh[i];

    return true;
}

bool aes128::aes_gcm_decrypt(
    const uint8_t iv12[12],
    const uint8_t* aad, size_t aad_len,
    const uint8_t* ct,  size_t ct_len,
    const uint8_t tag16[16],
    uint8_t* pt
) const{
    uint8_t J0[16] = {0};
    std::memcpy(J0, iv12, 12);
    J0[15] = 1;

#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        // Fused single pass: GHASH-of-ciphertext and CTR-decrypt share one
        // loop over the (already fully available) ciphertext instead of a
        // full GHASH pass followed by a full separate CTR pass — see
        // aesni_gcm_ctr_ghash_decrypt_128's comment.
        uint8_t expTag[16];
        aesni_gcm_ctr_ghash_decrypt_128(ct, pt, ct_len, J0, m_roundKeyBytes.data(),
                                         aad, aad_len, m_H, expTag);
        if (!ct_memeq(expTag, tag16, 16)) {
            // AEAD contract is all-or-nothing: the caller must never use
            // `pt` when this returns false. It was written speculatively
            // by the fused pass above, so wipe it as defense in depth.
            std::memset(pt, 0, ct_len);
            return false;
        }
        return true;
    }
#endif

    std::array<uint8_t,16> Sgh{};
    Sgh.fill(0);
    ghash(Sgh, aad, aad_len, ct, ct_len);

    uint8_t EJ0[16];
    encrypt_block(J0, EJ0);

    uint8_t expTag[16];
    for (int i = 0; i < 16; i++)
        expTag[i] = EJ0[i] ^ Sgh[i];

    if (!ct_memeq(expTag, tag16, 16))
        return false;

    uint8_t ctr[16];
    std::memcpy(ctr, J0, 16);

    size_t off = 0;
    while (off < ct_len) {
        inc32(ctr);
        uint8_t S[16];
        encrypt_block(ctr, S);

        size_t take = std::min<size_t>(16, ct_len - off);
        for (size_t i = 0; i < take; i++)
            pt[off+i] = ct[off+i] ^ S[i];
        off += take;
    }

    return true;
}

// ============================================================
// AES-256 Implementation
// ============================================================

uint8_t aes256::gf_mul(uint8_t a, uint8_t b) {
    uint8_t p = 0;
    for (int i = 0; i < 8; i++) {
        uint8_t mask = static_cast<uint8_t>(0 - (b & 1));
        p ^= (a & mask);
        uint8_t hi = static_cast<uint8_t>(a & 0x80);
        a <<= 1;
        a ^= static_cast<uint8_t>(0x1B & static_cast<uint8_t>(0 - (hi >> 7)));
        b >>= 1;
    }
    return p;
}

uint8_t aes256::gf_inv(uint8_t a) {
    uint8_t a2  = gf_mul(a,a);
    uint8_t a4  = gf_mul(a2,a2);
    uint8_t a8  = gf_mul(a4,a4);
    uint8_t a16 = gf_mul(a8,a8);
    uint8_t a32 = gf_mul(a16,a16);
    uint8_t a64 = gf_mul(a32,a32);
    uint8_t a128= gf_mul(a64,a64);

    uint8_t t = gf_mul(a128,a64);
    t = gf_mul(t,a32);
    t = gf_mul(t,a16);
    t = gf_mul(t,a8);
    t = gf_mul(t,a4);
    t = gf_mul(t,a2);
    return t;
}

uint8_t aes256::ct_sbox(uint8_t x) {
    uint8_t inv = gf_inv(x);
    uint8_t y = inv;
    y ^= static_cast<uint8_t>((inv << 1) | (inv >> 7));
    y ^= static_cast<uint8_t>((inv << 2) | (inv >> 6));
    y ^= static_cast<uint8_t>((inv << 3) | (inv >> 5));
    y ^= static_cast<uint8_t>((inv << 4) | (inv >> 4));
    y ^= 0x63;
    return y;
}

uint8_t aes256::ct_inv_sbox(uint8_t x) {
    uint8_t y = x;
    y ^= 0x63;
    uint8_t t = y;
    y = static_cast<uint8_t>(
        ((t << 1) | (t >> 7)) ^ ((t << 3) | (t >> 5)) ^ ((t << 6) | (t >> 2))
    );
    return gf_inv(y);
}

uint8_t aes256::gmul(uint8_t a, uint8_t b) {
    return gf_mul(a,b);
}

uint32_t aes256::rotWord(uint32_t w) {
    return (w << 8) | (w >> 24);
}

uint32_t aes256::subWord_ct(uint32_t w) {
    uint8_t b0 = ct_sbox(uint8_t(w >> 24));
    uint8_t b1 = ct_sbox(uint8_t(w >> 16));
    uint8_t b2 = ct_sbox(uint8_t(w >>  8));
    uint8_t b3 = ct_sbox(uint8_t(w));
    return (uint32_t(b0) << 24) |
           (uint32_t(b1) << 16) |
           (uint32_t(b2) <<  8) |
           (uint32_t(b3));
}

// Deliberately always the constant-time software S-box (subWord_ct), even
// when AES-NI is available: key expansion runs once per encryption
// level/direction/connection (see deriveInitialKeys()/deriveHandshakeKeys()/
// deriveApplicationKeys() in quic.cpp), never in the per-packet hot path —
// AES-NI's _mm_aeskeygenassist_si128 keygen path would only shorten an
// already-sub-microsecond, one-time-per-key operation, and AES-256's
// alternating (SubWord only on every other word) key schedule makes that
// path meaningfully more intricate to get right than AES-128's. Not worth
// the correctness risk for a speedup that never shows up in steady-state
// throughput.
void aes256::keyExpansion(const std::vector<uint8_t>& key) {
    if (key.size() != 32)
        throw std::invalid_argument("aes256::keyExpansion: key must be 32 bytes");

    static const uint8_t RCON[8] = {
        0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80
    };

    m_roundKeyBytes.resize(240); // 15 rounds * 16 bytes
    std::memcpy(m_roundKeyBytes.data(), key.data(), 32);

    uint32_t temp;
    int bytesGenerated = 32;
    int rconIter = 0;

    while (bytesGenerated < 240) {
        temp =
            (uint32_t(m_roundKeyBytes[bytesGenerated - 4]) << 24) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 3]) << 16) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 2]) <<  8) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 1]));

        if (bytesGenerated % 32 == 0) {
            temp = subWord_ct(rotWord(temp));
            temp ^= (uint32_t(RCON[rconIter]) << 24);
            rconIter++;
        } else if ((bytesGenerated % 32) == 16) {
            temp = subWord_ct(temp);
        }

        uint32_t prev =
            (uint32_t(m_roundKeyBytes[bytesGenerated - 32]) << 24) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 31]) << 16) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 30]) <<  8) |
            (uint32_t(m_roundKeyBytes[bytesGenerated - 29]));

        temp ^= prev;

        m_roundKeyBytes[bytesGenerated + 0] = uint8_t(temp >> 24);
        m_roundKeyBytes[bytesGenerated + 1] = uint8_t(temp >> 16);
        m_roundKeyBytes[bytesGenerated + 2] = uint8_t(temp >>  8);
        m_roundKeyBytes[bytesGenerated + 3] = uint8_t(temp);

        bytesGenerated += 4;
    }
}

void aes256::addRoundKey(uint8_t* s, const uint8_t* rk) {
    for (int i = 0; i < 16; i++) s[i] ^= rk[i];
}

void aes256::subBytes(uint8_t* s) {
    for (int i = 0; i < 16; i++) s[i] = ct_sbox(s[i]);
}

void aes256::invSubBytes(uint8_t* s) {
    for (int i = 0; i < 16; i++) s[i] = ct_inv_sbox(s[i]);
}

void aes256::shiftRows(uint8_t* s) {
    uint8_t t[16];
    std::memcpy(t, s, 16);

    s[0]  = t[0];  s[4]  = t[4];  s[8]  = t[8];  s[12] = t[12];
    s[1]  = t[5];  s[5]  = t[9];  s[9]  = t[13]; s[13] = t[1];
    s[2]  = t[10]; s[6]  = t[14]; s[10] = t[2];  s[14] = t[6];
    s[3]  = t[15]; s[7]  = t[3];  s[11] = t[7];  s[15] = t[11];
}

void aes256::invShiftRows(uint8_t* s) {
    uint8_t t[16];
    std::memcpy(t, s, 16);

    s[0]  = t[0];  s[4]  = t[4];  s[8]  = t[8];  s[12] = t[12];
    s[1]  = t[13]; s[5]  = t[1];  s[9]  = t[5];  s[13] = t[9];
    s[2]  = t[10]; s[6]  = t[14]; s[10] = t[2];  s[14] = t[6];
    s[3]  = t[7];  s[7]  = t[11]; s[11] = t[15]; s[15] = t[3];
}

void aes256::mixColumns(uint8_t* s) {
    for (int c = 0; c < 4; c++) {
        uint8_t* col = &s[4*c];
        uint8_t a0=col[0], a1=col[1], a2=col[2], a3=col[3];

        uint8_t r0 = xtime(a0) ^ (xtime(a1)^a1) ^ a2 ^ a3;
        uint8_t r1 = a0 ^ xtime(a1) ^ (xtime(a2)^a2) ^ a3;
        uint8_t r2 = a0 ^ a1 ^ xtime(a2) ^ (xtime(a3)^a3);
        uint8_t r3 = (xtime(a0)^a0) ^ a1 ^ a2 ^ xtime(a3);

        col[0]=r0; col[1]=r1; col[2]=r2; col[3]=r3;
    }
}

void aes256::invMixColumns(uint8_t* s) {
    for (int c = 0; c < 4; c++) {
        uint8_t* col = &s[4*c];
        uint8_t a0=col[0], a1=col[1], a2=col[2], a3=col[3];

        col[0] = gmul(0x0E,a0)^gmul(0x0B,a1)^gmul(0x0D,a2)^gmul(0x09,a3);
        col[1] = gmul(0x09,a0)^gmul(0x0E,a1)^gmul(0x0B,a2)^gmul(0x0D,a3);
        col[2] = gmul(0x0D,a0)^gmul(0x09,a1)^gmul(0x0E,a2)^gmul(0x0B,a3);
        col[3] = gmul(0x0B,a0)^gmul(0x0D,a1)^gmul(0x09,a2)^gmul(0x0E,a3);
    }
}

aes256::aes256(const std::vector<uint8_t>& key) {
    keyExpansion(key);
    uint8_t zero[16] = {0};
    encrypt_block(zero, m_H);
}

void aes256::encrypt_block(const uint8_t in[16], uint8_t out[16]) const {
#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        aesni_encrypt_256(in, out, m_roundKeyBytes.data());
        return;
    }
#endif
    uint8_t s[16];
    std::memcpy(s, in, 16);

    addRoundKey(s, rk(0));
    for (int r = 1; r < 14; r++) {
        subBytes(s);
        shiftRows(s);
        mixColumns(s);
        addRoundKey(s, rk(r));
    }
    subBytes(s);
    shiftRows(s);
    addRoundKey(s, rk(14));

    std::memcpy(out, s, 16);
}

void aes256::decrypt_block(const uint8_t in[16], uint8_t out[16]) const {
#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        aesni_decrypt_256(in, out, m_roundKeyBytes.data());
        return;
    }
#endif
    uint8_t s[16];
    std::memcpy(s, in, 16);

    addRoundKey(s, rk(14));
    for (int r = 13; r >= 1; r--) {
        invShiftRows(s);
        invSubBytes(s);
        addRoundKey(s, rk(r));
        invMixColumns(s);
    }
    invShiftRows(s);
    invSubBytes(s);
    addRoundKey(s, rk(0));

    std::memcpy(out, s, 16);
}

std::vector<uint8_t> aes256::encrypt(const std::vector<uint8_t>& pt) {
    if (pt.size() != 16) throw std::invalid_argument("encrypt: block must be 16 bytes");
    uint8_t out[16];
    encrypt_block(pt.data(), out);
    return std::vector<uint8_t>(out, out+16);
}

std::vector<uint8_t> aes256::decrypt(const std::vector<uint8_t>& ct) {
    if (ct.size() != 16) throw std::invalid_argument("decrypt: block must be 16 bytes");
    uint8_t out[16];
    decrypt_block(ct.data(), out);
    return std::vector<uint8_t>(out, out+16);
}

std::vector<uint8_t> aes256::encryptCBC(const std::vector<uint8_t>& pt,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("encryptCBC: iv must be 16 bytes");
    if (pt.empty() || (pt.size() % 16) != 0) throw std::invalid_argument("encryptCBC: pt multiple of 16");

    std::vector<uint8_t> out(pt.size());
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < pt.size(); off += 16) {
        uint8_t blk[16];
        for (int i = 0; i < 16; i++) blk[i] = pt[off+i] ^ prev[i];
        encrypt_block(blk, &out[off]);
        std::memcpy(prev, &out[off], 16);
    }
    return out;
}

std::vector<uint8_t> aes256::decryptCBC(const std::vector<uint8_t>& ct,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("decryptCBC: iv must be 16 bytes");
    if (ct.empty() || (ct.size() % 16) != 0) throw std::invalid_argument("decryptCBC: ct multiple of 16");

    std::vector<uint8_t> out(ct.size());
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < ct.size(); off += 16) {
        uint8_t dec[16];
        decrypt_block(&ct[off], dec);
        for (int i = 0; i < 16; i++) out[off+i] = dec[i] ^ prev[i];
        std::memcpy(prev, &ct[off], 16);
    }
    return out;
}

// CTS mode (Cipher Text Stealing) - RFC 3962 for Kerberos (AES-256)
std::vector<uint8_t> aes256::encryptCTS(const std::vector<uint8_t>& pt,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("encryptCTS: iv must be 16 bytes");
    if (pt.size() < 16)  throw std::invalid_argument("encryptCTS: plaintext must be >= 16 bytes");

    size_t n = pt.size();

    if (n == 16) {
        uint8_t blk[16], out[16];
        for (int i = 0; i < 16; i++) blk[i] = pt[i] ^ iv[i];
        encrypt_block(blk, out);
        return std::vector<uint8_t>(out, out + 16);
    }

    size_t nBlocks = (n + 15) / 16;
    std::vector<uint8_t> padded = pt;
    padded.resize(nBlocks * 16, 0);

    std::vector<uint8_t> cbcOut(nBlocks * 16);
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < nBlocks * 16; off += 16) {
        uint8_t blk[16];
        for (int i = 0; i < 16; i++) blk[i] = padded[off + i] ^ prev[i];
        encrypt_block(blk, &cbcOut[off]);
        std::memcpy(prev, &cbcOut[off], 16);
    }

    std::vector<uint8_t> result(n);
    size_t penultimateOff = (nBlocks - 2) * 16;
    size_t lastOff = (nBlocks - 1) * 16;
    size_t lastPlainLen = n - penultimateOff - 16;

    if (nBlocks > 2)
        std::memcpy(result.data(), cbcOut.data(), penultimateOff);

    std::memcpy(result.data() + penultimateOff, &cbcOut[lastOff], 16);
    std::memcpy(result.data() + penultimateOff + 16, &cbcOut[penultimateOff], lastPlainLen);

    return result;
}

std::vector<uint8_t> aes256::decryptCTS(const std::vector<uint8_t>& ct,
                                        const std::vector<uint8_t>& iv) {
    if (iv.size() != 16) throw std::invalid_argument("decryptCTS: iv must be 16 bytes");
    if (ct.size() < 16)  throw std::invalid_argument("decryptCTS: ciphertext must be >= 16 bytes");

    size_t n = ct.size();

    if (n == 16) {
        uint8_t dec[16];
        decrypt_block(ct.data(), dec);
        std::vector<uint8_t> out(16);
        for (int i = 0; i < 16; i++) out[i] = dec[i] ^ iv[i];
        return out;
    }

    size_t nBlocks = (n + 15) / 16;
    size_t penultimateOff = (nBlocks - 2) * 16;
    size_t lastPlainLen = n - penultimateOff - 16;

    uint8_t Cn_1[16];
    std::memcpy(Cn_1, ct.data() + penultimateOff, 16);

    uint8_t Cn_2_trunc[16] = {};
    std::memcpy(Cn_2_trunc, ct.data() + penultimateOff + 16, lastPlainLen);

    uint8_t decCn1[16];
    decrypt_block(Cn_1, decCn1);

    uint8_t Cn_2_full[16];
    std::memcpy(Cn_2_full, Cn_2_trunc, lastPlainLen);
    std::memcpy(Cn_2_full + lastPlainLen, decCn1 + lastPlainLen, 16 - lastPlainLen);

    std::vector<uint8_t> result(n);
    uint8_t prev[16];
    std::memcpy(prev, iv.data(), 16);

    for (size_t off = 0; off < penultimateOff; off += 16) {
        uint8_t dec[16];
        decrypt_block(ct.data() + off, dec);
        for (int i = 0; i < 16; i++) result[off + i] = dec[i] ^ prev[i];
        std::memcpy(prev, ct.data() + off, 16);
    }

    {
        uint8_t dec[16];
        decrypt_block(Cn_2_full, dec);
        for (int i = 0; i < 16; i++) result[penultimateOff + i] = dec[i] ^ prev[i];
        std::memcpy(prev, Cn_2_full, 16);
    }

    for (size_t i = 0; i < lastPlainLen; i++)
        result[penultimateOff + 16 + i] = decCn1[i] ^ prev[i];

    return result;
}

void aes256::gcm_mul(std::array<uint8_t,16>& X) const {
#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        gcm_mul_clmul(X, m_H);
        return;
    }
#endif
    uint64_t Xhi = load_be64(X.data());
    uint64_t Xlo = load_be64(X.data() + 8);
    
    uint64_t Hhi = load_be64(m_H);
    uint64_t Hlo = load_be64(m_H + 8);
    
    uint64_t Zhi = 0, Zlo = 0;
    
    for (int i = 0; i < 128; i++) {
        uint64_t xi;
        if (i < 64)
            xi = (Xhi >> (63 - i)) & 1;
        else
            xi = (Xlo >> (63 - (i - 64))) & 1;
        
        if (xi) {
            Zhi ^= Hhi;
            Zlo ^= Hlo;
        }
        
        uint64_t lsb = Hlo & 1;
        Hlo = (Hlo >> 1) | (Hhi << 63);
        Hhi = Hhi >> 1;
        
        if (lsb) {
            Hhi ^= 0xE100000000000000ULL;
        }
    }
    
    store_be64(X.data(), Zhi);
    store_be64(X.data() + 8, Zlo);
}

void aes256::ghash(std::array<uint8_t,16>& out,
                   const uint8_t* aad, size_t aad_len,
                   const uint8_t* ct,  size_t ct_len) const {
    std::array<uint8_t,16> Y{};
    Y.fill(0);

    auto update = [&](const uint8_t* data, size_t len) {
        while (len >= 16) {
            for (int i = 0; i < 16; i++) Y[i] ^= data[i];
            gcm_mul(Y);
            data += 16;
            len  -= 16;
        }
        if (len > 0) {
            std::array<uint8_t,16> last{};
            last.fill(0);
            std::memcpy(last.data(), data, len);
            for (int i = 0; i < 16; i++) Y[i] ^= last[i];
            gcm_mul(Y);
        }
    };

    update(aad, aad_len);
    update(ct,  ct_len);

    uint8_t lenblk[16] = {0};
    store_be64(lenblk,     uint64_t(aad_len) * 8);
    store_be64(lenblk + 8, uint64_t(ct_len)  * 8);

    for (int i = 0; i < 16; i++) Y[i] ^= lenblk[i];
    gcm_mul(Y);

    out = Y;
}

bool aes256::aes_gcm_encrypt(
    const uint8_t iv12[12],
    const uint8_t* aad, size_t aad_len,
    const uint8_t* pt,  size_t pt_len,
    uint8_t* ct,
    uint8_t tag16[16]
) const {
    uint8_t J0[16] = {0};
    std::memcpy(J0, iv12, 12);
    J0[15] = 1;

#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        aesni_gcm_ctr_ghash_256(pt, ct, pt_len, J0, m_roundKeyBytes.data(),
                                 aad, aad_len, m_H, tag16);
        return true;
    }
#endif
    {
        uint8_t ctr[16];
        std::memcpy(ctr, J0, 16);

        size_t off = 0;
        while (off < pt_len) {
            inc32(ctr);
            uint8_t S[16];
            encrypt_block(ctr, S);

            size_t take = std::min<size_t>(16, pt_len - off);
            for (size_t i = 0; i < take; i++)
                ct[off+i] = pt[off+i] ^ S[i];
            off += take;
        }
    }

    std::array<uint8_t,16> Sgh{};
    Sgh.fill(0);
    ghash(Sgh, aad, aad_len, ct, pt_len);

    uint8_t EJ0[16];
    encrypt_block(J0, EJ0);

    for (int i = 0; i < 16; i++)
        tag16[i] = EJ0[i] ^ Sgh[i];

    return true;
}

bool aes256::aes_gcm_decrypt(
    const uint8_t iv12[12],
    const uint8_t* aad, size_t aad_len,
    const uint8_t* ct,  size_t ct_len,
    const uint8_t tag16[16],
    uint8_t* pt
) const {
    uint8_t J0[16] = {0};
    std::memcpy(J0, iv12, 12);
    J0[15] = 1;

#ifdef NETPLUS_X86
    if (g_has_hw_aes) {
        // See aes128::aes_gcm_decrypt / aesni_gcm_ctr_ghash_decrypt_256's
        // comment: fused GHASH+CTR-decrypt in one pass.
        uint8_t expTag[16];
        aesni_gcm_ctr_ghash_decrypt_256(ct, pt, ct_len, J0, m_roundKeyBytes.data(),
                                         aad, aad_len, m_H, expTag);
        if (!ct_memeq(expTag, tag16, 16)) {
            std::memset(pt, 0, ct_len);
            return false;
        }
        return true;
    }
#endif

    std::array<uint8_t,16> Sgh{};
    Sgh.fill(0);
    ghash(Sgh, aad, aad_len, ct, ct_len);

    uint8_t EJ0[16];
    encrypt_block(J0, EJ0);

    uint8_t expTag[16];
    for (int i = 0; i < 16; i++)
        expTag[i] = EJ0[i] ^ Sgh[i];

    if (!ct_memeq(expTag, tag16, 16))
        return false;

    uint8_t ctr[16];
    std::memcpy(ctr, J0, 16);

    size_t off = 0;
    while (off < ct_len) {
        inc32(ctr);
        uint8_t S[16];
        encrypt_block(ctr, S);

        size_t take = std::min<size_t>(16, ct_len - off);
        for (size_t i = 0; i < take; i++)
            pt[off+i] = ct[off+i] ^ S[i];
        off += take;
    }

    return true;
}

// ============================================================
// AES-GCM Static wrapper class
// ============================================================

bool aes_gcm::encrypt(
    const aes& cipher,
    const uint8_t iv12[12],
    const uint8_t* aad, size_t aad_len,
    const uint8_t* pt,  size_t pt_len,
    uint8_t* ct,
    uint8_t tag16[16]
) {
    return cipher.aes_gcm_encrypt(iv12, aad, aad_len, pt, pt_len, ct, tag16);
}

bool aes_gcm::decrypt(
    const aes& cipher,
    const uint8_t iv12[12],
    const uint8_t* aad, size_t aad_len,
    const uint8_t* ct,  size_t ct_len,
    const uint8_t tag16[16],
    uint8_t* pt
) {
    return cipher.aes_gcm_decrypt(iv12, aad, aad_len, ct, ct_len, tag16, pt);
}

} // namespace netplus
