Merge pull request #660 from DataDog/zgu/safe-copy

Optimize SafeAccess:safeCopy() implementation in assembly

Co-authored-by: zhengyu123 <zhengyu.gu@datadoghq.com>

Author: gh-worker-dd-mergequeue-cf854d[bot]

ddprof-lib/src/main/cpp/safeAccess.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright 2025 Datadog, Inc
+ * Copyright 2025, 2026 Datadog, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,12 +16,44 @@
 
 
 #include "safeAccess.h"
+#include <cstdio>
+#include <cstdlib>
 #include <signal.h>
 #include <ucontext.h>
 
 extern "C" int safefetch32_cont(int* adr, int errValue);
 extern "C" int64_t safefetch64_cont(int64_t* adr, int64_t errValue);
 
+// safecopy_impl copies `len` bytes from `src` to `dst` one byte at a time and
+// returns 1 on success. If any load from `src` faults, the signal handler
+// redirects the pc anywhere in [safecopy_impl, safecopy_cont) to safecopy_cont,
+// which returns 0. Both are leaf routines with no stack frame, so the redirect
+// keeps the return address intact.
+extern "C" int safecopy_impl(void* dst, const void* src, size_t len);
+extern "C" int safecopy_cont(void* dst, const void* src, size_t len);
+
+#ifdef DEBUG
+// handle_safefetch protects safeCopy by treating any fault whose pc lands in
+// [safecopy_impl, safecopy_cont) as a recoverable read fault. That range is
+// only valid if the assembler/linker keeps safecopy_cont strictly above
+// safecopy_impl (they are emitted adjacently in one asm block). A future
+// toolchain change could in principle reorder them, collapsing the range to
+// empty and silently disabling fault protection. Assert the invariant once at
+// load time so such a regression aborts loudly at startup instead of turning
+// safeCopy into an unprotected crash. Debug builds only: this is a
+// toolchain-regression tripwire, not a runtime safety check.
+__attribute__((constructor))
+static void verify_safecopy_range() {
+  if ((uintptr_t)safecopy_cont <= (uintptr_t)safecopy_impl) {
+    fprintf(stderr,
+            "FATAL: safecopy_cont (%p) is not above safecopy_impl (%p); "
+            "safeCopy fault protection would be lost\n",
+            (void*)safecopy_cont, (void*)safecopy_impl);
+    abort();
+  }
+}
+#endif // DEBUG
+
 #ifdef __APPLE__
     #if defined(__x86_64__)
       #define current_pc context_rip
@@ -67,6 +99,25 @@ extern "C" int64_t safefetch64_cont(int64_t* adr, int64_t errValue);
         _safefetch64_cont:
             movq %rsi, %rax
             ret
+        .globl _safecopy_impl
+        .private_extern _safecopy_impl
+        _safecopy_impl:
+            xorl %ecx, %ecx
+        Lsafecopy_loop:
+            cmpq %rdx, %rcx
+            jae Lsafecopy_done
+            movzbl (%rsi,%rcx,1), %eax
+            movb %al, (%rdi,%rcx,1)
+            incq %rcx
+            jmp Lsafecopy_loop
+        Lsafecopy_done:
+            movl $1, %eax
+            ret
+        .globl _safecopy_cont
+        .private_extern _safecopy_cont
+        _safecopy_cont:
+            xorl %eax, %eax
+            ret
     )");
   #else
     asm(R"(
@@ -95,6 +146,27 @@ extern "C" int64_t safefetch64_cont(int64_t* adr, int64_t errValue);
         safefetch64_cont:
             movq %rsi, %rax
             ret
+        .globl safecopy_impl
+        .hidden safecopy_impl
+        .type safecopy_impl, %function
+        safecopy_impl:
+            xorl %ecx, %ecx
+        .Lsafecopy_loop:
+            cmpq %rdx, %rcx
+            jae .Lsafecopy_done
+            movzbl (%rsi,%rcx,1), %eax
+            movb %al, (%rdi,%rcx,1)
+            incq %rcx
+            jmp .Lsafecopy_loop
+        .Lsafecopy_done:
+            movl $1, %eax
+            ret
+        .globl safecopy_cont
+        .hidden safecopy_cont
+        .type safecopy_cont, %function
+        safecopy_cont:
+            xorl %eax, %eax
+            ret
     )");
   #endif // __APPLE__
 #elif defined(__aarch64__)
@@ -120,6 +192,25 @@ extern "C" int64_t safefetch64_cont(int64_t* adr, int64_t errValue);
         _safefetch64_cont:
             mov      x0, x1
             ret
+        .globl _safecopy_impl
+        .private_extern _safecopy_impl
+        _safecopy_impl:
+            mov      x3, xzr
+        Lsafecopy_loop:
+            cmp      x3, x2
+            b.hs     Lsafecopy_done
+            ldrb     w4, [x1, x3]
+            strb     w4, [x0, x3]
+            add      x3, x3, #1
+            b        Lsafecopy_loop
+        Lsafecopy_done:
+            mov      w0, #1
+            ret
+        .globl _safecopy_cont
+        .private_extern _safecopy_cont
+        _safecopy_cont:
+            mov      w0, #0
+            ret
     )");
   #else
     asm(R"(
@@ -148,61 +239,39 @@ extern "C" int64_t safefetch64_cont(int64_t* adr, int64_t errValue);
         safefetch64_cont:
             mov      x0, x1
             ret
+        .globl safecopy_impl
+        .hidden safecopy_impl
+        .type safecopy_impl, %function
+        safecopy_impl:
+            mov      x3, xzr
+        .Lsafecopy_loop:
+            cmp      x3, x2
+            b.hs     .Lsafecopy_done
+            ldrb     w4, [x1, x3]
+            strb     w4, [x0, x3]
+            add      x3, x3, #1
+            b        .Lsafecopy_loop
+        .Lsafecopy_done:
+            mov      w0, #1
+            ret
+        .globl safecopy_cont
+        .hidden safecopy_cont
+        .type safecopy_cont, %function
+        safecopy_cont:
+            mov      w0, #0
+            ret
     )");
   #endif
 #endif
 
 bool SafeAccess::safeCopy(void* dst, const void* src, size_t len) {
-  // Two-sentinel pattern (same as isReadable): a real-data word may equal
-  // one sentinel by chance, but not both — if both fetches return their
-  // sentinel, the access truly faulted.
-  //
-  // All safefetch32 loads issued here use 4-byte-aligned addresses. Pages
-  // are 4 KiB (or 16 KiB on Apple Silicon), both divisible by 4, so an
-  // aligned 4-byte load never spans a page boundary. The only fault
-  // possible is when the aligned address itself lies in an unmapped page;
-  // we never spuriously fault on an over-read past `src + len`.
-  static const int32_t SENT_A = (int32_t)0x55AA55AA;
-  static const int32_t SENT_B = (int32_t)0xAA55AA55;
-  uint8_t* d = (uint8_t*)dst;
-  const uint8_t* s = (const uint8_t*)src;
-  size_t i = 0;
-
-  // Front fixup: if `src` is not 4-byte aligned, fetch at the previous
-  // aligned address (1..3 bytes before src). That address lies in the
-  // same 4-byte word as src — and since pages are 4-byte aligned, in
-  // the same page as src. The leading k bytes of the fetched word lie
-  // before the caller's range and are discarded via the +k offset; they
-  // never reach `dst`.
-  size_t k = (uintptr_t)s & 3u;
-  if (k != 0 && i < len) {
-    int32_t* aligned = (int32_t*)(s - k);
-    int32_t v1 = safefetch32_impl(aligned, SENT_A);
-    int32_t v2 = safefetch32_impl(aligned, SENT_B);
-    if (v1 == SENT_A && v2 == SENT_B) {
-      return false;
-    }
-    size_t take = (4 - k < len) ? (4 - k) : len;
-    memcpy(d, ((const uint8_t*)&v1) + k, take);
-    i = take;
-  }
-
-  // Middle + tail: (s + i) is now 4-byte aligned. The final iteration may
-  // load up to 3 over-read bytes past `src + len`, but those bytes sit in
-  // the same 4-byte-aligned word and therefore the same page as the bytes
-  // we actually wanted — never a fault from the over-read alone.
-  while (i < len) {
-    int32_t* aligned = (int32_t*)(s + i);
-    int32_t v1 = safefetch32_impl(aligned, SENT_A);
-    int32_t v2 = safefetch32_impl(aligned, SENT_B);
-    if (v1 == SENT_A && v2 == SENT_B) {
-      return false;
-    }
-    size_t chunk = (len - i >= 4) ? 4 : (len - i);
-    memcpy(d + i, &v1, chunk);  // memcpy from local — no UAF risk
-    i += chunk;
-  }
-  return true;
+  // The copy runs entirely inside the safecopy_impl assembly stub, which
+  // reads `src` one byte at a time. If a load faults, handle_safefetch
+  // redirects execution to safecopy_cont, which returns 0. Because the copy
+  // is byte-granular it only ever touches bytes in [src, src+len) — there is
+  // no over-read past the requested range and no alignment reasoning needed.
+  // A fault mid-copy may leave up to len-1 bytes already written to `dst`.
+  return safecopy_impl(dst, src, len) != 0;
 }
 
 bool SafeAccess::handle_safefetch(int sig, void* context) {
@@ -215,6 +284,11 @@ bool SafeAccess::handle_safefetch(int sig, void* context) {
     } else if (pc == (uintptr_t)safefetch64_impl) {
       uc->current_pc = (uintptr_t)safefetch64_cont;
       return true;
+    } else if (pc >= (uintptr_t)safecopy_impl && pc < (uintptr_t)safecopy_cont) {
+      // Unlike safefetch, the faulting load can be at any pc inside the copy
+      // loop, so match the whole [safecopy_impl, safecopy_cont) range.
+      uc->current_pc = (uintptr_t)safecopy_cont;
+      return true;
     }
   }
   return false;

ddprof-lib/src/main/cpp/safeAccess.h

@@ -1,5 +1,6 @@
 /*
  * Copyright 2021 Andrei Pangin
+* Copyright 2026 Datadog, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -64,11 +65,12 @@ class SafeAccess {
     return safefetch64_impl(ptr, errorValue);
   }
 
-  // Copies up to len bytes from src to dst using safefetch32_impl so that a
-  // page-unmap or repurpose of src memory during the copy does not crash the
-  // process. Returns true on full success, false if any read faulted. dst must
-  // have at least len bytes capacity; reads from src may over-read up to 3
-  // bytes past src+len (over-read is also safefetch-protected).
+  // Copies len bytes from src to dst via the safecopy_impl assembly stub so
+  // that a page-unmap or repurpose of src memory during the copy does not
+  // crash the process. Returns true on full success, false if any read
+  // faulted (in which case dst may hold a partial prefix). dst must have at
+  // least len bytes capacity. The copy is byte-granular, so it never reads
+  // past src+len.
   NOINLINE
   static bool safeCopy(void* dst, const void* src, size_t len);
 

ddprof-lib/src/test/cpp/safefetch_ut.cpp

@@ -1,3 +1,8 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <climits>
@@ -162,16 +167,16 @@ TEST_F(SafeFetchTest, mprotectedMemory64) {
 
 // ---------------------------------------------------------------------------
 // SafeAccess::safeCopy — bulk variant of safeFetch{32,64} that copies a byte
-// range via the safefetch trampoline. Must:
+// range via the safecopy_impl assembly stub (byte-granular loads guarded by
+// handle_safefetch). Must:
 //   - return true and copy the bytes exactly when src is fully readable,
 //     including when [src, src+len) sits within a few bytes of an unmapped
-//     page boundary (aligned-load strategy keeps over-reads in-page)
+//     page boundary (byte-granular reads never over-read past src+len)
 //   - return false (no crash) when the requested range itself crosses into
 //     an unmapped page
-//   - handle unaligned src by fetching at the previous 4-byte aligned
-//     address and discarding the leading 1..3 bytes
-//   - never write past dst[len-1] even when len is not a multiple of 4
-//   - not mis-classify real data as a fault when it equals one sentinel
+//   - handle unaligned src (no alignment requirement on src)
+//   - never write past dst[len-1] regardless of len
+//   - copy real data faithfully regardless of its byte values
 // ---------------------------------------------------------------------------
 
 TEST_F(SafeFetchTest, safeCopy_happyPath) {
@@ -188,7 +193,7 @@ TEST_F(SafeFetchTest, safeCopy_zeroLength) {
 }
 
 TEST_F(SafeFetchTest, safeCopy_shortLength_doesNotOverwriteDst) {
-  // The internal 4-byte fetch must not overflow dst beyond len bytes.
+  // The copy must not write beyond len bytes into dst.
   const char src[] = "AB";
   char dst[8];
   memset(dst, 0x5A, sizeof(dst));
@@ -281,6 +286,40 @@ TEST_F(SafeFetchTest, safeCopy_requestedRangeCrossesUnmappedPage_returnsFalse) {
   munmap(region, page_size);
 }
 
+TEST_F(SafeFetchTest, safeCopy_partialPrefixCopiedBeforeFault) {
+  // When the requested range straddles the boundary into an unmapped page,
+  // safeCopy returns false but the byte-granular copy still writes every
+  // readable byte before the fault. Verify the readable prefix landed in dst.
+  long page_size = sysconf(_SC_PAGESIZE);
+  ASSERT_GT(page_size, 0);
+
+  void* region = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
+                      MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(region, MAP_FAILED);
+  ASSERT_EQ(0, munmap((char*)region + page_size, page_size));
+
+  char* mapped_end = (char*)region + page_size;
+  // Place src so exactly `prefix` bytes are readable and the rest fall in
+  // the unmapped page.
+  const size_t prefix = 5;
+  char* src = mapped_end - prefix;
+  static const char kPrefix[] = "HELLO";  // 5 known readable bytes
+  memcpy(src, kPrefix, prefix);
+
+  // Request more than the readable prefix so the copy faults partway.
+  const size_t requested = prefix + 4;
+  char dst[16];
+  memset(dst, 0x5A, sizeof(dst));
+  EXPECT_FALSE(SafeAccess::safeCopy(dst, src, requested));
+
+  // The readable prefix must have been copied faithfully before the fault.
+  EXPECT_EQ(0, memcmp(dst, kPrefix, prefix));
+  // Bytes past the prefix were never written (fault stopped the copy).
+  EXPECT_EQ((char)0x5A, dst[prefix]);
+
+  munmap(region, page_size);
+}
+
 TEST_F(SafeFetchTest, safeCopy_unalignedSource_allMisalignments) {
   // The front fixup must correctly extract leading bytes from the
   // previous-aligned-word fetch for every misalignment k ∈ {1, 2, 3}.
@@ -329,8 +368,8 @@ TEST_F(SafeFetchTest, safeCopy_unalignedShortAtPageEnd_stillSucceeds) {
 }
 
 TEST_F(SafeFetchTest, safeCopy_dataMatchingSingleSentinel_stillSucceeds) {
-  // The two-sentinel pattern must not mis-classify real data that happens
-  // to equal one of the sentinels. SENT_A is 0x55AA55AA.
+  // Real data is copied faithfully regardless of its byte values (the copy
+  // no longer uses sentinel values to detect faults).
   uint32_t real_data = 0x55AA55AA;
   char dst[4];
   ASSERT_TRUE(SafeAccess::safeCopy(dst, &real_data, 4));