How heap buffer overflow happens in C parallel_memcpy() and how to fix it

How Heap Buffer Overflow Happens in C parallel_memcpy() and How to Fix It

Summary

A critical heap buffer overflow was discovered in csrc/cpu/comm/shm.cpp, the shared-memory communication layer used in distributed deep learning workloads. The parallel_memcpy() function — called from at least three sites in the same file — copies an arbitrary number of bytes into fixed-size destination buffers without ever checking whether the copy length exceeds the buffer's capacity. A co-located malicious process with access to the shared memory segment could manipulate the chunk_size value to be larger than MAX_BUF_SIZE (32 MB), triggering a heap buffer overflow that can corrupt memory, crash the process, or, in the worst case, enable arbitrary code execution.

Introduction

The csrc/cpu/comm/shm.cpp file implements the CPU-side shared-memory collective communication primitives — allreduce, broadcast, and related operations — used by distributed training frameworks. It manages large, fixed-size shared memory buffers (MAX_BUF_SIZE = 32 MB, NAIVE_ALLREDUCE_THRESHOLD = 1 MB) that multiple processes on the same machine read from and write to simultaneously.

The problem is straightforward but severe: parallel_memcpy() is a thin wrapper around a multi-threaded memory copy, and it trusts the n_bytes argument completely. The callers at lines 512, 593, and 634 derive chunk_size from data_size calculations, but those calculations rely on values that live in the shared memory segment — values that any co-located process can overwrite.

// Vulnerable call site (representative, line ~512)
parallel_memcpy(symmetric_buffer[buf_idx][i],
                src + offset,
                chunk_size);   // ← chunk_size comes from shared memory; never validated

If chunk_size exceeds MAX_BUF_SIZE, the copy walks off the end of the heap allocation and into whatever happens to be adjacent — other buffers, metadata, or function pointers.

The Vulnerability Explained

What parallel_memcpy does

parallel_memcpy divides a copy job across multiple CPU threads to maximize memory bandwidth. Conceptually:

// Simplified representation of the vulnerable function (before fix)
void parallel_memcpy(void* to, void* from, size_t n_bytes) {
    // Splits n_bytes across worker threads and calls memcpy on each chunk.
    // No check: is n_bytes <= sizeof(*to) ?
    size_t per_thread = n_bytes / num_threads;
    for (int t = 0; t < num_threads; t++) {
        threads[t] = std::thread(memcpy,
                                 (char*)to   + t * per_thread,
                                 (char*)from + t * per_thread,
                                 per_thread);
    }
    // ... join threads
}

The destination buffers are allocated with a hard ceiling:

#define MAX_BUF_SIZE            (32 * 1024 * 1024)   // 32 MB
#define NAIVE_ALLREDUCE_THRESHOLD (1 * 1024 * 1024)  // 1 MB

symmetric_buffer and distributed_buffer entries point into these fixed allocations. If chunk_size > MAX_BUF_SIZE, the write overflows.

The attacker's path

This is a 2-step exploit chain:

1. Shared memory manipulation: The attacker runs a process on the same machine and maps the same shared memory segment (e.g., via /dev/shm or a POSIX shm_open handle). They overwrite the data_size or related bookkeeping fields inside the shared segment to an inflated value — say, 64 MB or 320 MB.

2. Triggering the overflow: The victim process computes chunk_size from the attacker-controlled data_size, then calls parallel_memcpy(symmetric_buffer[buf_idx][i], src, chunk_size). Because chunk_size (64 MB) exceeds MAX_BUF_SIZE (32 MB), the copy writes 32 MB of attacker-influenced data beyond the end of the destination heap buffer.

Real-world impact

• Heap corruption leading to crash (denial of service against the training job).
• Overwriting adjacent heap objects such as other allreduce_workspace structs, potentially redirecting control flow.
• Data poisoning of gradient buffers in adjacent memory, silently corrupting model training without a crash.

In a multi-tenant GPU cluster or shared HPC environment, this is a realistic threat: co-located jobs share the same physical host and can reach the same shared memory namespace.

The Fix

Buffer-length check in parallel_memcpy callers

The primary fix adds a guard that clamps chunk_size to the destination buffer's declared capacity before passing it to parallel_memcpy. This ensures that no matter what value lives in shared memory, the copy never exceeds the allocation.

// After fix — representative call site
size_t safe_chunk = std::min(chunk_size, static_cast<size_t>(MAX_BUF_SIZE));
parallel_memcpy(symmetric_buffer[buf_idx][i],
                src + offset,
                safe_chunk);

The invariant enforced: buffer reads never exceed the declared length.

Switching from malloc to calloc for pointer arrays

The diff also changes the initialization of the pointer arrays that hold references to all ranks' shared buffers:

Before:

// malloc — allocated memory contains garbage values
workspace          = (struct allreduce_workspace**)malloc(size * sizeof(struct allreduce_workspace*));
symmetric_buffer[0] = (char**)malloc(size * sizeof(char**));
symmetric_buffer[1] = (char**)malloc(size * sizeof(char**));
distributed_buffer[0] = (char**)malloc(size * sizeof(char**));
distributed_buffer[1] = (char**)malloc(size * sizeof(char**));

After:

// calloc — zero-initializes every pointer; comment explains the invariant
workspace          = (struct allreduce_workspace**)calloc(size, sizeof(struct allreduce_workspace*));
symmetric_buffer[0] = (char**)calloc(size, sizeof(char*));
symmetric_buffer[1] = (char**)calloc(size, sizeof(char*));
distributed_buffer[0] = (char**)calloc(size, sizeof(char*));
distributed_buffer[1] = (char**)calloc(size, sizeof(char*));

Why this matters:

• malloc leaves memory uninitialized. If the loop that populates these arrays is interrupted (e.g., by a signal or an early error return), any later code that iterates over the array could dereference a garbage pointer — undefined behavior that can be exploited.
• calloc zero-initializes all entries. A NULL pointer dereference is a clean, detectable crash rather than a silent use of an attacker-influenced address.
• The comment // calloc used for defensive zero-init; the loop below writes every element before any read makes the invariant explicit for future maintainers.

Also note the subtle type correction: sizeof(char**) → sizeof(char*). The original code over-allocated each element by the size of an extra pointer level (8 bytes on 64-bit systems × size elements), a minor waste that the fix corrects.

Before/After Summary

Key Takeaways

• parallel_memcpy() in shm.cpp must never receive a chunk_size derived directly from shared memory without clamping it to MAX_BUF_SIZE first. Shared memory is an untrusted IPC channel.
• The three call sites at lines 512, 593, and 634 were all vulnerable — fixing only one would have left two exploitable paths open.
• malloc for pointer arrays that are populated in a subsequent loop is a latent use-before-init risk. calloc is the correct choice and costs almost nothing.
• Heap buffer overflows in distributed training infrastructure can corrupt gradient buffers silently, making the attack impact extend beyond crashes to subtly wrong model weights.
• Regression tests with adversarial sizes (2× and 10× the buffer limit) are essential for shared-memory code — boundary conditions at exact limits are not sufficient.

How Orbis AppSec Detected This

• Source: The data_size field inside the shared memory segment (workspace_buf->data_size), which any co-located process can overwrite via the shared memory handle.
• Sink: parallel_memcpy(symmetric_buffer[buf_idx][i], src + offset, chunk_size) at lines 512, 593, and 634 in csrc/cpu/comm/shm.cpp, where chunk_size flows from the attacker-controlled data_size without any bounds check.
• Missing control: No validation that chunk_size ≤ MAX_BUF_SIZE (or NAIVE_ALLREDUCE_THRESHOLD) before the copy; no assertion, no clamp, no early return.
• CWE: CWE-120 — Buffer Copy without Checking Size of Input ('Classic Buffer Overflow').
• Fix: Added a bounds check capping chunk_size to the destination buffer's declared capacity before each parallel_memcpy call, and replaced malloc with calloc for zero-safe initialization of the pointer arrays.

Orbis AppSec automatically detected this vulnerability and opened a pull request with the fix. Try Orbis AppSec on your repositories to find and fix issues like this automatically.

Conclusion

The parallel_memcpy overflow in csrc/cpu/comm/shm.cpp is a textbook example of why shared memory must be treated as untrusted input. The function itself was not malicious — it was doing exactly what it was told. The flaw was the implicit assumption that callers would always supply a safe size. In a shared-memory environment, that assumption is never safe: any co-located process can break it.

The fix is small — a bounds clamp before each copy and a switch from malloc to calloc — but the security improvement is substantial. Combined with the regression test suite that probes adversarial sizes up to 10× the buffer limit, this change makes the collective communication layer meaningfully more resilient against co-located attackers.

When writing C/C++ code that operates on shared memory, always ask: "What happens if a value I read from this segment is 10× larger than I expect?" If the answer is a heap overflow, add the bounds check before shipping.