Case Study 15.2: AES-NI Encryption — Hardware-Accelerated AES in Assembly

The Goal

This case study implements AES-128 encryption in assembly using AES-NI instructions. We cover the complete implementation: key schedule expansion, single-block encryption, and a CTR-mode stream cipher. The result is a functionally correct, constant-time AES implementation in approximately 60 lines of NASM.

AES-128 Background

AES-128 uses a 128-bit (16-byte) key and operates on 128-bit blocks. The algorithm:

1. Key expansion: Expand the 16-byte key into 11 round keys (176 bytes total)
2. Block encryption: - Initial key whitening: plaintext XOR round_key[0] - 9 rounds of AESENC (SubBytes → ShiftRows → MixColumns → XOR round key) - Final round: AESENCLAST (SubBytes → ShiftRows → XOR round key, no MixColumns)

Each AES-NI instruction performs one complete round in hardware. The total for AES-128: one PXOR + nine AESENC + one AESENCLAST = 11 XOR-style operations.

Part 1: Key Schedule Expansion

The key schedule takes the original 128-bit key and derives 10 additional round keys. Each round key is derived from the previous one using a non-linear transformation.

The Key Expansion Formula

For AES-128, round key derivation uses AESKEYGENASSIST which computes SubWord(RotWord(w)) on the high dword, then XORs with round constant (RCON). The formula:

round_key[i][0] = round_key[i-1][0] XOR SubWord(RotWord(round_key[i-1][3])) XOR RCON[i]
round_key[i][1] = round_key[i-1][1] XOR round_key[i][0]
round_key[i][2] = round_key[i-1][2] XOR round_key[i][1]
round_key[i][3] = round_key[i-1][3] XOR round_key[i][2]

The AESKEYGENASSIST instruction computes the XOR with RCON and the S-box substitution; we implement the XOR accumulation with shifts and XORs.

NASM Macro: KEY_EXPAND

; KEY_EXPAND: generate the next AES-128 round key
; Parameters:
;   %1 = destination address for this round key
;   %2 = RCON value (round constant: 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36)
; On entry:  xmm0 = previous round key
; On exit:   xmm0 = new round key (also stored at [%1])
; Clobbers:  xmm1, xmm2

%macro KEY_EXPAND 2
    ; AESKEYGENASSIST: computes SubWord(RotWord(xmm0[127:96])) XOR RCON
    ; The result we need is in bits [127:96] of xmm1 (the high dword)
    aeskeygenassist xmm1, xmm0, %2

    ; We need xmm1[127:96] broadcast to all four 32-bit lanes
    ; PSHUFD with imm 0xFF selects element 3 (top dword) to all positions
    pshufd  xmm1, xmm1, 0xFF
    ; xmm1 = [T, T, T, T] where T = SubWord(RotWord(prev[3])) XOR RCON

    ; Now build the XOR accumulation chain:
    ; new_key[0] = xmm0[0] XOR T
    ; new_key[1] = xmm0[1] XOR xmm0[0] XOR T
    ; new_key[2] = xmm0[2] XOR xmm0[1] XOR xmm0[0] XOR T
    ; new_key[3] = xmm0[3] XOR xmm0[2] XOR xmm0[1] XOR xmm0[0] XOR T
    ;
    ; This is a prefix XOR applied to the previous key, then XOR with T.
    ; Implement by shifting xmm0 and XORing cumulatively:

    movaps  xmm2, xmm0
    pslldq  xmm2, 4         ; xmm2 = [0, xmm0[0], xmm0[1], xmm0[2]]
    xorps   xmm0, xmm2      ; xmm0[i] ^= xmm0[i-1]

    movaps  xmm2, xmm0
    pslldq  xmm2, 8         ; shift by 2 dwords
    xorps   xmm0, xmm2      ; complete prefix XOR accumulation
    ; now xmm0[i] = XOR of all previous round key dwords up to position i

    xorps   xmm0, xmm1      ; XOR with the AESKEYGENASSIST result

    movaps  [%1], xmm0      ; store the new round key
%endmacro

Complete Key Schedule Function

; void aes128_key_expand(uint8_t *key_schedule, const uint8_t *key)
; rdi = key_schedule output (176 bytes: 11 * 16)
; rsi = original 16-byte key

section .text
global aes128_key_expand

aes128_key_expand:
    ; Load the original key as round key 0
    movdqu  xmm0, [rsi]
    movaps  [rdi], xmm0          ; key_schedule[0] = original key

    ; Generate round keys 1-10 using the macro
    KEY_EXPAND  [rdi + 16],  0x01
    KEY_EXPAND  [rdi + 32],  0x02
    KEY_EXPAND  [rdi + 48],  0x04
    KEY_EXPAND  [rdi + 64],  0x08
    KEY_EXPAND  [rdi + 80],  0x10
    KEY_EXPAND  [rdi + 96],  0x20
    KEY_EXPAND  [rdi + 112], 0x40
    KEY_EXPAND  [rdi + 128], 0x80
    KEY_EXPAND  [rdi + 144], 0x1b
    KEY_EXPAND  [rdi + 160], 0x36

    ret

Key schedule trace for key = 2b7e1516 28aed2a6 abf71588 09cf4f3c:

Round 0:  2b7e1516 28aed2a6 abf71588 09cf4f3c
Round 1:  a0fafe17 88542cb1 23a33939 2a6c7605
Round 2:  f2c295f2 7a96b943 5935807a 7359f67f
Round 3:  3d80477d 4716fe3e 1e237e44 6d7a883b
...
Round 10: 13111d7f e3944a17 f307a78b 4d2b30c5

(These match the FIPS 197 AES standard test vectors.)

Part 2: Single Block Encryption

; void aes128_encrypt_block(uint8_t *ct, const uint8_t *pt, const uint8_t *ks)
; rdi = ciphertext (16 bytes output)
; rsi = plaintext (16 bytes input)
; rdx = key schedule (176 bytes, from aes128_key_expand)

section .text
global aes128_encrypt_block

aes128_encrypt_block:
    ; Load plaintext
    movdqu  xmm0, [rsi]

    ; Initial key whitening (round 0)
    pxor    xmm0, [rdx]

    ; Rounds 1-9: AESENC
    aesenc  xmm0, [rdx + 16]
    aesenc  xmm0, [rdx + 32]
    aesenc  xmm0, [rdx + 48]
    aesenc  xmm0, [rdx + 64]
    aesenc  xmm0, [rdx + 80]
    aesenc  xmm0, [rdx + 96]
    aesenc  xmm0, [rdx + 112]
    aesenc  xmm0, [rdx + 128]
    aesenc  xmm0, [rdx + 144]

    ; Final round: AESENCLAST (no MixColumns)
    aesenclast xmm0, [rdx + 160]

    ; Store ciphertext
    movdqu  [rdi], xmm0
    ret

Verification against FIPS 197 test vector: - Key: 2b7e151628aed2a6abf7158809cf4f3c - Plaintext: 6bc1bee22e409f96e93d7e117393172a - Expected ciphertext: 3ad77bb40d7a3660a89ecaf32466ef97

The AESENC instruction is specified to produce exactly the AES standard output; any correct key schedule produces a correct ciphertext matching FIPS 197.

Part 3: AES-128 CTR Mode

CTR (Counter) mode turns AES into a stream cipher. A counter block is encrypted, producing a keystream; the keystream is XORed with plaintext to produce ciphertext. The same operation (just XOR with the same keystream) decrypts.

Counter block: [nonce (8 bytes)][counter (8 bytes)]
Block 0: encrypt(nonce || 0), XOR with plaintext[0:15]
Block 1: encrypt(nonce || 1), XOR with plaintext[16:31]
...

; void aes128_ctr_encrypt(uint8_t *out, const uint8_t *in, size_t len,
;                         const uint8_t *ks, const uint8_t *nonce)
; rdi = output buffer
; rsi = input buffer
; rdx = length in bytes
; rcx = key schedule (176 bytes)
; r8  = nonce (8 bytes — we build the full 16-byte counter block internally)

section .data
align 16
ctr_increment: dq 1, 0       ; little-endian: increment low qword by 1

section .text
global aes128_ctr_encrypt

aes128_ctr_encrypt:
    push    rbp
    mov     rbp, rsp
    push    rbx
    push    r12
    push    r13
    push    r14
    push    r15
    sub     rsp, 16             ; space for counter block on stack
    and     rsp, -16            ; 16-byte align

    ; Build initial counter block on stack: [nonce | counter=0]
    movq    xmm5, [r8]          ; load 8-byte nonce into low qword of xmm5
    pxor    xmm6, xmm6          ; zero counter
    ; xmm5 = nonce in low 64 bits
    ; We need: [nonce (low 64 bits)][counter (high 64 bits)]
    punpcklqdq xmm5, xmm6       ; xmm5 = [nonce, 0] = initial counter block

    movdqa  [rsp], xmm5         ; save counter block to stack

    ; Load the counter increment vector
    movdqa  xmm7, [rel ctr_increment]  ; xmm7 = [1, 0] (add 1 to low qword)

    ; Save key schedule pointer
    mov     r12, rcx

    ; Process full 16-byte blocks
    xor     r13, r13            ; offset into buffers
    mov     r14, rdx            ; remaining bytes

.block_loop:
    cmp     r14, 16
    jb      .partial_block

    ; Load current counter block
    movdqa  xmm0, [rsp]

    ; Encrypt the counter block (same as aes128_encrypt_block inline)
    pxor    xmm0, [r12]
    aesenc  xmm0, [r12 + 16]
    aesenc  xmm0, [r12 + 32]
    aesenc  xmm0, [r12 + 48]
    aesenc  xmm0, [r12 + 64]
    aesenc  xmm0, [r12 + 80]
    aesenc  xmm0, [r12 + 96]
    aesenc  xmm0, [r12 + 112]
    aesenc  xmm0, [r12 + 128]
    aesenc  xmm0, [r12 + 144]
    aesenclast xmm0, [r12 + 160]
    ; xmm0 = keystream block

    ; XOR with plaintext
    movdqu  xmm1, [rsi + r13]
    pxor    xmm0, xmm1
    movdqu  [rdi + r13], xmm0

    ; Increment counter (add 1 to low 64 bits)
    movdqa  xmm5, [rsp]
    paddq   xmm5, xmm7           ; xmm5.lo += 1
    movdqa  [rsp], xmm5

    add     r13, 16
    sub     r14, 16
    jmp     .block_loop

.partial_block:
    ; Handle remaining 1-15 bytes
    test    r14, r14
    jz      .done

    ; Encrypt counter block to get keystream
    movdqa  xmm0, [rsp]
    pxor    xmm0, [r12]
    aesenc  xmm0, [r12 + 16]
    aesenc  xmm0, [r12 + 32]
    aesenc  xmm0, [r12 + 48]
    aesenc  xmm0, [r12 + 64]
    aesenc  xmm0, [r12 + 80]
    aesenc  xmm0, [r12 + 96]
    aesenc  xmm0, [r12 + 112]
    aesenc  xmm0, [r12 + 128]
    aesenc  xmm0, [r12 + 144]
    aesenclast xmm0, [r12 + 160]
    ; Store keystream to a temp buffer on stack
    movdqa  [rsp - 32], xmm0     ; careful: rsp is already moved down by 16

    ; XOR byte by byte
    xor     r15, r15
.partial_byte_loop:
    cmp     r15, r14
    jae     .done
    movzx   eax, byte [rsi + r13 + r15]
    xor     al, [rsp - 32 + r15]
    mov     [rdi + r13 + r15], al
    inc     r15
    jmp     .partial_byte_loop

.done:
    add     rsp, 16
    pop     r15
    pop     r14
    pop     r13
    pop     r12
    pop     rbx
    pop     rbp
    ret

Security note: The partial block loop processes one byte at a time. For a constant-time implementation, the partial block handling should avoid data-dependent branching (which jae on r15 introduces). Production implementations use branchless techniques for the final partial block.

Performance Analysis

Throughput Measurements

The single-block implementation above achieves roughly one AES-NI instruction per cycle — but AESENC has a 4-cycle latency. Processing one block means 10 dependent AESENC instructions at 4 cycles each = 40 cycles minimum per block (16 bytes), or 2.5 cycles/byte.

Pipelining: The Key to Performance

Since CTR mode blocks are independent (each counter value is independent), we can pipeline 4 or more blocks simultaneously:

; Process 4 CTR blocks in parallel:
movdqa  xmm0, ctr0     ; counter block 0
movdqa  xmm1, ctr1     ; counter block 1 (= ctr0 + 1)
movdqa  xmm2, ctr2     ; counter block 2 (= ctr0 + 2)
movdqa  xmm3, ctr3     ; counter block 3 (= ctr0 + 3)

pxor    xmm0, [ks]     ; round 0 whitening for all 4 blocks simultaneously
pxor    xmm1, [ks]
pxor    xmm2, [ks]
pxor    xmm3, [ks]

aesenc  xmm0, [ks+16]  ; round 1: all 4 blocks
aesenc  xmm1, [ks+16]  ; processor executes these in parallel
aesenc  xmm2, [ks+16]
aesenc  xmm3, [ks+16]
; ... continue for 9 rounds ...

aesenclast xmm0, [ks+160]
aesenclast xmm1, [ks+160]
aesenclast xmm2, [ks+160]
aesenclast xmm3, [ks+160]
; XOR with plaintext...

With 4-wide pipelining, the throughput improves ~4× because the CPU can execute the independent AESENC instructions on different pipelines. Modern Intel CPUs have 2 AES execution units, so pipelining 2+ blocks in parallel saturates the units and achieves ~0.5 cycles/AESENC throughput.

Correctness Verification

To verify the implementation against standard test vectors:

FIPS 197, Appendix B:

Key:       2b7e151628aed2a6 abf7158809cf4f3c
Plaintext: 3243f6a8885a308d 313198a2e0370734
Expected:  3925841d02dc09fb dc118597196a0b32

Running aes128_key_expand on this key and aes128_encrypt_block on this plaintext must produce exactly the expected ciphertext. Any deviation indicates a bug in the key schedule or the encryption loop — likely a wrong RCON value or incorrect PSHUFD immediate.

What the Instructions Actually Do

AESENC xmm_dst, xmm_src performs one complete AES round in hardware:

1. SubBytes: Apply the AES S-box to each of the 16 bytes. The S-box is a bijective non-linear function implemented in hardware using the multiplicative inverse in GF(2^8). This is what software T-table implementations compute with table lookups — and what makes software AES vulnerable to cache-timing attacks.

2. ShiftRows: Cyclically left-shift each row of the 4×4 byte state matrix by 0, 1, 2, 3 positions.

3. MixColumns: Apply a linear transformation to each column of the state, treating each column as a polynomial in GF(2^8). This provides diffusion.

4. AddRoundKey: XOR the state with the round key.

All four transformations execute in a single instruction. The hardware can implement SubBytes as a combinational circuit rather than a table lookup, making execution time independent of the data values — hence no cache-timing vulnerability.

AESENCLAST skips MixColumns (the final round of AES intentionally omits it for the decrypt-side to work cleanly) but performs the other three steps.

The Security Property

The critical security guarantee of AES-NI:

The execution time of AESENC and AESENCLAST does not depend on the value of the key material or plaintext.

This is guaranteed by the Intel and AMD architectures. The instruction executes in a fixed number of cycles regardless of input values. Software AES using lookup tables cannot make this guarantee because which cache lines are accessed depends on the key — and cache access timing is observable by a co-located process.

This is why all modern TLS implementations, OpenSSH, and disk encryption software use AES-NI when available. The performance improvement (10-40×) is secondary to the elimination of a fundamental side-channel vulnerability.