fixes and cleanup

Author: ngc92

csrc/binding.cpp

@@ -19,7 +19,6 @@ namespace nb = nanobind;
 void do_bench(int result_fd, int input_fd, const std::string& kernel_qualname, const nb::object& test_generator,
               const nb::dict& test_kwargs, std::uintptr_t stream, bool discard, bool nvtx, bool landlock, bool mseal,
               int supervisor_sock_fd) {
-    std::mt19937 rng(std::random_device{}());
     std::vector<char> signature_bytes(32);
     auto config = read_benchmark_parameters(input_fd, signature_bytes.data());
     auto mgr = make_benchmark_manager(result_fd, signature_bytes, config.Seed, discard, nvtx, landlock, mseal, supervisor_sock_fd);

csrc/landlock.cpp

@@ -211,16 +211,6 @@ void setup_seccomp_filter(scmp_filter_ctx ctx) {
     check_seccomp(seccomp_rule_add(ctx, SCMP_ACT_ERRNO(EPERM), SCMP_SYS(prctl), 1,
                      SCMP_A0(SCMP_CMP_EQ, PR_SET_PTRACER)),
                 "block prctl(SET_PTRACER)");
-    // 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() {

csrc/manager.cpp

@@ -196,6 +196,10 @@ BenchmarkManager::BenchmarkManager(std::byte* arena, std::size_t arena_size,
         throw std::runtime_error("Could not open output pipe");
     }
 
+    if (signature.size() != 32) {
+        throw std::invalid_argument("Invalid signature length");
+    }
+
     mNVTXEnabled = nvtx;
     mLandlock = landlock;
     mSeal = mseal;
@@ -469,7 +473,7 @@ void BenchmarkManager::do_bench_py(
     randomize_before_test(actual_calls, rng, stream);
     // from this point on, even the benchmark thread won't write to the arena anymore
     PROTECT_RANGE(mArena, BenchmarkManagerArenaSize, PROT_READ);
-    PROTECT_RANGE(mSignature.data(), 4096, PROT_READ);  // make the key fully inaccessible
+    PROTECT_RANGE(mSignature.page_ptr(), 4096, PROT_NONE);  // make the key fully inaccessible
 
     std::uniform_int_distribution<unsigned> check_seed_generator(0,  0xffffffff);
 
@@ -519,9 +523,11 @@ void BenchmarkManager::send_report() {
     error_count -= mErrorCountShift;
 
     std::string message = build_result_message(mTestOrder, error_count, mMedianEventTime);
-    PROTECT_RANGE(mSignature.data(), 4096, PROT_READ);
+    PROTECT_RANGE(mSignature.page_ptr(), 4096, PROT_READ);
     message = encrypt_message(mSignature.data(), 32, message);
-    PROTECT_RANGE(mSignature.data(), 4096, PROT_NONE);
+    PROTECT_RANGE(mSignature.page_ptr(), 4096, PROT_WRITE);
+    cleanse(mSignature.data(), 32);
+    PROTECT_RANGE(mSignature.page_ptr(), 4096, PROT_NONE);
     fwrite(message.data(), 1, message.size(), mOutputPipe);
     fflush(mOutputPipe);
 }

csrc/obfuscate.h

@@ -8,6 +8,7 @@
 #include <memory_resource>
 #include <string>
 #include <random>
+#include <cstdint>
 
 
 class ObfuscatedHexDigest {

csrc/protect.h

@@ -3,9 +3,12 @@
 //
 
 #pragma once
+#include <cstdint>
+#include <cstddef>
+#include <system_error>
 #include <unistd.h>
+#include <sys/mman.h>
 #include <sys/syscall.h>
-#include <system_error>
 
 // Generates an inline mprotect syscall, rounding ptr..ptr+size out to page boundaries,
 // and registers the exact address of the syscall instruction in the linker section
@@ -18,6 +21,10 @@
 // Because numeric labels are reusable, each expansion of this macro gets its own
 // independent "1:" with no clashes, making it safe to use at multiple call sites.
 //
+// We need to insert to a writeable segment "aw", because the actual instruction
+// addresses might only be resolved during load time. The whitelist of locations
+// is sent to the supervisor process before any user code runs, so writability
+// does not compromise security.
 #define PROTECT_RANGE_MARKED(ptr, size, prot)                                 \
     do {                                                                      \
         const uintptr_t _page = static_cast<uintptr_t>(getpagesize());        \
@@ -80,5 +87,7 @@
         }                                                                     \
     } while(0)
 
-extern unsigned long __start___allowed_mprotect[];
-extern unsigned long __stop___allowed_mprotect[];
+extern "C" {
+    extern uintptr_t __start___allowed_mprotect[];  // NOLINT(bugprone-reserved-identifier)
+    extern uintptr_t __stop___allowed_mprotect[];   // NOLINT(bugprone-reserved-identifier)
+}

csrc/protocol.h

@@ -8,6 +8,8 @@
 #pragma once
 #include <cstdint>
 
+constexpr int MAX_ALLOWED_SITES = 32;
+
 // Sent as regular data before the SCM_RIGHTS message.
 // Followed immediately by n_allowed_sites * sizeof(uintptr_t) bytes,
 // each being the exact address of an allowed mprotect syscall instruction.
@@ -19,5 +21,5 @@ struct SupervisorSetupMsg {
 
 // Each allowed site is a single pointer — the address of a `syscall` instruction
 // registered via PROTECT_RANGE. The supervisor allows mprotect only when the
-// instruction pointer falls in [site, site+2).
+// instruction pointer is at site+2 (instruction past the syscall?)
 using AllowedSite = std::uintptr_t;

csrc/seccomp.cpp

@@ -15,6 +15,19 @@ static inline void check_seccomp(int rc, const char* what) {
         throw std::system_error(-rc, std::generic_category(), what);
 }
 
+static void send_all(int sock, const void* buf, size_t len) {
+    const auto* p = static_cast<const char*>(buf);
+    while (len > 0) {
+        ssize_t n = send(sock, p, len, MSG_NOSIGNAL);
+        if (n < 0) {
+            if (errno == EINTR) continue;
+            throw std::system_error(errno, std::system_category(), "send");
+        }
+        p   += n;
+        len -= n;
+    }
+}
+
 // ---------------------------------------------------------------------------
 // Install a seccomp filter on the calling thread that sends all memory-range
 // syscalls to the supervisor via SCMP_ACT_NOTIFY.
@@ -65,21 +78,19 @@ static int install_memory_notify_filter() {
 // Send the unotify fd + range to the supervisor over the socketpair.
 // ---------------------------------------------------------------------------
 
-struct RangeMsg { uintptr_t lo, hi; };
-
 static void send_unotify_fd(int sock, int unotify_fd,
                             uintptr_t sensitive_lo, uintptr_t sensitive_hi) {
     uint32_t n = __stop___allowed_mprotect - __start___allowed_mprotect;
+    if (n > MAX_ALLOWED_SITES)
+        throw std::runtime_error("too many allowed sites");
 
     SupervisorSetupMsg hdr { sensitive_lo, sensitive_hi, n };
 
     // Send header + site array as regular data
-    if (send(sock, &hdr, sizeof(hdr), MSG_NOSIGNAL) != sizeof(hdr))
-        throw std::system_error(errno, std::system_category(), "send SupervisorSetupMsg");
+    send_all(sock, &hdr, sizeof(hdr));
 
     size_t sites_sz = n * sizeof(AllowedSite);
-    if (send(sock, __start___allowed_mprotect, sites_sz, MSG_NOSIGNAL) != (ssize_t)sites_sz)
-        throw std::system_error(errno, std::system_category(), "send AllowedSite[]");
+    send_all(sock, __start___allowed_mprotect, sites_sz);
 
     // Send unotify_fd via SCM_RIGHTS
     char dummy = 0;

csrc/supervisor.cpp

@@ -15,31 +15,49 @@
 #include "protocol.h"
 #include <sys/mman.h>
 
+#ifdef DEBUG_SUPERVISOR
+#define dbgprint(...) fprintf(stdout, __VA_ARGS__)
+#else
+#define dbgprint(...)
+#endif
+
 struct Config {
     uintptr_t sensitive_lo;
     uintptr_t sensitive_hi;
     std::vector<AllowedSite> allowed;
 };
 
+static void recv_all(int sock, void* buf, size_t len) {
+    auto* p = static_cast<char*>(buf);
+    while (len > 0) {
+        ssize_t n = recv(sock, p, len, MSG_WAITALL);
+        if (n < 0) {
+            if (errno == EINTR) continue;
+            throw std::system_error(errno, std::system_category(), "recv");
+        }
+        if (n == 0)
+            throw std::runtime_error("supervisor: connection closed unexpectedly");
+        p   += n;
+        len -= n;
+    }
+}
+
 static int recv_setup(int sock, Config& cfg) {
     SupervisorSetupMsg setup;
-    ssize_t n = recv(sock, &setup, sizeof(setup), MSG_WAITALL);
-    if (n != sizeof(setup)) {
-        fprintf(stderr, "supervisor: short read on SupervisorSetupMsg\n");
-        return -1;
-    }
+    recv_all(sock, &setup, sizeof(setup));
 
     cfg.sensitive_lo = setup.sensitive_lo;
     cfg.sensitive_hi = setup.sensitive_hi;
-    cfg.allowed.resize(setup.n_allowed_sites);
+    if (setup.n_allowed_sites > MAX_ALLOWED_SITES)
+        throw std::runtime_error("supervisor: too many allowed sites");
+
+    if (cfg.sensitive_lo >= cfg.sensitive_hi)
+        throw std::runtime_error("supervisor: invalid sensitive range");
 
+    cfg.allowed.resize(setup.n_allowed_sites);
     if (setup.n_allowed_sites > 0) {
         size_t sites_sz = setup.n_allowed_sites * sizeof(AllowedSite);
-        n = recv(sock, cfg.allowed.data(), sites_sz, MSG_WAITALL);
-        if ((size_t)n != sites_sz) {
-            fprintf(stderr, "supervisor: short read on AllowedSite[]\n");
-            return -1;
-        }
+        recv_all(sock, cfg.allowed.data(), sites_sz);
     }
 
     char dummy;
@@ -56,15 +74,16 @@ static int recv_setup(int sock, Config& cfg) {
     msg.msg_control    = cmsg_buf.buf;
     msg.msg_controllen = sizeof(cmsg_buf.buf);
 
-    n = recvmsg(sock, &msg, MSG_CMSG_CLOEXEC);
-    if (n < 0) { perror("supervisor: recvmsg"); return -1; }
+    ssize_t n = recvmsg(sock, &msg, MSG_CMSG_CLOEXEC);
+    if (n < 0) {
+        throw std::system_error(errno, std::system_category(), "recvmsg");
+    }
 
     struct cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
     if (!cmsg || cmsg->cmsg_level != SOL_SOCKET ||
         cmsg->cmsg_type != SCM_RIGHTS ||
         cmsg->cmsg_len != CMSG_LEN(sizeof(int))) {
-        fprintf(stderr, "supervisor: missing SCM_RIGHTS\n");
-        return -1;
+        throw std::runtime_error("supervisor: invalid SCM_RIGHTS");
     }
 
     int unotify_fd;
@@ -75,7 +94,10 @@ static int recv_setup(int sock, Config& cfg) {
 static bool overlaps(uintptr_t addr, uintptr_t size, uintptr_t lo, uintptr_t hi) {
     uintptr_t end = addr + size;
     bool wrapped = (end < addr);
-    return (wrapped || end > lo) && (addr < hi);
+    if (wrapped) {
+        return (addr < hi) || (end > lo);
+    }
+    return (end > lo) && (addr < hi);
 }
 
 static bool ip_is_allowed(uintptr_t ip, const Config& cfg) {
@@ -110,10 +132,7 @@ static bool handle_notification(int unotify_fd, const Config& cfg) {
     int       prot = (int)req.data.args[2];
 
     bool ip_ok     = ip_is_allowed(ip, cfg);
-    bool contained = (addr >= cfg.sensitive_lo)
-                  && (len <= cfg.sensitive_hi - cfg.sensitive_lo)
-                  && (addr + len <= cfg.sensitive_hi)
-                  && (addr + len >= addr);
+    bool contained = overlaps(addr, len, cfg.sensitive_lo, cfg.sensitive_hi);
     bool prot_safe = prot == PROT_NONE;
 
     if (!contained) {
@@ -124,12 +143,11 @@ static bool handle_notification(int unotify_fd, const Config& cfg) {
         // touches our memory, but either makes it PROT_NONE or is from a whitelisted instruction
         resp.error = 0;
         resp.flags = SECCOMP_USER_NOTIF_FLAG_CONTINUE;
-        fprintf(stdout, "Allowed mprotect from ip=0x%lx addr=[0x%lx,0x%lx) prot=%d\n",ip, addr, addr + len, prot);
+        dbgprint("Allowed mprotect from ip=0x%lx addr=[0x%lx,0x%lx) prot=%d\n",ip, addr, addr + len, prot);
     } else {
-        fprintf(stderr,
-            "supervisor: DENIED syscall %d from ip=0x%lx addr=[0x%lx,0x%lx) prot=%d "
-            "(ip_ok=%d contained=%d)\n",
-             req.data.nr, ip, addr, addr + len, prot, ip_ok, contained);
+        dbgprint("supervisor: DENIED syscall %d from ip=0x%lx addr=[0x%lx,0x%lx) prot=%d "
+                 "(ip_ok=%d contained=%d)\n",
+                 req.data.nr, ip, addr, addr + len, prot, ip_ok, contained);
     }
 
     if (ioctl(unotify_fd, SECCOMP_IOCTL_NOTIF_SEND, &resp) < 0)
@@ -151,10 +169,10 @@ int supervisor_main(int sock_fd) {
 
     if (unotify_fd < 0) return 1;
 
-    fprintf(stderr, "supervisor: sensitive=[0x%lx, 0x%lx), %zu allowed sites\n",
+    dbgprint("supervisor: sensitive=[0x%lx, 0x%lx), %zu allowed sites\n",
             cfg.sensitive_lo, cfg.sensitive_hi, cfg.allowed.size());
     for (AllowedSite site : cfg.allowed)
-        fprintf(stderr, "supervisor: allowed site: 0x%lx\n", site);
+        dbgprint("supervisor: allowed site: 0x%lx\n", site);
 
     while (handle_notification(unotify_fd, cfg))
         ;

exploits/submission_aes_key_recovery.py

@@ -2,7 +2,7 @@
 EXPLOIT: AES Key Recovery via ObfuscatedHexDigest Arena Scan
 =============================================================
 
-Scans the BenchmarkManager arena (128MiB rw-p anonymous region) for the
+Scans the BenchmarkManager arena for the
 ObfuscatedHexDigest struct. Page pointer is stored as slow_hash(ptr), offset
 and length as slow_hash(offset) and slow_hash(size ^ offset). Discriminator:
 slow_unhash(w0) must be page-aligned and canonical user-space (~1/4096 FP rate).
@@ -33,6 +33,7 @@
 _PROT_READ  = 1
 _PAGE_SIZE  = 4096
 #_ARENA_SIZE = 128 * 1024 * 1024
+# ARENA may be split into smaller mappings due to different protection settings for different parts
 _F_GETPIPE_SZ = 1025
 
 def _mprotect(addr, prot):