Move the remainder of the Rust coverage logic into the noxfile (#8936)

3 files changed, 90 insertions, 20 deletions

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 273d5d808..261ccdb7a 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -59,6 +59,7 @@ jobs:
         uses: actions/setup-python@v4.6.0
         with:
           python-version: ${{ matrix.PYTHON.VERSION }}
+      - run: rustup component add llvm-tools-preview
       - name: Clone wycheproof
         timeout-minutes: 2
         uses: ./.github/actions/wycheproof
@@ -147,6 +148,8 @@ jobs:
           - {IMAGE: "ubuntu-jammy:aarch64", NOXSESSION: "tests", RUNNER: [self-hosted, Linux, ARM64]}
           - {IMAGE: "alpine:aarch64", NOXSESSION: "tests-nocoverage", RUNNER: [self-hosted, Linux, ARM64]}
     timeout-minutes: 15
+    env:
+      RUSTUP_HOME: /root/.rustup
     steps:
       - name: Ridiculous alpine workaround for actions support on arm64
         run: |
@@ -185,7 +188,6 @@ jobs:
       - run: /venv/bin/python -m pip install -c ci-constraints-requirements.txt 'nox'
       - run: '/venv/bin/nox -v --install-only'
         env:
-          RUSTUP_HOME: /root/.rustup
           CARGO_TARGET_DIR: ${{ format('{0}/src/rust/target/', github.workspace) }}
           # OPENSSL_ENABLE_SHA1_SIGNATURES is for CentOS 9 Stream
           OPENSSL_ENABLE_SHA1_SIGNATURES: 1
@@ -294,21 +296,6 @@ jobs:
         env:
           COLUMNS: 80
           CARGO_TARGET_DIR: ${{ format('{0}/src/rust/target/', github.workspace) }}
-      - name: Process coverage data
-        run: |
-            set -xe
-            "$(rustc --print target-libdir)/../bin/llvm-profdata" merge -sparse $(find . -iname "*.profraw") -o rust-cov.profdata
-            COV_UUID=$(python3 -c "import uuid; print(uuid.uuid4())")
-
-            "$(rustc --print target-libdir)/../bin/llvm-cov" export \
-              .nox/tests/lib/python${{ matrix.PYTHON }}/site-packages/cryptography/hazmat/bindings/_rust.abi3.so \
-              $(cat rust-tests.txt | awk '{print "-object " $0}') \
-              -instr-profile=rust-cov.profdata \
-              --ignore-filename-regex='/.cargo/' \
-              --ignore-filename-regex='/rustc/' \
-              --ignore-filename-regex='/.rustup/toolchains/' --format=lcov > "${COV_UUID}.lcov"
-
-            sed -E -i 's/SF:(.*)\/src\/rust\/(.*)/SF:src\/rust\/\2/g' "${COV_UUID}.lcov"
       - uses: ./.github/actions/upload-coverage
 
   macos:
@@ -346,6 +333,7 @@ jobs:
         with:
           python-version: ${{ matrix.PYTHON.VERSION }}
           architecture: 'x64' # we force this right now so that it will install the universal2 on arm64
+      - run: rustup component add llvm-tools-preview
 
       - run: python -m pip install -c ci-constraints-requirements.txt 'nox'
 
@@ -405,6 +393,7 @@ jobs:
         with:
           python-version: ${{ matrix.PYTHON.VERSION }}
           architecture: ${{ matrix.WINDOWS.ARCH }}
+      - run: rustup component add llvm-tools-preview
       - name: Cache rust and pip
         uses: ./.github/actions/cache
         timeout-minutes: 2
diff --git a/.gitignore b/.gitignore
index 7a00ba471..035b15ccd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,4 @@ htmlcov/
 *.py[cdo]
 .hypothesis/
 target/
+.rust-cov/
\ No newline at end of file
diff --git a/noxfile.py b/noxfile.py
index 93b10cd33..86a6a68b6 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -4,7 +4,14 @@
 
 from __future__ import annotations
 
+import glob
+import itertools
 import json
+import pathlib
+import re
+import sys
+import typing
+import uuid
 
 import nox
 
@@ -32,12 +39,15 @@ def tests(session: nox.Session) -> None:
     if session.name == "tests-randomorder":
         extras += ",test-randomorder"
 
+    prof_location = (
+        pathlib.Path(".") / ".rust-cov" / str(uuid.uuid4())
+    ).absolute()
     if session.name != "tests-nocoverage":
         session.env.update(
             {
                 "RUSTFLAGS": "-Cinstrument-coverage "
                 + session.env.get("RUSTFLAGS", ""),
-                "LLVM_PROFILE_FILE": ".rust-cov/cov-%p.profraw",
+                "LLVM_PROFILE_FILE": str(prof_location / "cov-%p.profraw"),
             }
         )
 
@@ -65,6 +75,13 @@ def tests(session: nox.Session) -> None:
         "tests/",
     )
 
+    if session.name != "tests-nocoverage":
+        [rust_so] = glob.glob(
+            f"{session.virtualenv.location}/**/cryptography/hazmat/bindings/_rust.*",
+            recursive=True,
+        )
+        process_rust_coverage(session, [rust_so], prof_location)
+
 
 @nox.session
 def docs(session: nox.Session) -> None:
@@ -149,11 +166,14 @@ def flake(session: nox.Session) -> None:
 
 @nox.session
 def rust(session: nox.Session) -> None:
+    prof_location = (
+        pathlib.Path(".") / ".rust-cov" / str(uuid.uuid4())
+    ).absolute()
     session.env.update(
         {
             "RUSTFLAGS": "-Cinstrument-coverage  "
             + session.env.get("RUSTFLAGS", ""),
-            "LLVM_PROFILE_FILE": ".rust-cov/cov-%p.profraw",
+            "LLVM_PROFILE_FILE": str(prof_location / "cov-%p.profraw"),
         }
     )
 
@@ -187,5 +207,65 @@ def rust(session: nox.Session) -> None:
             if data.get("profile", {}).get("test", False):
                 rust_tests.extend(data["filenames"])
 
-        with open("rust-tests.txt", "w") as f:
-            f.write("\n".join(rust_tests))
+        process_rust_coverage(session, rust_tests, prof_location)
+
+
+LCOV_SOURCEFILE_RE = re.compile(
+    r"^SF:.*[\\/]src[\\/]rust[\\/](.*)$", flags=re.MULTILINE
+)
+BIN_EXT = ".exe" if sys.platform == "win32" else ""
+
+
+def process_rust_coverage(
+    session: nox.Session,
+    rust_binaries: typing.List[str],
+    prof_raw_location: pathlib.Path,
+) -> None:
+    # Hitting weird issues merging Windows and Linux Rust coverage, so just
+    # say the hell with it.
+    if sys.platform == "win32":
+        return
+
+    target_libdir = session.run(
+        "rustc", "--print", "target-libdir", external=True, silent=True
+    )
+    if target_libdir is not None:
+        target_bindir = pathlib.Path(target_libdir).parent / "bin"
+
+        profraws = [
+            str(prof_raw_location / p)
+            for p in prof_raw_location.glob("*.profraw")
+        ]
+        session.run(
+            str(target_bindir / ("llvm-profdata" + BIN_EXT)),
+            "merge",
+            "-sparse",
+            *profraws,
+            "-o",
+            "rust-cov.profdata",
+            external=True,
+        )
+
+        lcov_data = session.run(
+            str(target_bindir / ("llvm-cov" + BIN_EXT)),
+            "export",
+            rust_binaries[0],
+            *itertools.chain.from_iterable(
+                ["-object", b] for b in rust_binaries[1:]
+            ),
+            "-instr-profile=rust-cov.profdata",
+            "--ignore-filename-regex=[/\\].cargo[/\\]",
+            "--ignore-filename-regex=[/\\]rustc[/\\]",
+            "--ignore-filename-regex=[/\\].rustup[/\\]toolchains[/\\]",
+            "--ignore-filename-regex=[/\\]target[/\\]",
+            "--format=lcov",
+            silent=True,
+            external=True,
+        )
+        assert isinstance(lcov_data, str)
+        lcov_data = LCOV_SOURCEFILE_RE.sub(
+            lambda m: "SF:src/rust/" + m.group(1).replace("\\", "/"),
+            lcov_data.replace("\r\n", "\n"),
+        )
+        with open(f"{uuid.uuid4()}.lcov", "w") as f:
+            f.write(lcov_data)