add simple seccomp setup

Author: ngc92

.github/workflows/wheel.yml

@@ -39,7 +39,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Install deps
-        run: apt update && apt install -y git g++-13
+        run: apt update && apt install -y git g++-13 libseccomp-dev pkg-config
 
       - name: Install the latest version of uv
         uses: astral-sh/setup-uv@v7

CMakeLists.txt

@@ -7,6 +7,9 @@ set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CUDA_ARCHITECTURES "80;90")
 
+find_package(PkgConfig REQUIRED)
+pkg_check_modules(LIBSECCOMP REQUIRED IMPORTED_TARGET libseccomp)
+
 FetchContent_Declare(
         nanobind
         QUIET
@@ -28,7 +31,7 @@ nanobind_add_module(_pygpubench
         csrc/landlock.cpp
         csrc/obfuscate.cpp
 )
-target_link_libraries(_pygpubench PUBLIC Python::Module CUDA::cudart)
+target_link_libraries(_pygpubench PUBLIC Python::Module CUDA::cudart PkgConfig::LIBSECCOMP)
 # set a bunch of hardening options to make it harder to tamper with the executable
 target_compile_options(_pygpubench PUBLIC -fPIC -pie -fstack-protector-strong -fcf-protection=full -ftrivial-auto-var-init=zero -Wl,-z,relro,-z,now )
 target_compile_definitions(_pygpubench PUBLIC -D_GLIBCXX_ASSERTIONS -D_FORTIFY_SOURCE=3)

csrc/landlock.cpp

@@ -16,6 +16,7 @@
 #include <sys/syscall.h>
 #include <unistd.h>
 #include <linux/landlock.h>
+#include <seccomp.h>
 #include <system_error>
 #include <unordered_set>
 #include <utility>
@@ -114,21 +115,6 @@ void install_landlock() {
     allow_path(ruleset_fd, "/tmp", RW);
     allow_path(ruleset_fd, "/dev", RW); // needed for /dev/null etc, used e.g., by triton
 
-    // Prevent ptrace and /proc/self/mem tampering
-    if (prctl(PR_SET_DUMPABLE, 0) < 0) {
-        throw std::system_error(errno, std::system_category(), "prctl(PR_SET_DUMPABLE)");
-    }
-
-    // Prevent gaining privileges (if attacker tries setuid exploits)
-    if (prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) < 0) {
-        throw std::system_error(errno, std::system_category(), "prctl(PR_SET_NO_NEW_PRIVS)");
-    };
-    // no new executable code pages
-    // note: this also prevents thread creating, which breaks torch.compile
-    // workaround: run torch.compile once from trusted python code, then the thread already
-    //             exists at this point. does not seem reliable, so disabled for now
-    // prctl(PR_SET_MDWE, PR_MDWE_REFUSE_EXEC_GAIN, 0, 0, 0);
-
     landlock_restrict_self(ruleset_fd, 0);
 }
 
@@ -192,4 +178,63 @@ void seal_executable_mappings() {
     for (auto& r : to_seal) {
         mseal(reinterpret_cast<void*>(r.start), r.end - r.start, r.src);
     }
-}
+}
+
+static inline void check_seccomp(int rc, const char* what) {
+    if (rc < 0)
+        throw std::system_error(-rc, std::generic_category(), what);
+}
+
+void setup_seccomp_filter(scmp_filter_ctx ctx) {
+    check_seccomp(seccomp_rule_add(ctx, SCMP_ACT_ERRNO(EPERM), SCMP_SYS(ptrace), 0),
+                      "block ptrace");
+
+    check_seccomp(seccomp_rule_add(ctx, SCMP_ACT_ERRNO(EPERM), SCMP_SYS(prctl), 2,
+                      SCMP_A0(SCMP_CMP_EQ, PR_SET_DUMPABLE),
+                      SCMP_A1(SCMP_CMP_EQ, 1)),
+                  "block prctl(SET_DUMPABLE, 1)");
+
+    check_seccomp(seccomp_rule_add(ctx, SCMP_ACT_ERRNO(EPERM), SCMP_SYS(prctl), 1,
+                      SCMP_A0(SCMP_CMP_EQ, PR_SET_SECCOMP)),
+                  "block prctl(SET_SECCOMP)");
+    // TODO figure out what else we can and should block
+    /*
+    check_seccomp(seccomp_rule_add(ctx, SCMP_ACT_ERRNO(EPERM), SCMP_SYS(mprotect), 1,
+                      SCMP_A2(SCMP_CMP_MASKED_EQ, PROT_WRITE, PROT_WRITE)),
+                  "block mprotect+WRITE");
+
+    check_seccomp(seccomp_rule_add(ctx, SCMP_ACT_ERRNO(EPERM), SCMP_SYS(pkey_mprotect), 1,
+                      SCMP_A2(SCMP_CMP_MASKED_EQ, PROT_WRITE, PROT_WRITE)),
+                  "block pkey_mprotect+WRITE");
+    */
+}
+
+void install_seccomp_filter() {
+    scmp_filter_ctx ctx = seccomp_init(SCMP_ACT_ALLOW);
+    if (!ctx) throw std::runtime_error("seccomp_init failed");
+    try {
+        setup_seccomp_filter(ctx);
+    }  catch (...) {
+        seccomp_release(ctx);
+        throw;
+    }
+
+    // Prevent ptrace and /proc/self/mem tampering
+    if (prctl(PR_SET_DUMPABLE, 0) < 0) {
+        throw std::system_error(errno, std::system_category(), "prctl(PR_SET_DUMPABLE)");
+    }
+
+    // Prevent gaining privileges (if attacker tries setuid exploits)
+    if (prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) < 0) {
+        throw std::system_error(errno, std::system_category(), "prctl(PR_SET_NO_NEW_PRIVS)");
+    };
+    // no new executable code pages
+    // note: this also prevents thread creating, which breaks torch.compile
+    // workaround: run torch.compile once from trusted python code, then the thread already
+    //             exists at this point. does not seem reliable, so disabled for now
+    // prctl(PR_SET_MDWE, PR_MDWE_REFUSE_EXEC_GAIN, 0, 0, 0);
+
+    int rc = seccomp_load(ctx);
+    seccomp_release(ctx);
+    check_seccomp(rc, "seccomp_load");
+}

csrc/manager.cpp

@@ -22,6 +22,7 @@ extern void clear_cache(void* dummy_memory, int size, bool discard, cudaStream_t
 extern void install_landlock();
 extern bool mseal_supported();
 extern void seal_executable_mappings();
+extern void install_seccomp_filter();
 
 static void check_check_approx_match_dispatch(unsigned* result, void* expected_data, nb::dlpack::dtype expected_type,
                                        const nb_cuda_array& received, float r_tol, float a_tol, unsigned seed, std::size_t n_bytes, cudaStream_t stream) {
@@ -293,6 +294,8 @@ void BenchmarkManager::do_bench_py(
         seal_executable_mappings();
     }
 
+    install_seccomp_filter();
+
     // at this point, we call user code as we import the kernel (executing arbitrary top-level code)
     // after this, we cannot trust python anymore
     nb::callable kernel = kernel_from_qualname(kernel_qualname);