From f88eeb67cebeec7367c9339b576249517edafca8 Mon Sep 17 00:00:00 2001
From: "Matthias J. Kannwischer" <matthias@kannwischer.eu>
Date: Tue, 3 Dec 2024 15:19:33 +0800
Subject: [PATCH] CBMC: add `tests cbmc` command

Currently to run all CBMC proofs, we have to run

cd cbmc/proofs
MLKEM_K=2 python3 run-cbmc-proofs.py
MLKEM_K=3 python3 run-cbmc-proofs.py
MLKEM_K=4 python3 run-cbmc-proofs.py

This commit wraps all that into a single wrapper:
tests cbmc

It is also added to CI.

Signed-off-by: Matthias J. Kannwischer <matthias@kannwischer.eu>
---
 .github/actions/cbmc/action.yml |  3 +--
 scripts/lib/mlkem_test.py       | 20 ++++++++++++++++++++
 scripts/tests                   | 23 +++++++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/.github/actions/cbmc/action.yml b/.github/actions/cbmc/action.yml
index b76f5138b..193c9fa61 100644
--- a/.github/actions/cbmc/action.yml
+++ b/.github/actions/cbmc/action.yml
@@ -46,7 +46,6 @@ runs:
       - name: Run CBMC proofs (MLKEM_K=${{ inputs.mlkem_k }})
         shell: ${{ env.SHELL }}
         run: |
-          cd cbmc/proofs;
           echo "::group::cbmc_${{ inputs.mlkem_k }}"
-          MLKEM_K=${{ inputs.mlkem_k }} ./run-cbmc-proofs.py --summarize --no-coverage -j8;
+          tests cbmc --k ${{ inputs.mlkem_k }};
           echo "::endgroup::"
diff --git a/scripts/lib/mlkem_test.py b/scripts/lib/mlkem_test.py
index 4853aac69..d935e0b1c 100644
--- a/scripts/lib/mlkem_test.py
+++ b/scripts/lib/mlkem_test.py
@@ -50,6 +50,7 @@ def __init__(self):
         self.run = True
         self.exec_wrapper = ""
         self.run_as_root = ""
+        self.k = "ALL"
 
 
 class Base:
@@ -709,3 +710,22 @@ def all(opt: bool):
             exit_code = exit_code or all(True)
 
         exit(exit_code)
+
+    def cbmc(self, k):
+        config_logger(self.verbose)
+        def run_cbmc(mlkem_k):
+            envvars = {"MLKEM_K": mlkem_k}
+            cpucount = os.cpu_count()
+            p = subprocess.Popen(
+                ["python3", "run-cbmc-proofs.py", "--summarize",  "--no-coverage", f"-j{cpucount}"],
+                cwd="cbmc/proofs",
+                env=os.environ.copy() | envvars,
+            )
+            p.communicate()
+            assert p.returncode == 0
+        if k == "ALL":
+            run_cbmc("2")
+            run_cbmc("3")
+            run_cbmc("4")
+        else: 
+            run_cbmc(k)
diff --git a/scripts/tests b/scripts/tests
index 5c8708053..e559663e4 100755
--- a/scripts/tests
+++ b/scripts/tests
@@ -291,5 +291,28 @@ def all(
     Tests(opts).all(func, kat, nistkat, acvp)
 
 
+@cli.command(
+    short_help="Run the CBMC proofs for all parameter sets",
+    context_settings={"show_default": True},
+)
+@click.make_pass_decorator(Options, ensure=True)
+@add_options(
+    [
+        click.option(
+            "--k",
+            expose_value=False,
+            nargs=1,
+            type=click.Choice(["2", "3", "4", "ALL"]),
+            show_default=True,
+            default="ALL",
+            help="MLKEM parameter set (MLKEM_K).",
+            callback=__callback("k"),
+        )
+    ]
+)
+def cbmc(opts: Options):
+    Tests(opts).cbmc(opts.k)
+
+
 if __name__ == "__main__":
     cli()