and break the exploit again

Author: ngc92

csrc/binding.cpp

@@ -11,8 +11,8 @@
 namespace nb = nanobind;
 
 
-void do_bench(int result_fd, const std::string& kernel_qualname, const nb::object& test_generator, const nb::dict& test_kwargs, int repeats, std::uint64_t seed, std::uintptr_t stream, bool discard, bool nvtx) {
-    BenchmarkManager mgr(result_fd, seed, discard, nvtx);
+void do_bench(int result_fd, int signature_fd, const std::string& kernel_qualname, const nb::object& test_generator, const nb::dict& test_kwargs, int repeats, std::uint64_t seed, std::uintptr_t stream, bool discard, bool nvtx) {
+    BenchmarkManager mgr(result_fd, signature_fd, seed, discard, nvtx);
     auto [args, expected] = mgr.setup_benchmark(nb::cast<nb::callable>(test_generator), test_kwargs, repeats);
     mgr.do_bench_py(kernel_qualname, args, expected, reinterpret_cast<cudaStream_t>(stream));
 }

csrc/manager.cpp

@@ -52,7 +52,7 @@ static nb::callable kernel_from_qualname(const std::string& qualname) {
     return nb::cast<nb::callable>(mod.attr(attr.c_str()));
 }
 
-BenchmarkManager::BenchmarkManager(int result_fd, std::uint64_t seed, bool discard, bool nvtx) {
+BenchmarkManager::BenchmarkManager(int result_fd, int signature_fd, std::uint64_t seed, bool discard, bool nvtx) {
     int device;
     CUDA_CHECK(cudaGetDevice(&device));
     CUDA_CHECK(cudaDeviceGetAttribute(&mL2CacheSize, cudaDevAttrL2CacheSize, device));
@@ -63,6 +63,11 @@ BenchmarkManager::BenchmarkManager(int result_fd, std::uint64_t seed, bool disca
     mNVTXEnabled = nvtx;
     mDiscardCache = discard;
     mSeed = seed;
+    char sig_buf[256];
+    FILE* sig_file = fdopen(signature_fd, "r");
+    fgets(sig_buf, sizeof(sig_buf), sig_file);
+    fclose(sig_file);
+    mSignature = std::string(sig_buf);
 }
 
 BenchmarkManager::~BenchmarkManager() {
@@ -371,6 +376,7 @@ void BenchmarkManager::do_bench_py(const std::string& kernel_qualname, const std
         CUDA_CHECK(cudaEventElapsedTime(&duration, mStartEvents.at(i), mEndEvents.at(i)));
         fprintf(mOutputFile, "%d\t%f\n", test_order.at(i) - 1, duration * 1000);
     }
+    fprintf(mOutputFile, "signature\t%s", mSignature.c_str());
     fflush(mOutputFile);
 
     // cleanup events

csrc/manager.h

@@ -19,7 +19,7 @@ using nb_cuda_array = nb::ndarray<nb::c_contig, nb::device::cuda>;
 
 class BenchmarkManager {
 public:
-    BenchmarkManager(int result_fd, std::uint64_t seed, bool discard, bool nvtx);
+    BenchmarkManager(int result_fd, int signature_fd, std::uint64_t seed, bool discard, bool nvtx);
     ~BenchmarkManager();
     std::pair<std::vector<nb::tuple>, std::vector<nb::tuple>> setup_benchmark(const nb::callable& generate_test_case, const nb::dict& kwargs, int repeats);
     void do_bench_py(const std::string& kernel_qualname, const std::vector<nb::tuple>& args, const std::vector<nb::tuple>& expected, cudaStream_t stream);
@@ -67,6 +67,7 @@ class BenchmarkManager {
     std::vector<Expected> mExpectedOutputs;
 
     FILE* mOutputFile;
+    std::string mSignature;
 
     static ShadowArgumentList make_shadow_args(const nb::tuple& args, cudaStream_t stream);
 

python/pygpubench/__init__.py

@@ -4,6 +4,7 @@
 import multiprocessing as mp
 import os
 import traceback
+import secrets
 
 from typing import Optional
 
@@ -24,12 +25,13 @@
 ]
 
 
-def do_bench_impl(out_fd: "multiprocessing.Pipe", qualname: str, test_generator: TestGeneratorInterface,
+def do_bench_impl(out_fd: "multiprocessing.Pipe", signature: "multiprocessing.Pipe", qualname: str, test_generator: TestGeneratorInterface,
                   test_args: dict, repeats: int, seed: int, stream: int = None, discard: bool = True,
-                  nvtx: bool = False, tb_conn=None):
+                  nvtx: bool = False, tb_conn: "multiprocessing.Pipe" = None):
     """
     Benchmarks the kernel referred to by `qualname` against the test case returned by `test_generator`.
     :param out_fd: Writable file descriptor to which benchmark results are written.
+    :param signature: Authentication token read by the C++ layer before untrusted code runs.
     :param qualname: Fully qualified name of the kernel object, e.g. ``my_package.my_module.kernel``.
     :param test_generator: A function that takes the test arguments (including a seed) and returns a test case; i.e., a tuple of (input, expected)
     :param test_args: keyword arguments to be passed to `test_generator`. Seed will be generated automatically.
@@ -48,6 +50,7 @@ def do_bench_impl(out_fd: "multiprocessing.Pipe", qualname: str, test_generator:
         with DeterministicContext():
             _pygpubench.do_bench(
                 out_fd.fileno(),
+                signature.fileno(),
                 qualname,
                 test_generator,
                 test_args,
@@ -141,6 +144,11 @@ def do_bench_isolated(
     read_fd = result_parent.fileno()
     write_fd = result_child.fileno()
 
+    sig_r, sig_w = ctx.Pipe(duplex=False)
+    signature = secrets.token_hex(16)
+    os.write(sig_w.fileno(), signature.encode())
+    sig_w.close()
+
     try:
         import fcntl
         # F_SETPIPE_SZ is Linux-specific (1032); fall back silently on other OSes.
@@ -159,6 +167,7 @@ def do_bench_isolated(
         target=do_bench_impl,
         args=(
             result_child,
+            sig_r,
             qualname,
             test_generator,
             test_args,
@@ -204,6 +213,7 @@ def do_bench_isolated(
     parent_tb_conn.close()
 
     results = BenchmarkResult(None, [-1] * repeats, None, False)
+    has_signature = False
     for line in raw.decode().splitlines():
         parts = line.strip().split('\t')
         if len(parts) == 2 and parts[0].isdigit():
@@ -214,5 +224,12 @@ def do_bench_isolated(
             results.event_overhead_us = float(parts[1].split()[0])
         elif parts[0] == "error-count":
             results.errors = int(parts[1])
+        elif parts[0] == "signature":
+            if signature != parts[1]:
+                raise AssertionError(f"Invalid signature")
+            has_signature = True
+    if not has_signature:
+        raise RuntimeError(f"No signature found in output")
+
     results.full = all((t > 0 for t in results.time_us))
     return results