Add AArch64 NEON runtime fallback for ML-KEM and ML-DSA (#3353)

On aarch64, the ML-KEM/ML-DSA native primitives called the NEON assembly
unconditionally, with no runtime path to the C reference implementation.
This was fine since NEON assembly is a mandatory part of the base
aarch64 architecture.

However, FIPS requires us to include a way to test the C fallback
implementation on relevant architectures. This adds a runtime NEON
capability gate on every native entry point so that when NEON is
unavailable (reachable via `OPENSSL_armcap=0`) execution falls back to
the C reference path, mirroring the existing x86-64/AVX2 gate.

This was implemented with the shadow-header pattern already used for
x86-64 ML-DSA (`mldsa_x86_64_meta.h`). Since upstream mlkem-native does
not have this, I've created two similar files `mlkem_aarch64_meta.h` /
`mldsa_aarch64_meta.h` which mirror the vendored aarch64 meta.h. This
allows no edits to the vendored mlkem/ or mldsa/ trees.

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license and the ISC license.

Author: samuel40791765

crypto/fipsmodule/ml_dsa/mldsa_aarch64_meta.h

@@ -0,0 +1,324 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR ISC
+
+// AWS-LC-owned replacement for the vendored mldsa/native/aarch64/meta.h.
+//
+// The vendored header calls the NEON assembly unconditionally. This copy is
+// identical except that every native entry point is guarded by a runtime NEON
+// check (mld_sys_check_capability(MLD_SYS_CAP_NEON)); when NEON is absent it
+// returns MLD_NATIVE_FUNC_FALLBACK and the frontend uses its C reference path
+// (e.g. under OPENSSL_armcap=0). Kept outside the vendored `mldsa/` tree, which
+// importer.sh clobbers on re-import, and selected via mldsa_native_backend.h.
+// Keep the toggles/identifier/asm signatures in sync with the vendored header.
+
+// Reuse the vendored include-guard name on purpose: mldsa_native_bcm.c
+// #undef's it between per-level rounds to re-emit the per-level native
+// functions. A different guard name would break that re-emission.
+#ifndef MLD_NATIVE_AARCH64_META_H
+#define MLD_NATIVE_AARCH64_META_H
+
+/* Set of primitives that this backend replaces */
+#define MLD_USE_NATIVE_NTT
+#define MLD_USE_NATIVE_INTT
+#define MLD_USE_NATIVE_REJ_UNIFORM
+#define MLD_USE_NATIVE_REJ_UNIFORM_ETA2
+#define MLD_USE_NATIVE_REJ_UNIFORM_ETA4
+#define MLD_USE_NATIVE_POLY_DECOMPOSE_32
+#define MLD_USE_NATIVE_POLY_DECOMPOSE_88
+#define MLD_USE_NATIVE_POLY_CADDQ
+#define MLD_USE_NATIVE_POLY_USE_HINT_32
+#define MLD_USE_NATIVE_POLY_USE_HINT_88
+#define MLD_USE_NATIVE_POLY_CHKNORM
+#define MLD_USE_NATIVE_POLYZ_UNPACK_17
+#define MLD_USE_NATIVE_POLYZ_UNPACK_19
+#define MLD_USE_NATIVE_POINTWISE_MONTGOMERY
+#define MLD_USE_NATIVE_POLYVECL_POINTWISE_ACC_MONTGOMERY_L4
+#define MLD_USE_NATIVE_POLYVECL_POINTWISE_ACC_MONTGOMERY_L5
+#define MLD_USE_NATIVE_POLYVECL_POINTWISE_ACC_MONTGOMERY_L7
+
+/* Identifier for this backend so that source and assembly files
+ * in the build can be appropriately guarded. */
+#define MLD_ARITH_BACKEND_AARCH64
+
+
+#if !defined(__ASSEMBLER__)
+#include "mldsa/native/api.h"
+#include "mldsa/native/aarch64/src/arith_native_aarch64.h"
+
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_ntt_native(int32_t data[MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_ntt_aarch64_asm(data, mld_aarch64_ntt_zetas_layer123456,
+                      mld_aarch64_ntt_zetas_layer78);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_intt_native(int32_t data[MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_intt_aarch64_asm(data, mld_aarch64_intt_zetas_layer78,
+                       mld_aarch64_intt_zetas_layer123456);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_rej_uniform_native(int32_t *r, unsigned len,
+                                             const uint8_t *buf,
+                                             unsigned buflen)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON) || len != MLDSA_N ||
+      buflen % 24 != 0)
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
+  /* Safety: outlen is at most MLDSA_N, hence, this cast is safe. */
+  return (int)mld_rej_uniform_aarch64_asm(r, buf, buflen,
+                                          mld_rej_uniform_table);
+}
+
+#if !defined(MLD_CONFIG_NO_KEYPAIR_API)
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLDSA_ETA == 2
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_rej_uniform_eta2_native(int32_t *r, unsigned len,
+                                                  const uint8_t *buf,
+                                                  unsigned buflen)
+{
+  uint64_t outlen;
+  /* AArch64 implementation assumes specific buffer lengths */
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON) || len != MLDSA_N ||
+      buflen != MLD_AARCH64_REJ_UNIFORM_ETA2_BUFLEN)
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  /* Constant time: Inputs and outputs to this function are secret.
+   * It is safe to leak which coefficients are accepted/rejected.
+   * The assembly implementation must not leak any other information about the
+   * accepted coefficients. Constant-time testing cannot cover this, and we
+   * hence have to manually verify the assembly.
+   * We declassify prior the input data and mark the outputs as secret.
+   */
+  MLD_CT_TESTING_DECLASSIFY(buf, buflen);
+  outlen = mld_rej_uniform_eta2_aarch64_asm(r, buf, buflen,
+                                            mld_rej_uniform_eta_table);
+  MLD_CT_TESTING_SECRET(r, sizeof(int32_t) * outlen);
+  /* Safety: outlen is at most MLDSA_N and, hence, this cast is safe. */
+  return (int)outlen;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLDSA_ETA == 2 */
+
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLDSA_ETA == 4
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_rej_uniform_eta4_native(int32_t *r, unsigned len,
+                                                  const uint8_t *buf,
+                                                  unsigned buflen)
+{
+  uint64_t outlen;
+  /* AArch64 implementation assumes specific buffer lengths */
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON) || len != MLDSA_N ||
+      buflen != MLD_AARCH64_REJ_UNIFORM_ETA4_BUFLEN)
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  /* Constant time: Inputs and outputs to this function are secret.
+   * It is safe to leak which coefficients are accepted/rejected.
+   * The assembly implementation must not leak any other information about the
+   * accepted coefficients. Constant-time testing cannot cover this, and we
+   * hence have to manually verify the assembly.
+   * We declassify prior the input data and mark the outputs as secret.
+   */
+  MLD_CT_TESTING_DECLASSIFY(buf, buflen);
+  outlen = mld_rej_uniform_eta4_aarch64_asm(r, buf, buflen,
+                                            mld_rej_uniform_eta_table);
+  MLD_CT_TESTING_SECRET(r, sizeof(int32_t) * outlen);
+  /* Safety: outlen is at most MLDSA_N and, hence, this cast is safe. */
+  return (int)outlen;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLDSA_ETA == 4 */
+#endif /* !MLD_CONFIG_NO_KEYPAIR_API */
+
+#if !defined(MLD_CONFIG_NO_SIGN_API)
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || \
+    (MLD_CONFIG_PARAMETER_SET == 65 || MLD_CONFIG_PARAMETER_SET == 87)
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_poly_decompose_32_native(int32_t *a1, int32_t *a0)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_poly_decompose_32_aarch64_asm(a1, a0);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 65 \
+          || MLD_CONFIG_PARAMETER_SET == 87 */
+
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLD_CONFIG_PARAMETER_SET == 44
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_poly_decompose_88_native(int32_t *a1, int32_t *a0)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_poly_decompose_88_aarch64_asm(a1, a0);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 44 \
+        */
+#endif /* !MLD_CONFIG_NO_SIGN_API */
+
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_poly_caddq_native(int32_t a[MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_poly_caddq_aarch64_asm(a);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+
+#if !defined(MLD_CONFIG_NO_VERIFY_API)
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || \
+    (MLD_CONFIG_PARAMETER_SET == 65 || MLD_CONFIG_PARAMETER_SET == 87)
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_poly_use_hint_32_native(int32_t *a, const int32_t *h)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_poly_use_hint_32_aarch64_asm(a, h);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 65 \
+          || MLD_CONFIG_PARAMETER_SET == 87 */
+
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLD_CONFIG_PARAMETER_SET == 44
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_poly_use_hint_88_native(int32_t *a, const int32_t *h)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_poly_use_hint_88_aarch64_asm(a, h);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 44 \
+        */
+#endif /* !MLD_CONFIG_NO_VERIFY_API */
+
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_poly_chknorm_native(const int32_t *a, int32_t B)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  return mld_poly_chknorm_aarch64_asm(a, B);
+}
+
+#if !defined(MLD_CONFIG_NO_SIGN_API) || !defined(MLD_CONFIG_NO_VERIFY_API)
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLD_CONFIG_PARAMETER_SET == 44
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_polyz_unpack_17_native(int32_t *r, const uint8_t *buf)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_polyz_unpack_17_aarch64_asm(r, buf, mld_polyz_unpack_17_indices);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 44 \
+        */
+
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || \
+    (MLD_CONFIG_PARAMETER_SET == 65 || MLD_CONFIG_PARAMETER_SET == 87)
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_polyz_unpack_19_native(int32_t *r, const uint8_t *buf)
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_polyz_unpack_19_aarch64_asm(r, buf, mld_polyz_unpack_19_indices);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 65 \
+          || MLD_CONFIG_PARAMETER_SET == 87 */
+#endif /* !MLD_CONFIG_NO_SIGN_API || !MLD_CONFIG_NO_VERIFY_API */
+
+#if !defined(MLD_CONFIG_NO_SIGN_API) || !defined(MLD_CONFIG_NO_VERIFY_API) || \
+    defined(MLD_CONFIG_REDUCE_RAM) || defined(MLD_UNIT_TEST)
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_poly_pointwise_montgomery_native(
+    int32_t a[MLDSA_N], const int32_t b[MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_poly_pointwise_montgomery_aarch64_asm(a, b);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* !MLD_CONFIG_NO_SIGN_API || !MLD_CONFIG_NO_VERIFY_API || \
+          MLD_CONFIG_REDUCE_RAM || MLD_UNIT_TEST */
+
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLDSA_L == 4
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_polyvecl_pointwise_acc_montgomery_l4_native(
+    int32_t w[MLDSA_N], const int32_t u[4][MLDSA_N],
+    const int32_t v[4][MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_polyvecl_pointwise_acc_montgomery_l4_aarch64_asm(w, u, v);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLDSA_L == 4 */
+
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLDSA_L == 5
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_polyvecl_pointwise_acc_montgomery_l5_native(
+    int32_t w[MLDSA_N], const int32_t u[5][MLDSA_N],
+    const int32_t v[5][MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_polyvecl_pointwise_acc_montgomery_l5_aarch64_asm(w, u, v);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLDSA_L == 5 */
+
+#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLDSA_L == 7
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_polyvecl_pointwise_acc_montgomery_l7_native(
+    int32_t w[MLDSA_N], const int32_t u[7][MLDSA_N],
+    const int32_t v[7][MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_NEON))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_polyvecl_pointwise_acc_montgomery_l7_aarch64_asm(w, u, v);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLDSA_L == 7 */
+
+#endif /* !__ASSEMBLER__ */
+#endif /* !MLD_NATIVE_AARCH64_META_H */

crypto/fipsmodule/ml_dsa/mldsa_native_backend.h

@@ -10,7 +10,9 @@
     (defined(OPENSSL_LINUX) || defined(OPENSSL_APPLE))
 
 #if defined(OPENSSL_AARCH64)
-#include "mldsa/native/aarch64/meta.h"
+// AWS-LC-owned replacement for mldsa/native/aarch64/meta.h that adds a runtime
+// NEON capability gate; falls back to the C reference impl when NEON is absent.
+#include "mldsa_aarch64_meta.h"
 #elif defined(OPENSSL_X86_64) && !defined(MY_ASSEMBLER_IS_TOO_OLD_FOR_AVX)
 #include "mldsa_x86_64_meta.h"
 #endif

crypto/fipsmodule/ml_dsa/mldsa_native_config.h

@@ -6,6 +6,7 @@
 
 #if !defined(__ASSEMBLER__)
 #include "../../internal.h"
+#include "../cpucap/internal.h"
 #endif
 
 // Namespacing: All symbols are of the form mldsa*. Level-specific
@@ -34,18 +35,29 @@
 #define MLD_CONFIG_KEYGEN_PCT
 #endif
 
-// Map the CPU capability function to the ones used by AWS-LC
+// Map the CPU capability function to the ones used by AWS-LC.
+// MLD_SYS_CAP_NEON is defined here (not in the vendored mldsa/sys.h enum)
+// with a value that does not collide with AVX2 == 0 / SHA3 == 1.
+#if !defined(__ASSEMBLER__)
+#define MLD_SYS_CAP_NEON 2
+#endif
 #define MLD_CONFIG_CUSTOM_CAPABILITY_FUNC
 #if !defined(__ASSEMBLER__)
 #include <stdint.h>
 #include "mldsa/sys.h"
-static MLD_INLINE int mld_sys_check_capability(mld_sys_cap cap)
+static MLD_INLINE int mld_sys_check_capability(int cap)
 {
 #if defined(MLD_SYS_X86_64)
   if (cap == MLD_SYS_CAP_AVX2)
   {
     return CRYPTO_is_AVX2_capable();
   }
+#endif
+#if defined(MLD_SYS_AARCH64)
+  if (cap == MLD_SYS_CAP_NEON)
+  {
+    return CRYPTO_is_NEON_capable();
+  }
 #endif
   return 0;
 }