Add runtime dispatch (FIPS202)

willieyz · willieyz · commit b6f985146fdc · 2025-11-20T18:33:00.000+08:00
Signed-off-by: willieyz &lt;willie.zhao@chelpis.com&gt;
diff --git a/.github/actions/config-variations/action.yml b/.github/actions/config-variations/action.yml
@@ -7,7 +7,7 @@ inputs:
     description: 'GitHub token'
     required: true
   tests:
-    description: 'List of tests to run (space-separated IDs) or "all" for all tests. Available IDs: pct-enabled, pct-enabled-broken, custom-zeroize, native-cap-ON, native-cap-OFF, native-cap-CPUID_AVX2, no-asm, serial-fips202, custom-randombytes, custom-memcpy, custom-memset, custom-stdlib'
+    description: 'List of tests to run (space-separated IDs) or "all" for all tests. Available IDs: pct-enabled, pct-enabled-broken, custom-zeroize, native-cap-ON, native-cap-OFF, native-cap-ID_AA64PFR1_EL1, native-cap-CPUID_AVX2, no-asm, serial-fips202, custom-randombytes, custom-memcpy, custom-memset, custom-stdlib'
     required: false
     default: 'all'
   opt:
@@ -85,6 +85,19 @@ runs:
         acvp: true
         opt: ${{ inputs.opt }}
         examples: false # Some examples use a custom config themselves
+    - name: "Custom native capability functions (ID_AA64PFR1_EL1 detection)"
+      if: ${{ (inputs.tests == 'all' || contains(inputs.tests, 'native-cap-ID_AA64PFR1_EL1')) && runner.os == 'Linux' && runner.arch == 'ARM64' }}
+      uses: ./.github/actions/multi-functest
+      with:
+        gh_token: ${{ inputs.gh_token }}
+        compile_mode: native
+        cflags: "-std=c11 -march=armv8.4-a+sha3 -D_GNU_SOURCE -DMLD_CONFIG_FILE=\\\\\\\"../../test/custom_native_capability_config_ID_AA64PFR1_EL1.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        ldflags: "-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        func: true
+        kat: true
+        acvp: true
+        opt: ${{ inputs.opt }}
+        examples: false # Some examples use a custom config themselves
     - name: "Custom native capability functions (CPUID AVX2 detection)"
       if: ${{ (inputs.tests == 'all' || contains(inputs.tests, 'native-cap-CPUID_AVX2')) && runner.os == 'Linux' && runner.arch == 'X64' }}
       uses: ./.github/actions/multi-functest
diff --git a/BIBLIOGRAPHY.md b/BIBLIOGRAPHY.md
@@ -43,6 +43,7 @@ source code and documentation.
   - [test/custom_native_capability_config_0.h](test/custom_native_capability_config_0.h)
   - [test/custom_native_capability_config_1.h](test/custom_native_capability_config_1.h)
   - [test/custom_native_capability_config_CPUID_AVX2.h](test/custom_native_capability_config_CPUID_AVX2.h)
+  - [test/custom_native_capability_config_ID_AA64PFR1_EL1.h](test/custom_native_capability_config_ID_AA64PFR1_EL1.h)
   - [test/custom_randombytes_config.h](test/custom_randombytes_config.h)
   - [test/custom_stdlib_config.h](test/custom_stdlib_config.h)
   - [test/custom_zeroize_config.h](test/custom_zeroize_config.h)
@@ -94,6 +95,7 @@ source code and documentation.
   - [test/custom_native_capability_config_0.h](test/custom_native_capability_config_0.h)
   - [test/custom_native_capability_config_1.h](test/custom_native_capability_config_1.h)
   - [test/custom_native_capability_config_CPUID_AVX2.h](test/custom_native_capability_config_CPUID_AVX2.h)
+  - [test/custom_native_capability_config_ID_AA64PFR1_EL1.h](test/custom_native_capability_config_ID_AA64PFR1_EL1.h)
   - [test/custom_randombytes_config.h](test/custom_randombytes_config.h)
   - [test/custom_stdlib_config.h](test/custom_stdlib_config.h)
   - [test/custom_zeroize_config.h](test/custom_zeroize_config.h)
diff --git a/dev/fips202/aarch64/x1_scalar.h b/dev/fips202/aarch64/x1_scalar.h
@@ -13,10 +13,12 @@
 #define MLD_FIPS202_AARCH64_NEED_X1_SCALAR
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x1_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x1_native(uint64_t *state)
 {
   mld_keccak_f1600_x1_scalar_asm(state, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/dev/fips202/aarch64/x1_v84a.h b/dev/fips202/aarch64/x1_v84a.h
@@ -17,10 +17,17 @@
 #define MLD_FIPS202_AARCH64_NEED_X1_V84A
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x1_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x1_native(uint64_t *state)
 {
+  if (!mld_sys_check_capability(MLD_SYS_CAP_SHA3))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
   mld_keccak_f1600_x1_v84a_asm(state, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/dev/fips202/aarch64/x2_v84a.h b/dev/fips202/aarch64/x2_v84a.h
@@ -17,13 +17,20 @@
 #define MLD_FIPS202_AARCH64_NEED_X2_V84A
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
 
 
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)
 {
+  if (!mld_sys_check_capability(MLD_SYS_CAP_SHA3))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
   mld_keccak_f1600_x2_v84a_asm(state + 0 * 25, mld_keccakf1600_round_constants);
   mld_keccak_f1600_x2_v84a_asm(state + 2 * 25, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/dev/fips202/aarch64/x4_v8a_scalar.h b/dev/fips202/aarch64/x4_v8a_scalar.h
@@ -13,11 +13,13 @@
 #define MLD_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)
 {
   mld_keccak_f1600_x4_scalar_v8a_hybrid_asm(state,
                                             mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/dev/fips202/aarch64/x4_v8a_v84a_scalar.h b/dev/fips202/aarch64/x4_v8a_v84a_scalar.h
@@ -17,11 +17,18 @@
 #define MLD_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)
 {
+  if (!mld_sys_check_capability(MLD_SYS_CAP_SHA3))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
   mld_keccak_f1600_x4_scalar_v8a_v84a_hybrid_asm(
       state, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/mldsa/mldsa_native.c b/mldsa/mldsa_native.c
@@ -481,6 +481,8 @@
 #if defined(MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202)
 /* mldsa/src/fips202/native/api.h */
 #undef MLD_FIPS202_NATIVE_API_H
+#undef MLD_NATIVE_FUNC_FALLBACK
+#undef MLD_NATIVE_FUNC_SUCCESS
 /* mldsa/src/fips202/native/auto.h */
 #undef MLD_FIPS202_NATIVE_AUTO_H
 #if defined(MLD_SYS_AARCH64)
diff --git a/mldsa/src/fips202/keccakf1600.c b/mldsa/src/fips202/keccakf1600.c
@@ -113,16 +113,17 @@ void mld_keccakf1600x4_xor_bytes(uint64_t *state, const unsigned char *data0,
 void mld_keccakf1600x4_permute(uint64_t *state)
 {
 #if defined(MLD_USE_FIPS202_X4_NATIVE)
-  mld_keccak_f1600_x4_native(state);
-#else
+  if (mld_keccak_f1600_x4_native(state) == MLD_NATIVE_FUNC_SUCCESS)
+  {
+    return;
+  }
+#endif /* MLD_USE_FIPS202_X4_NATIVE */
   mld_keccakf1600_permute(state + MLD_KECCAK_LANES * 0);
   mld_keccakf1600_permute(state + MLD_KECCAK_LANES * 1);
   mld_keccakf1600_permute(state + MLD_KECCAK_LANES * 2);
   mld_keccakf1600_permute(state + MLD_KECCAK_LANES * 3);
-#endif /* !MLD_USE_FIPS202_X4_NATIVE */
 }
 
-#if !defined(MLD_USE_FIPS202_X1_NATIVE)
 static const uint64_t mld_KeccakF_RoundConstants[MLD_KECCAK_NROUNDS] = {
     (uint64_t)0x0000000000000001ULL, (uint64_t)0x0000000000008082ULL,
     (uint64_t)0x800000000000808aULL, (uint64_t)0x8000000080008000ULL,
@@ -137,7 +138,7 @@ static const uint64_t mld_KeccakF_RoundConstants[MLD_KECCAK_NROUNDS] = {
     (uint64_t)0x8000000080008081ULL, (uint64_t)0x8000000000008080ULL,
     (uint64_t)0x0000000080000001ULL, (uint64_t)0x8000000080008008ULL};
 
-void mld_keccakf1600_permute(uint64_t *state)
+static void mld_keccakf1600_permute_c(uint64_t *state)
 {
   unsigned round;
 
@@ -402,12 +403,17 @@ void mld_keccakf1600_permute(uint64_t *state)
   state[23] = Aso;
   state[24] = Asu;
 }
-#else  /* !MLD_USE_FIPS202_X1_NATIVE */
+
 void mld_keccakf1600_permute(uint64_t *state)
 {
-  mld_keccak_f1600_x1_native(state);
-}
+#if defined(MLD_USE_FIPS202_X1_NATIVE)
+  if (mld_keccak_f1600_x1_native(state) == MLD_NATIVE_FUNC_SUCCESS)
+  {
+    return;
+  }
 #endif /* MLD_USE_FIPS202_X1_NATIVE */
+  mld_keccakf1600_permute_c(state);
+}
 
 #else /* !MLD_CONFIG_MULTILEVEL_NO_SHARED */
 
diff --git a/mldsa/src/fips202/native/aarch64/x1_scalar.h b/mldsa/src/fips202/native/aarch64/x1_scalar.h
@@ -13,10 +13,12 @@
 #define MLD_FIPS202_AARCH64_NEED_X1_SCALAR
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x1_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x1_native(uint64_t *state)
 {
   mld_keccak_f1600_x1_scalar_asm(state, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/mldsa/src/fips202/native/aarch64/x1_v84a.h b/mldsa/src/fips202/native/aarch64/x1_v84a.h
@@ -17,10 +17,17 @@
 #define MLD_FIPS202_AARCH64_NEED_X1_V84A
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x1_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x1_native(uint64_t *state)
 {
+  if (!mld_sys_check_capability(MLD_SYS_CAP_SHA3))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
   mld_keccak_f1600_x1_v84a_asm(state, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/mldsa/src/fips202/native/aarch64/x2_v84a.h b/mldsa/src/fips202/native/aarch64/x2_v84a.h
@@ -17,13 +17,20 @@
 #define MLD_FIPS202_AARCH64_NEED_X2_V84A
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
 
 
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)
 {
+  if (!mld_sys_check_capability(MLD_SYS_CAP_SHA3))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
   mld_keccak_f1600_x2_v84a_asm(state + 0 * 25, mld_keccakf1600_round_constants);
   mld_keccak_f1600_x2_v84a_asm(state + 2 * 25, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/mldsa/src/fips202/native/aarch64/x4_v8a_scalar.h b/mldsa/src/fips202/native/aarch64/x4_v8a_scalar.h
@@ -13,11 +13,13 @@
 #define MLD_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)
 {
   mld_keccak_f1600_x4_scalar_v8a_hybrid_asm(state,
                                             mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/mldsa/src/fips202/native/aarch64/x4_v8a_v84a_scalar.h b/mldsa/src/fips202/native/aarch64/x4_v8a_v84a_scalar.h
@@ -17,11 +17,18 @@
 #define MLD_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/fips202_native_aarch64.h"
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)
 {
+  if (!mld_sys_check_capability(MLD_SYS_CAP_SHA3))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
   mld_keccak_f1600_x4_scalar_v8a_v84a_hybrid_asm(
       state, mld_keccakf1600_round_constants);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/mldsa/src/fips202/native/api.h b/mldsa/src/fips202/native/api.h
@@ -16,6 +16,14 @@
 #include <stdint.h>
 #include "../../cbmc.h"
 
+/* Backends must return MLD_NATIVE_FUNC_SUCCESS upon success. */
+#define MLD_NATIVE_FUNC_SUCCESS (0)
+/* Backends may return MLD_NATIVE_FUNC_FALLBACK to signal to the frontend that
+ * the target/parameters are unsupported; typically, this would be because of
+ * dependencies on CPU features not detected on the host CPU. In this case,
+ * the frontend falls back to the default C implementation. */
+#define MLD_NATIVE_FUNC_FALLBACK (-1)
+
 /*
  * This is the C<->native interface allowing for the drop-in
  * of custom Keccak-F1600 implementations.
@@ -30,10 +38,10 @@
  */
 
 #if defined(MLD_USE_FIPS202_X1_NATIVE)
-static MLD_INLINE void mld_keccak_f1600_x1_native(uint64_t *state);
+static MLD_INLINE int mld_keccak_f1600_x1_native(uint64_t *state);
 #endif
 #if defined(MLD_USE_FIPS202_X4_NATIVE)
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state);
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state);
 #endif
 
 #endif /* !MLD_FIPS202_NATIVE_API_H */
diff --git a/mldsa/src/fips202/native/x86_64/xkcp.h b/mldsa/src/fips202/native/x86_64/xkcp.h
@@ -13,12 +13,18 @@
 
 #if !defined(__ASSEMBLER__)
 #include <stdint.h>
+#include "../api.h"
 #include "src/KeccakP_1600_times4_SIMD256.h"
 
 #define MLD_USE_FIPS202_X4_NATIVE
-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)
+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)
 {
+  if (!mld_sys_check_capability(MLD_SYS_CAP_AVX2))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
   mld_keccakf1600x4_permute24(state);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 #endif /* !__ASSEMBLER__ */
 
diff --git a/test/configs.yml b/test/configs.yml
@@ -246,6 +246,40 @@ configs:
           }
           #endif /* !__ASSEMBLER__ */
 
+  - path: test/custom_native_capability_config_ID_AA64PFR1_EL1.h
+    description: "Test configuration with ARM system register capability detection"
+    defines:
+      MLD_CONFIG_CUSTOM_CAPABILITY_FUNC:
+        content: |
+          #define MLD_CONFIG_CUSTOM_CAPABILITY_FUNC
+          #if !defined(__ASSEMBLER__)
+          #include <stdint.h>
+          #include "../mldsa/src/sys.h"
+
+          #if !defined(MLD_SYS_AARCH64) || !defined(MLD_SYS_LINUX)
+          #error This configuration is only suitable for Linux/AArch64 systems
+          #endif
+
+          static MLD_INLINE int mld_sys_check_capability(mld_sys_cap cap)
+          {
+            if (cap == MLD_SYS_CAP_SHA3)
+            {
+              uint64_t id_aa64pfr1_el1;
+
+              /* Read ID_AA64PFR1_EL1 system register */
+              __asm__ volatile("mrs %0, id_aa64pfr1_el1" : "=r"(id_aa64pfr1_el1));
+
+              /* Extract SHA3 field (bits 35:32) and check if SHA3 is supported */
+              /* SHA3 field: 0b0000 = not implemented, 0b0001 = SHA3 implemented */
+              uint64_t sha3_field = (id_aa64pfr1_el1 >> 32) & 0xF;
+
+              return (sha3_field == 1) ? 1 : 0;
+            }
+
+            /* Default to 0 (conservative) for unknown capabilities */
+            return 0;
+          }
+          #endif /* !__ASSEMBLER__ */
   # Example configs
 
   - path: examples/monolithic_build/config_44.h
diff --git a/test/custom_native_capability_config_ID_AA64PFR1_EL1.h b/test/custom_native_capability_config_ID_AA64PFR1_EL1.h

Original file line number	Diff line number	Diff line change
`@@ -13,10 +13,12 @@`
`13`	`13`	`#define MLD_FIPS202_AARCH64_NEED_X1_SCALAR`
`14`	`14`
`15`	`15`	`#if !defined(__ASSEMBLER__)`
	`16`	`+#include "../api.h"`
`16`	`17`	`#include "src/fips202_native_aarch64.h"`
`17`		`-static MLD_INLINE void mld_keccak_f1600_x1_native(uint64_t *state)`
	`18`	`+static MLD_INLINE int mld_keccak_f1600_x1_native(uint64_t *state)`
`18`	`19`	`{`
`19`	`20`	`mld_keccak_f1600_x1_scalar_asm(state, mld_keccakf1600_round_constants);`
	`21`	`+ return MLD_NATIVE_FUNC_SUCCESS;`
`20`	`22`	`}`
`21`	`23`	`#endif /* !__ASSEMBLER__ */`
`22`	`24`
Original file line number	Diff line number	Diff line change
`@@ -13,11 +13,13 @@`
`13`	`13`	`#define MLD_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID`
`14`	`14`
`15`	`15`	`#if !defined(__ASSEMBLER__)`
	`16`	`+#include "../api.h"`
`16`	`17`	`#include "src/fips202_native_aarch64.h"`
`17`		`-static MLD_INLINE void mld_keccak_f1600_x4_native(uint64_t *state)`
	`18`	`+static MLD_INLINE int mld_keccak_f1600_x4_native(uint64_t *state)`
`18`	`19`	`{`
`19`	`20`	`mld_keccak_f1600_x4_scalar_v8a_hybrid_asm(state,`
`20`	`21`	`mld_keccakf1600_round_constants);`
	`22`	`+ return MLD_NATIVE_FUNC_SUCCESS;`
`21`	`23`	`}`
`22`	`24`	`#endif /* !__ASSEMBLER__ */`
`23`	`25`