From ec57fa1452bc6e8191dea33d43fe2b4253bf6335 Mon Sep 17 00:00:00 2001 From: Zhengyu Gu Date: Tue, 14 Jul 2026 15:49:06 +0000 Subject: [PATCH 1/2] Optimize SafeAccess:safeCopy() implementation in assembly --- ddprof-lib/src/main/cpp/safeAccess.cpp | 150 +++++++++++++++-------- ddprof-lib/src/main/cpp/safeAccess.h | 11 +- ddprof-lib/src/test/cpp/safefetch_ut.cpp | 18 +-- 3 files changed, 115 insertions(+), 64 deletions(-) diff --git a/ddprof-lib/src/main/cpp/safeAccess.cpp b/ddprof-lib/src/main/cpp/safeAccess.cpp index ce650f3c2a..8f4576530b 100644 --- a/ddprof-lib/src/main/cpp/safeAccess.cpp +++ b/ddprof-lib/src/main/cpp/safeAccess.cpp @@ -22,6 +22,14 @@ 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 __APPLE__ #if defined(__x86_64__) #define current_pc context_rip @@ -67,6 +75,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 +122,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 +168,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 +215,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 +260,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; diff --git a/ddprof-lib/src/main/cpp/safeAccess.h b/ddprof-lib/src/main/cpp/safeAccess.h index 43ed9ce3ae..151e6ee75d 100644 --- a/ddprof-lib/src/main/cpp/safeAccess.h +++ b/ddprof-lib/src/main/cpp/safeAccess.h @@ -64,11 +64,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); diff --git a/ddprof-lib/src/test/cpp/safefetch_ut.cpp b/ddprof-lib/src/test/cpp/safefetch_ut.cpp index f1cf4ae7f0..a5bb9941f5 100644 --- a/ddprof-lib/src/test/cpp/safefetch_ut.cpp +++ b/ddprof-lib/src/test/cpp/safefetch_ut.cpp @@ -162,16 +162,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 +188,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)); @@ -329,8 +329,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)); From bef3bc45ab4d4ec4348129c694b5e1f9d37bc521 Mon Sep 17 00:00:00 2001 From: Zhengyu Gu Date: Tue, 14 Jul 2026 19:28:29 +0000 Subject: [PATCH 2/2] Review comments --- ddprof-lib/src/main/cpp/safeAccess.cpp | 26 +++++++++++++++- ddprof-lib/src/main/cpp/safeAccess.h | 1 + ddprof-lib/src/test/cpp/safefetch_ut.cpp | 39 ++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 1 deletion(-) diff --git a/ddprof-lib/src/main/cpp/safeAccess.cpp b/ddprof-lib/src/main/cpp/safeAccess.cpp index 8f4576530b..f423af9d9c 100644 --- a/ddprof-lib/src/main/cpp/safeAccess.cpp +++ b/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,6 +16,8 @@ #include "safeAccess.h" +#include +#include #include #include @@ -30,6 +32,28 @@ extern "C" int64_t safefetch64_cont(int64_t* adr, int64_t errValue); 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 diff --git a/ddprof-lib/src/main/cpp/safeAccess.h b/ddprof-lib/src/main/cpp/safeAccess.h index 151e6ee75d..14a009aad3 100644 --- a/ddprof-lib/src/main/cpp/safeAccess.h +++ b/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. diff --git a/ddprof-lib/src/test/cpp/safefetch_ut.cpp b/ddprof-lib/src/test/cpp/safefetch_ut.cpp index a5bb9941f5..769f891e9f 100644 --- a/ddprof-lib/src/test/cpp/safefetch_ut.cpp +++ b/ddprof-lib/src/test/cpp/safefetch_ut.cpp @@ -1,3 +1,8 @@ +/* + * Copyright 2026, Datadog, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + #include #include #include @@ -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}.