Clear ML-KEM Keccak stack scratch

tob-joe · codex · tob-joe · commit eaa80926eac6 · 2026-07-07T17:20:07.000-04:00
Wipe the AVX2, AArch64 scalar-hybrid, and Armv8.1-M MVE x4 Keccak
scratch frames before returning, and refresh the covered HOL Light proof
copies for the changed instruction streams.

This keeps state-derived stack scratch from surviving past the permutation
return path.

Verification:

- nix develop .#hol_light-cross-aarch64 -c make -B -C proofs/hol_light/aarch64 mlkem/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.correct

- nix develop .#hol_light-cross-aarch64 -c make -B -C proofs/hol_light/aarch64 mlkem/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.correct

- nix develop .#hol_light-cross-x86_64 -c make -B -C proofs/hol_light/x86_64 mlkem/keccak_f1600_x4_avx2_asm.correct

Co-authored-by: Codex &lt;codex@openai.com&gt;
Signed-off-by: Joe Doyle &lt;joseph.doyle@trailofbits.com&gt;
diff --git a/dev/fips202/aarch64/src/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.S b/dev/fips202/aarch64/src/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.S
@@ -335,6 +335,15 @@
     add sp, sp, #(STACK_SIZE)
 .endm
 
+.macro clear_stack
+    eor v0.16b, v0.16b, v0.16b
+    .set .Lclear_stack_offset, 0
+    .rept STACK_SIZE / 16
+    str q0, [sp, #.Lclear_stack_offset]
+    .set .Lclear_stack_offset, .Lclear_stack_offset + 16
+    .endr
+.endm
+
 .macro eor5 dst, src0, src1, src2, src3, src4
     eor \dst, \src0, \src1
     eor \dst, \dst,  \src2
@@ -2453,6 +2462,7 @@ keccak_f1600_x4_v8a_scalar_hybrid_done:
 
     restore_vregs
     restore_gprs
+    clear_stack
     free_stack
     ret
 
diff --git a/dev/fips202/aarch64/src/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.S b/dev/fips202/aarch64/src/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.S
@@ -354,6 +354,15 @@
     add sp, sp, #(STACK_SIZE)
 .endm
 
+.macro clear_stack
+    eor v0.16b, v0.16b, v0.16b
+    .set .Lclear_stack_offset, 0
+    .rept STACK_SIZE / 16
+    str q0, [sp, #.Lclear_stack_offset]
+    .set .Lclear_stack_offset, .Lclear_stack_offset + 16
+    .endr
+.endm
+
 .macro eor5 dst, src0, src1, src2, src3, src4
     eor \dst, \src0, \src1
     eor \dst, \dst,  \src2
@@ -2284,6 +2293,7 @@ keccak_f1600_x4_v8a_v84a_scalar_hybrid_done:
 
     restore_vregs
     restore_gprs
+    clear_stack
     free_stack
     ret
 
diff --git a/dev/fips202/armv81m/src/keccak_f1600_x4_mve.S b/dev/fips202/armv81m/src/keccak_f1600_x4_mve.S
@@ -1626,6 +1626,14 @@ MLK_ASM_FN_SYMBOL(keccak_f1600_x4_mve_asm)
 
     le lr, keccak_f1600_x4_mve_asm_roundstart
 keccak_f1600_x4_mve_asm_roundend:
+    // Clear the 0x80-byte scratch frame before releasing it.
+    movs r0, #0
+    vdup.32 q0, r0
+    .set .Lclear_stack_offset, 0
+    .rept 8
+    vstrw.32 q0, [r13, #.Lclear_stack_offset]
+    .set .Lclear_stack_offset, .Lclear_stack_offset + 16
+    .endr
     add sp, #8*16
 
     vpop {d8-d15}
diff --git a/dev/fips202/x86_64/src/keccak_f1600_x4_avx2_asm.S b/dev/fips202/x86_64/src/keccak_f1600_x4_avx2_asm.S
@@ -651,6 +651,14 @@ MLK_ASM_FN_SYMBOL(keccak_f1600_x4_avx2_asm)
    vmovhpd %xmm3, 0x188(%rdi)
    vmovq %xmm15, 0x250(%rdi)
    vmovhpd %xmm15, 0x318(%rdi)
+
+   // Clear the scratch frame before releasing it.
+   vpxor %ymm0, %ymm0, %ymm0
+   .set .Lclear_stack_offset, 0
+   .rept 24
+   vmovdqu %ymm0, .Lclear_stack_offset(%rsp)
+   .set .Lclear_stack_offset, .Lclear_stack_offset + 32
+   .endr
    movq   %r11, %rsp
    ret
 
diff --git a/mlkem/src/fips202/native/aarch64/src/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.S b/mlkem/src/fips202/native/aarch64/src/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.S
@@ -1063,6 +1063,21 @@ Lkeccak_f1600_x4_v8a_scalar_hybrid_done:
         ldp x29, x30, [sp, #0x80]
         .cfi_restore x29
         .cfi_restore x30
+        eor v0.16b, v0.16b, v0.16b
+        str q0, [sp]
+        str q0, [sp, #0x10]
+        str q0, [sp, #0x20]
+        str q0, [sp, #0x30]
+        str q0, [sp, #0x40]
+        str q0, [sp, #0x50]
+        str q0, [sp, #0x60]
+        str q0, [sp, #0x70]
+        str q0, [sp, #0x80]
+        str q0, [sp, #0x90]
+        str q0, [sp, #0xa0]
+        str q0, [sp, #0xb0]
+        str q0, [sp, #0xc0]
+        str q0, [sp, #0xd0]
         add sp, sp, #0xe0
         .cfi_adjust_cfa_offset -0xe0
         ret
diff --git a/mlkem/src/fips202/native/aarch64/src/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.S b/mlkem/src/fips202/native/aarch64/src/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.S
@@ -972,6 +972,21 @@ Lkeccak_f1600_x4_v8a_v84a_scalar_hybrid_done:
         ldp x29, x30, [sp, #0x80]
         .cfi_restore x29
         .cfi_restore x30
+        eor v0.16b, v0.16b, v0.16b
+        str q0, [sp]
+        str q0, [sp, #0x10]
+        str q0, [sp, #0x20]
+        str q0, [sp, #0x30]
+        str q0, [sp, #0x40]
+        str q0, [sp, #0x50]
+        str q0, [sp, #0x60]
+        str q0, [sp, #0x70]
+        str q0, [sp, #0x80]
+        str q0, [sp, #0x90]
+        str q0, [sp, #0xa0]
+        str q0, [sp, #0xb0]
+        str q0, [sp, #0xc0]
+        str q0, [sp, #0xd0]
         add sp, sp, #0xe0
         .cfi_adjust_cfa_offset -0xe0
         ret
diff --git a/mlkem/src/fips202/native/armv81m/src/keccak_f1600_x4_mve.S b/mlkem/src/fips202/native/armv81m/src/keccak_f1600_x4_mve.S
@@ -681,6 +681,13 @@ Lkeccak_f1600_x4_mve_asm_roundend_pre:
         le lr, Lkeccak_f1600_x4_mve_asm_roundstart @ imm = #-0x8c0
 
 Lkeccak_f1600_x4_mve_asm_roundend:
+        movs r0, #0x0
+        vdup.32 q0, r0
+        .set .Lkeccak_f1600_x4_mve_clear_stack_offset, 0
+        .rept 8
+        vstrw.32 q0, [sp, #.Lkeccak_f1600_x4_mve_clear_stack_offset]
+        .set .Lkeccak_f1600_x4_mve_clear_stack_offset, .Lkeccak_f1600_x4_mve_clear_stack_offset + 16
+        .endr
         add sp, #0x80
         .cfi_adjust_cfa_offset -0x80
         vpop {d8, d9, d10, d11, d12, d13, d14, d15}
@@ -705,7 +712,6 @@ Lkeccak_f1600_x4_mve_asm_roundend:
         .cfi_restore lr
         .cfi_adjust_cfa_offset -0x24
         .cfi_endproc
-        nop
 
 MLK_ASM_FN_SIZE(keccak_f1600_x4_mve_asm)
 
diff --git a/mlkem/src/fips202/native/x86_64/src/keccak_f1600_x4_avx2_asm.S b/mlkem/src/fips202/native/x86_64/src/keccak_f1600_x4_avx2_asm.S
@@ -473,6 +473,12 @@ LLkeccak_f1600_x4_avx2:
         vmovhpd %xmm3, 0x188(%rdi)
         vmovq %xmm15, 0x250(%rdi)
         vmovhpd %xmm15, 0x318(%rdi)
+        vpxor %ymm0, %ymm0, %ymm0
+        .set .Lkeccak_f1600_x4_avx2_clear_stack_offset, 0
+        .rept 24
+        vmovdqu %ymm0, .Lkeccak_f1600_x4_avx2_clear_stack_offset(%rsp)
+        .set .Lkeccak_f1600_x4_avx2_clear_stack_offset, .Lkeccak_f1600_x4_avx2_clear_stack_offset + 32
+        .endr
         movq %r11, %rsp
         .cfi_def_cfa_register %rsp
         retq
diff --git a/proofs/hol_light/aarch64/mlkem/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.S b/proofs/hol_light/aarch64/mlkem/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.S
@@ -1065,6 +1065,21 @@ Lkeccak_f1600_x4_v8a_scalar_hybrid_done:
         ldp x29, x30, [sp, #0x80]
         .cfi_restore x29
         .cfi_restore x30
+        eor v0.16b, v0.16b, v0.16b
+        str q0, [sp]
+        str q0, [sp, #0x10]
+        str q0, [sp, #0x20]
+        str q0, [sp, #0x30]
+        str q0, [sp, #0x40]
+        str q0, [sp, #0x50]
+        str q0, [sp, #0x60]
+        str q0, [sp, #0x70]
+        str q0, [sp, #0x80]
+        str q0, [sp, #0x90]
+        str q0, [sp, #0xa0]
+        str q0, [sp, #0xb0]
+        str q0, [sp, #0xc0]
+        str q0, [sp, #0xd0]
         add sp, sp, #0xe0
         .cfi_adjust_cfa_offset -0xe0
         ret
diff --git a/proofs/hol_light/aarch64/mlkem/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.S b/proofs/hol_light/aarch64/mlkem/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.S
@@ -973,6 +973,21 @@ Lkeccak_f1600_x4_v8a_v84a_scalar_hybrid_done:
         ldp x29, x30, [sp, #0x80]
         .cfi_restore x29
         .cfi_restore x30
+        eor v0.16b, v0.16b, v0.16b
+        str q0, [sp]
+        str q0, [sp, #0x10]
+        str q0, [sp, #0x20]
+        str q0, [sp, #0x30]
+        str q0, [sp, #0x40]
+        str q0, [sp, #0x50]
+        str q0, [sp, #0x60]
+        str q0, [sp, #0x70]
+        str q0, [sp, #0x80]
+        str q0, [sp, #0x90]
+        str q0, [sp, #0xa0]
+        str q0, [sp, #0xb0]
+        str q0, [sp, #0xc0]
+        str q0, [sp, #0xd0]
         add sp, sp, #0xe0
         .cfi_adjust_cfa_offset -0xe0
         ret
diff --git a/proofs/hol_light/aarch64/proofs/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.ml b/proofs/hol_light/aarch64/proofs/keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.ml
@@ -987,6 +987,21 @@ let keccak_f1600_x4_v8a_scalar_mc = define_assert_from_elf
   0xa9466bf9;       (* arm_LDP X25 X26 SP (Immediate_Offset (iword (&96))) *)
   0xa94773fb;       (* arm_LDP X27 X28 SP (Immediate_Offset (iword (&112))) *)
   0xa9487bfd;       (* arm_LDP X29 X30 SP (Immediate_Offset (iword (&128))) *)
+  0x6e201c00;       (* arm_EOR_VEC Q0 Q0 Q0 128 *)
+  0x3d8003e0;       (* arm_STR Q0 SP (Immediate_Offset (word 0)) *)
+  0x3d8007e0;       (* arm_STR Q0 SP (Immediate_Offset (word 16)) *)
+  0x3d800be0;       (* arm_STR Q0 SP (Immediate_Offset (word 32)) *)
+  0x3d800fe0;       (* arm_STR Q0 SP (Immediate_Offset (word 48)) *)
+  0x3d8013e0;       (* arm_STR Q0 SP (Immediate_Offset (word 64)) *)
+  0x3d8017e0;       (* arm_STR Q0 SP (Immediate_Offset (word 80)) *)
+  0x3d801be0;       (* arm_STR Q0 SP (Immediate_Offset (word 96)) *)
+  0x3d801fe0;       (* arm_STR Q0 SP (Immediate_Offset (word 112)) *)
+  0x3d8023e0;       (* arm_STR Q0 SP (Immediate_Offset (word 128)) *)
+  0x3d8027e0;       (* arm_STR Q0 SP (Immediate_Offset (word 144)) *)
+  0x3d802be0;       (* arm_STR Q0 SP (Immediate_Offset (word 160)) *)
+  0x3d802fe0;       (* arm_STR Q0 SP (Immediate_Offset (word 176)) *)
+  0x3d8033e0;       (* arm_STR Q0 SP (Immediate_Offset (word 192)) *)
+  0x3d8037e0;       (* arm_STR Q0 SP (Immediate_Offset (word 208)) *)
   0x910383ff;       (* arm_ADD SP SP (rvalue (word 224)) *)
   0xd65f03c0        (* arm_RET X30 *)
 ];;
@@ -1017,7 +1032,7 @@ let KECCAK_F1600_X4_V8A_SCALAR_PREAMBLE_LENGTH = new_definition
   `KECCAK_F1600_X4_V8A_SCALAR_PREAMBLE_LENGTH = 44`;;
 
 let KECCAK_F1600_X4_V8A_SCALAR_POSTAMBLE_LENGTH = new_definition
-  `KECCAK_F1600_X4_V8A_SCALAR_POSTAMBLE_LENGTH = 48`;;
+  `KECCAK_F1600_X4_V8A_SCALAR_POSTAMBLE_LENGTH = 108`;;
 
 let KECCAK_F1600_X4_V8A_SCALAR_CORE_START = new_definition
   `KECCAK_F1600_X4_V8A_SCALAR_CORE_START = KECCAK_F1600_X4_V8A_SCALAR_PREAMBLE_LENGTH`;;
@@ -1483,7 +1498,7 @@ let KECCAK_F1600_X4_V8A_SCALAR_SUBROUTINE_CORRECT = prove
     (WORDLIST_FROM_MEMORY_CONV THENC
      ONCE_DEPTH_CONV NORMALIZE_RELATIVE_ADDRESS_CONV) in
   CONV_TAC TWEAK_CONV THEN
-  ARM_ADD_RETURN_STACK_TAC ~pre_post_nsteps:(11,11) KECCAK_F1600_X4_V8A_SCALAR_EXEC
+  ARM_ADD_RETURN_STACK_TAC ~pre_post_nsteps:(11,26) KECCAK_F1600_X4_V8A_SCALAR_EXEC
    (CONV_RULE TWEAK_CONV (CONV_RULE LENGTH_SIMPLIFY_CONV KECCAK_F1600_X4_V8A_SCALAR_CORRECT))
   `[D8; D9; D10; D11; D12; D13; D14; D15;
     X19; X20; X21; X22; X23; X24; X25; X26; X27; X28; X29; X30]` 224);;
diff --git a/proofs/hol_light/aarch64/proofs/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.ml b/proofs/hol_light/aarch64/proofs/keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.ml
@@ -893,6 +893,21 @@ let keccak_f1600_x4_v8a_v84a_scalar_mc = define_assert_from_elf
   0xa9466bf9;       (* arm_LDP X25 X26 SP (Immediate_Offset (iword (&96))) *)
   0xa94773fb;       (* arm_LDP X27 X28 SP (Immediate_Offset (iword (&112))) *)
   0xa9487bfd;       (* arm_LDP X29 X30 SP (Immediate_Offset (iword (&128))) *)
+  0x6e201c00;       (* arm_EOR_VEC Q0 Q0 Q0 128 *)
+  0x3d8003e0;       (* arm_STR Q0 SP (Immediate_Offset (word 0)) *)
+  0x3d8007e0;       (* arm_STR Q0 SP (Immediate_Offset (word 16)) *)
+  0x3d800be0;       (* arm_STR Q0 SP (Immediate_Offset (word 32)) *)
+  0x3d800fe0;       (* arm_STR Q0 SP (Immediate_Offset (word 48)) *)
+  0x3d8013e0;       (* arm_STR Q0 SP (Immediate_Offset (word 64)) *)
+  0x3d8017e0;       (* arm_STR Q0 SP (Immediate_Offset (word 80)) *)
+  0x3d801be0;       (* arm_STR Q0 SP (Immediate_Offset (word 96)) *)
+  0x3d801fe0;       (* arm_STR Q0 SP (Immediate_Offset (word 112)) *)
+  0x3d8023e0;       (* arm_STR Q0 SP (Immediate_Offset (word 128)) *)
+  0x3d8027e0;       (* arm_STR Q0 SP (Immediate_Offset (word 144)) *)
+  0x3d802be0;       (* arm_STR Q0 SP (Immediate_Offset (word 160)) *)
+  0x3d802fe0;       (* arm_STR Q0 SP (Immediate_Offset (word 176)) *)
+  0x3d8033e0;       (* arm_STR Q0 SP (Immediate_Offset (word 192)) *)
+  0x3d8037e0;       (* arm_STR Q0 SP (Immediate_Offset (word 208)) *)
   0x910383ff;       (* arm_ADD SP SP (rvalue (word 224)) *)
   0xd65f03c0        (* arm_RET X30 *)
 ];;
@@ -922,7 +937,7 @@ let KECCAK_F1600_X4_V8A_V84A_SCALAR_PREAMBLE_LENGTH = new_definition
   `KECCAK_F1600_X4_V8A_V84A_SCALAR_PREAMBLE_LENGTH = 44`;;
 
 let KECCAK_F1600_X4_V8A_V84A_SCALAR_POSTAMBLE_LENGTH = new_definition
-  `KECCAK_F1600_X4_V8A_V84A_SCALAR_POSTAMBLE_LENGTH = 48`;;
+  `KECCAK_F1600_X4_V8A_V84A_SCALAR_POSTAMBLE_LENGTH = 108`;;
 
 let KECCAK_F1600_X4_V8A_V84A_SCALAR_CORE_START = new_definition
   `KECCAK_F1600_X4_V8A_V84A_SCALAR_CORE_START = KECCAK_F1600_X4_V8A_V84A_SCALAR_PREAMBLE_LENGTH`;;
@@ -1388,7 +1403,7 @@ let KECCAK_F1600_X4_V8A_V84A_SCALAR_SUBROUTINE_CORRECT = prove
     (WORDLIST_FROM_MEMORY_CONV THENC
      ONCE_DEPTH_CONV NORMALIZE_RELATIVE_ADDRESS_CONV) in
   CONV_TAC TWEAK_CONV THEN
-  ARM_ADD_RETURN_STACK_TAC ~pre_post_nsteps:(11,11) KECCAK_F1600_X4_V8A_V84A_SCALAR_EXEC
+  ARM_ADD_RETURN_STACK_TAC ~pre_post_nsteps:(11,26) KECCAK_F1600_X4_V8A_V84A_SCALAR_EXEC
    (CONV_RULE TWEAK_CONV (CONV_RULE LENGTH_SIMPLIFY_CONV KECCAK_F1600_X4_V8A_V84A_SCALAR_CORRECT))
   `[D8; D9; D10; D11; D12; D13; D14; D15;
     X19; X20; X21; X22; X23; X24; X25; X26; X27; X28; X29; X30]` 224);;
diff --git a/proofs/hol_light/x86_64/mlkem/keccak_f1600_x4_avx2_asm.S b/proofs/hol_light/x86_64/mlkem/keccak_f1600_x4_avx2_asm.S
@@ -476,6 +476,12 @@ LLkeccak_f1600_x4_avx2:
         vmovhpd %xmm3, 0x188(%rdi)
         vmovq %xmm15, 0x250(%rdi)
         vmovhpd %xmm15, 0x318(%rdi)
+        vpxor %ymm0, %ymm0, %ymm0
+        .set .Lkeccak_f1600_x4_avx2_clear_stack_offset, 0
+        .rept 24
+        vmovdqu %ymm0, .Lkeccak_f1600_x4_avx2_clear_stack_offset(%rsp)
+        .set .Lkeccak_f1600_x4_avx2_clear_stack_offset, .Lkeccak_f1600_x4_avx2_clear_stack_offset + 32
+        .endr
         movq %r11, %rsp
         .cfi_def_cfa_register %rsp
         retq
diff --git a/proofs/hol_light/x86_64/proofs/keccak_f1600_x4_avx2_asm.ml b/proofs/hol_light/x86_64/proofs/keccak_f1600_x4_avx2_asm.ml
