AArch64: Add ABI checker

This commit adds an ABI checker for all AArch64 assembly in mlkem-native.

The purpose of an ABI checker is to ensure that assembly files obey
the required calling conventions. For AArch64 specifically, it is
required by the AAPCS that x19-x30 and d8-d15 are unmodified.

It is important to test this since failure to obey the AAPCS can lead
to rare but fatal bugs depending on the exact compiler and compilation
settings.

On a high level, the ABI checker works as follows:

1. An assembly callstub

  void asm_call_stub(struct register_state *input,
                     struct register_state *output,
                     void (*function_ptr)(void))

is implemented which calls the function specified by function_ptr
on the input register state input, and returns the output register
state in output. Importantly, it does _not_ assume that function_ptr
obeys the AAPCS, so must be implemented in assembly. This is done for
AArch64 only so far.

2. Randomized ABI checker

Based on asm_call_stub, we implement an ABI checker which calls
the target function with random register inputs and checks that the
output register state has retained the callee saved registers.

3. Pointer inputs

Most functions read/write memory specified via parameters. Calling
those functions with random input would lead to segmentation faults.
The ABI checker therefore does not randomize those arguments, but
instead sets up valid pointers of the right size for them.

For now, we always align the input pointers via MLK_ALIGN. A future
extension should take alignment constraints into account and vary
the alignment of the inputs within those constraints.

The ABI checker gets the information about which inputs are pointers
from the YAML specification of AArch64 assembly.

Some implementation details:

- autogen is extended to generate one test/abicheck/check_FUNCNAME.c
  file per function under test. All tests are declared and enumerated
  in the autogenerated test/abicheck/checks_all.h.
- The handwritten test/abicheck/abicheck.c invokes all per-function
  ABI checks.
- We integrate the abicheck build by extending test/mk/components.mk
  with specific build rules. This allows us to partly piggy back on
  existing Makefile infrastructure, such as setting execution wrappers
  or cross prefixes. On the other hand, a lot of custom hackery is
  currently needed to get past the preprocessor directives guarding
  individual assembly files. This should be improved in the future.
- scripts/tests is extended to allow for tests abicheck, and
  tests all supports --abicheck and --no-abicheck. By default,
  the abicheck _is_ run with tests all.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>

Author: hanno-becker

Makefile

@@ -17,7 +17,7 @@
 	clean quickcheck check-defined-CYCLES \
 	size_512 size_768 size_1024 size \
 	run_size_512 run_size_768 run_size_1024 run_size \
-	host_info
+	host_info abicheck run_abicheck
 
 SHELL := /usr/bin/env bash
 .DEFAULT_GOAL := build
@@ -46,7 +46,7 @@ quickcheck: test
 build: func kat acvp wycheproof
 	$(Q)echo "  Everything builds fine!"
 
-test: run_kat run_func run_acvp run_wycheproof run_unit run_alloc run_rng_fail
+test: run_kat run_func run_acvp run_wycheproof run_unit run_alloc run_rng_fail run_abicheck
 	$(Q)echo "  Everything checks fine!"
 
 # Detect available SHA256 command
@@ -251,6 +251,16 @@ run_size: \
 	run_size_768 \
 	run_size_1024
 
+ifeq ($(OPT),1)
+abicheck: $(ABICHECK_DIR)/bin/abicheck
+
+run_abicheck: abicheck
+	$(W) $(ABICHECK_DIR)/bin/abicheck
+else
+abicheck:
+run_abicheck:
+endif
+
 # Display host and compiler feature detection information
 # Shows which architectural features are supported by both the compiler and host CPU
 # Usage: make host_info [AUTO=0|1] [CROSS_PREFIX=...]

scripts/autogen

@@ -4159,6 +4159,239 @@ def gen_test_configs():
         gen_test_config(config_spec["path"], config_spec, default_config)
 
 
+def extract_yaml_from_assembly(assembly_file):
+    """Extract YAML metadata from assembly file."""
+    with open(assembly_file, "r") as f:
+        content = f.read()
+
+    yaml_match = re.search(r"/\*yaml\s*\n(.*?)\n\*/", content, re.DOTALL)
+    if not yaml_match:
+        raise ValueError(f"No YAML metadata found in {assembly_file}")
+
+    return yaml.safe_load(yaml_match.group(1))
+
+
+def extract_arch_flags_from_assembly(assembly_file):
+    """Extract architecture-specific preprocessor flags from assembly file."""
+    with open(assembly_file, "r") as f:
+        content = f.read()
+
+    arch_flags = []
+    if re.search(r"#if\s+defined\(__ARM_FEATURE_SHA3\)", content):
+        arch_flags.append("__ARM_FEATURE_SHA3")
+    return arch_flags
+
+
+def resolve_buffer_size(reg_info, abi_data):
+    """Resolve buffer size from register info, handling cross-references."""
+    size = reg_info.get("size_bytes")
+    if isinstance(size, str):
+        if size in abi_data:
+            ref_reg = abi_data[size]
+            if ref_reg.get("type") == "scalar" and "test_with" in ref_reg:
+                return ref_reg["test_with"]
+            raise ValueError(f"Buffer {reg_info} references non-scalar register {size}")
+        return int(size)
+    elif isinstance(size, int):
+        return size
+    raise ValueError(f"Cannot resolve buffer size from {reg_info}")
+
+
+def gen_abicheck():
+    """Generate ABI checker tests for all aarch64 assembly functions."""
+
+    base_path = pathlib.Path(".")
+    patterns = [
+        "mlkem/src/fips202/native/aarch64/src/*.S",
+        "mlkem/src/native/aarch64/src/*.S",
+    ]
+
+    aarch64_asm_files = []
+    for pattern in patterns:
+        aarch64_asm_files.extend(base_path.glob(pattern))
+
+    if not aarch64_asm_files:
+        return
+
+    aarch64_asm_files = sorted(str(f) for f in aarch64_asm_files)
+    generated_functions = []
+
+    for assembly_file in aarch64_asm_files:
+        yaml_data = extract_yaml_from_assembly(assembly_file)
+        arch_flags = extract_arch_flags_from_assembly(assembly_file)
+
+        function_name = yaml_data.get("Name")
+        c_function_name = "mlk_" + function_name
+
+        def gen_c_test(
+            function_name=function_name,
+            c_function_name=c_function_name,
+            yaml_data=yaml_data,
+            arch_flags=arch_flags,
+        ):
+            yield from gen_header()
+            yield '#include "../../mlkem/src/sys.h"'
+            yield ""
+            yield "#if defined(MLK_SYS_AARCH64)"
+            yield ""
+
+            for flag in arch_flags:
+                yield f"#if defined({flag})"
+                yield ""
+
+            yield "#include <stdio.h>"
+            yield "#include <string.h>"
+            yield ""
+            yield '#include "../notrandombytes/notrandombytes.h"'
+            yield '#include "abicheckutil.h"'
+            yield '#include "checks_all.h"'
+            yield ""
+            yield "typedef struct register_state register_state;"
+            yield ""
+            yield "#define NUM_TESTS 3"
+            yield ""
+            yield yaml_data.get("Signature") + ";"
+            yield ""
+            yield f"int check_{function_name}(void)"
+            yield "{"
+
+            yield "  int test_iter;"
+            yield "  register_state input_state, output_state;"
+            yield "  int violations;"
+
+            abi_data = yaml_data.get("ABI", {})
+            sorted_registers = sorted(abi_data.items())
+
+            buffer_info = []
+            for reg_name, reg_info in sorted_registers:
+                if reg_info.get("type") == "buffer":
+                    size_bytes = resolve_buffer_size(reg_info, abi_data)
+                    buffer_name = f"buf_{reg_name}"
+                    description = reg_info.get("description", "")
+                    yield f"  MLK_ALIGN uint8_t {buffer_name}[{size_bytes}]; /* {description} */"
+                    buffer_info.append((reg_name, buffer_name, size_bytes))
+
+            yield ""
+            yield "  for (test_iter = 0; test_iter < NUM_TESTS; test_iter++)"
+            yield "  {"
+            yield "    /* Initialize random register state */"
+            yield "    init_random_register_state(&input_state);"
+            yield ""
+
+            for reg_name, buffer_name, size_bytes in buffer_info:
+                yield f"    /* Initialize buffer for {reg_name} */"
+                yield f"    randombytes({buffer_name}, {size_bytes});"
+
+            yield ""
+            yield "    /* Set up register state for function arguments */"
+
+            for reg_name, reg_info in sorted_registers:
+                reg_num = int(reg_name[1:])
+                if reg_info.get("type") == "buffer":
+                    yield f"    input_state.gpr[{reg_num}] = (uint64_t){f'buf_{reg_name}'};"
+                elif reg_info.get("type") == "scalar":
+                    test_with = reg_info.get("test_with")
+                    if test_with:
+                        yield f"    input_state.gpr[{reg_num}] = {test_with};"
+
+            yield ""
+            yield "    /* Call function through ABI test stub */"
+            yield f"    asm_call_stub(&input_state, &output_state, (void (*)(void)){c_function_name});"
+            yield ""
+            yield "    /* Check ABI compliance */"
+            yield "    violations = check_aarch64_aapcs_compliance(&input_state, &output_state);"
+            yield "    if (violations > 0) {"
+            yield f'      fprintf(stderr, "ABI test FAILED for {function_name} (iteration %d): %d violations\\n",'
+            yield "             test_iter + 1, violations);"
+            yield "      return 1;"
+            yield "    }"
+            yield "  }"
+            yield ""
+            yield "  return 0;"
+            yield "}"
+            yield ""
+
+            for flag in reversed(arch_flags):
+                yield f"#else /* !{flag} */"
+                yield ""
+                yield '#include "../../mlkem/src/common.h"'
+                yield f"MLK_EMPTY_CU(check_{function_name})"
+                yield ""
+                yield f"#endif /* {flag} */"
+                yield ""
+
+            yield "#else /* !MLK_SYS_AARCH64 */"
+            yield ""
+            yield '#include "../../mlkem/src/common.h"'
+            yield f"MLK_EMPTY_CU(check_{function_name})"
+            yield ""
+            yield "#endif /* MLK_SYS_AARCH64 */"
+            yield ""
+
+        output_file = f"test/abicheck/check_{function_name}.c"
+        update_file(output_file, "\n".join(gen_c_test()), force_format=True)
+        generated_functions.append((function_name, arch_flags))
+
+    # Generate checks_all.h
+    def gen_checks_all_header():
+        yield from gen_header()
+        yield ""
+        yield "#ifndef CHECKS_ALL_H"
+        yield "#define CHECKS_ALL_H"
+        yield ""
+        yield "#include <stddef.h>"
+        yield '#include "../../mlkem/src/sys.h"'
+        yield ""
+        yield "/* Array of all check functions */"
+        yield "typedef struct"
+        yield "{"
+        yield "  const char *name;"
+        yield "  int (*check_func)(void);"
+        yield "} abicheck_entry_t;"
+        yield ""
+        yield "#if defined(MLK_SYS_AARCH64)"
+        yield ""
+        yield "/* Function prototypes for all ABI checks (only on AArch64) */"
+
+        for func_name, arch_flags in generated_functions:
+            for flag in arch_flags:
+                yield f"#if defined({flag})"
+            yield f"int check_{func_name}(void);"
+            for flag in reversed(arch_flags):
+                yield f"#endif"
+
+        yield ""
+        yield "static const abicheck_entry_t all_checks[] = {"
+
+        for func_name, arch_flags in generated_functions:
+            for flag in arch_flags:
+                yield f"#if defined({flag})"
+            yield f'    {{"{func_name}", check_{func_name}}},'
+            for flag in reversed(arch_flags):
+                yield f"#endif"
+
+        yield "    {NULL, NULL} /* Sentinel */"
+        yield "};"
+        yield ""
+        yield "#else /* !MLK_SYS_AARCH64 */"
+        yield ""
+        yield "/* Empty array for non-AArch64 platforms */"
+        yield "static const abicheck_entry_t all_checks[] = {"
+        yield "    {NULL, NULL} /* Sentinel */"
+        yield "};"
+        yield ""
+        yield "#endif /* MLK_SYS_AARCH64 */"
+        yield ""
+        yield "#endif /* CHECKS_ALL_H */"
+        yield ""
+
+    update_file(
+        "test/abicheck/checks_all.h",
+        "\n".join(gen_checks_all_header()),
+        force_format=True,
+    )
+
+
 def _main():
     slothy_choices = [
         "ntt",
@@ -4277,6 +4510,7 @@ def _main():
         ("Generate undefs", gen_undefs),
         ("Generate test configs", gen_test_configs),
         ("Check macro typos", check_macro_typos),
+        ("Generate ABI checker tests", gen_abicheck),
         ("Generate preprocessor comments", gen_preprocessor_comments),
         # Formatting should be the last step
         ("Format files", lambda: format_files(args.dry_run)),

scripts/tests

@@ -214,6 +214,7 @@ class TEST_TYPES(Enum):
     ALLOC = 20
     RNG_FAIL = 21
     WYCHEPROOF = 22
+    ABICHECK = 23
 
     def is_benchmark(self):
         return self in [TEST_TYPES.BENCH, TEST_TYPES.BENCH_COMPONENTS]
@@ -294,6 +295,8 @@ class TEST_TYPES(Enum):
             return "Alloc Test"
         if self == TEST_TYPES.RNG_FAIL:
             return "RNG Failure Test"
+        if self == TEST_TYPES.ABICHECK:
+            return "ABI Compliance Test"
 
     def make_dir(self):
         if self == TEST_TYPES.BRING_YOUR_OWN_FIPS202:
@@ -365,6 +368,8 @@ class TEST_TYPES(Enum):
             return "alloc"
         if self == TEST_TYPES.RNG_FAIL:
             return "rng_fail"
+        if self == TEST_TYPES.ABICHECK:
+            return "abicheck"
 
     def make_run_target(self, scheme):
         t = self.make_target()
@@ -830,6 +835,17 @@ class Tests:
         if resultss is None:
             self.check_fail()
 
+    def abicheck(self):
+        """Run ABI compliance tests for assembly functions."""
+        if not self.do_opt():
+            return
+
+        self._compile_schemes(TEST_TYPES.ABICHECK, True)
+        if self.args.run:
+            self._run_scheme(TEST_TYPES.ABICHECK, True, None)
+
+        self.check_fail()
+
     def all(self):
         func = self.args.func
         kat = self.args.kat
@@ -840,6 +856,7 @@ class Tests:
         unit = self.args.unit
         alloc = self.args.alloc
         rng_fail = self.args.rng_fail
+        abicheck = self.args.abicheck
 
         def _all(opt):
             if func is True:
@@ -858,6 +875,8 @@ class Tests:
                 self._compile_schemes(TEST_TYPES.ALLOC, opt)
             if rng_fail is True:
                 self._compile_schemes(TEST_TYPES.RNG_FAIL, opt)
+            if abicheck is True and opt:
+                self._compile_schemes(TEST_TYPES.ABICHECK, opt)
 
             if self.args.check_namespace is True:
                 p = subprocess.run(
@@ -887,6 +906,8 @@ class Tests:
                 self._run_schemes(TEST_TYPES.ALLOC, opt)
             if rng_fail is True:
                 self._run_schemes(TEST_TYPES.RNG_FAIL, opt)
+            if abicheck is True and opt:
+                self._run_scheme(TEST_TYPES.ABICHECK, opt, None)
 
         if self.do_no_opt():
             _all(False)
@@ -1308,6 +1329,21 @@ def cli():
         help="Do not run RNG failure tests",
     )
 
+    abicheck_group = all_parser.add_mutually_exclusive_group()
+    abicheck_group.add_argument(
+        "--abicheck",
+        action="store_true",
+        dest="abicheck",
+        help="Run ABI compliance tests",
+        default=True,
+    )
+    abicheck_group.add_argument(
+        "--no-abicheck",
+        action="store_false",
+        dest="abicheck",
+        help="Do not run ABI compliance tests",
+    )
+
     # acvp arguments
     acvp_parser = cmd_subparsers.add_parser(
         "acvp", help="Run ACVP client", parents=[common_parser]
@@ -1385,6 +1421,13 @@ def cli():
         parents=[common_parser],
     )
 
+    # abicheck arguments
+    cmd_subparsers.add_parser(
+        "abicheck",
+        help="Run ABI compliance tests for assembly functions",
+        parents=[common_parser],
+    )
+
     # cbmc arguments
     cbmc_parser = cmd_subparsers.add_parser(
         "cbmc",
@@ -1576,6 +1619,8 @@ def cli():
         Tests(args).alloc()
     elif args.cmd == "rng_fail":
         Tests(args).rng_fail()
+    elif args.cmd == "abicheck":
+        Tests(args).abicheck()
 
 
 if __name__ == "__main__":