mm/fork: same-VA COW stack for deep call chains (opt-in)

Give a forked child its parent's EXACT user stack -- same virtual
addresses, private physical pages -- instead of a relocated copy with a
biased SP.  A relocated+biased stack leaves every saved rbp, return
address and &local in the copied frames pointing at the parent's stack
(off by the bias), so a deep unwind in the child dereferences the parent
and faults.  tst-fork-deep (fork 12 frames deep, child unwinds all the
way up and _exit()s a checked value) is the test that catches this; it
now passes.  Gated behind CONF_fork (default n); conf_fork=0 is unchanged
and builds/boots clean (fork code fully excluded, tst-mmap passes).

Four coupled changes make same-VA correct, layered on the Stage 2 COW
address space and the fork-child lifecycle fix (detached child +
dispose(child) in the reaper):

1. core/mmu.cc: clone_address_space() PRIVATIZES the forking thread's
   live stack VA range into the child AS -- fresh private pages that
   byte-copy the parent's stack, mapped writable (not COW: OSv runs
   kernel code with irqs off on the app stack, a COW fault there is
   illegal) at the SAME VA.  The parent's PTEs are untouched.

2. arch/x64/fork.cc + libc/process/fork.cc + include/osv/fork.hh: the
   child resumes with the caller's FULL callee-saved register context
   (rbx, rbp, r12-r15, rsp, rip), captured at fork() entry into
   osv::fork_resume_ctx and restored by the child trampoline off a
   scratch base register (so the restore never clobbers its own base).
   fork() jmps past its own epilogue, so without this the caller would
   resume with fork()'s internal register values -- a deep chain that
   keeps an accumulator in a callee-saved reg across fork() computes
   garbage.

3. arch/x64/arch-switch.hh: defer the fork CR3 switch to EXACTLY the
   rsp/rbp swap inside switch_to()'s asm.  Loading the incoming AS's
   CR3 early, while still executing on the outgoing thread's stack (the
   fpu-cw/mxcsr scratch saves), resolved those shared-VA stack writes
   through the wrong address space and clobbered a blocked thread's live
   stack (root-caused via hardware watchpoint: switch_to wrote garbage
   into a blocked parent's on-stack user_mutex slot in AS0).

4. core/lfmutex.cc + core/mmu.cc: a wait_record on a shared kernel
   condvar/mutex must live in the identity-mapped kernel heap (coherent
   in every AS) whenever a cross-address-space wake is possible.  Extend
   the Option-A rule from 'current thread is a fork child' to also
   cover 'any child address space is live' (live_child_address_spaces
   counter), so the parent in AS0 -- woken by a child that walks the
   queue through its own page tables -- is safe too.

libc/process/execve.cc: on a fork child, move the thread onto its own
kernel stack before switching to AS0 and running the target, so
reloading CR3 to AS0 cannot alias the same-VA app stack onto the
parent's physical page (which would corrupt a blocked parent).  The
child address space is destroyed HERE, on the kernel stack, exactly once:
the reap-time cleanup (libc/process/fork.cc) then sees the thread's AS is
AS0 (!= child_as) and skips its own destroy, so destroy_address_space()
runs once whether the child exits normally (reaper destroys) or execs
(execve destroys) -- no leak, no double-free.  A pre-exec existence check
keeps a failed exec returning to the caller.

Validated on x86-64 (KVM): tst-fork 10/10, tst-fork-cow 6/6,
tst-fork-deep 0 failures, tst-execve 3/3 with clean shutdown (no hang, no
double-free at reap), on smp=1 and smp=4, 5/5 repeat runs; conf_fork=0
builds and boots clean (fork excluded, tst-mmap passes).

Author: gburd

arch/aarch64/fork.cc

@@ -22,8 +22,9 @@
 extern "C" void __osv_run_atfork_child();
 
 sched::thread *fork_thread(void *caller_ret, void *caller_sp,
-                           void **out_stack_to_free)
+                           void *resume_ctx, void **out_stack_to_free)
 {
+    (void)resume_ctx;   // aarch64 uses caller_sp bias-copy (same-VA is x86-64)
     auto parent = sched::thread::current();
     auto parent_pinned_cpu = parent->pinned() ? sched::cpu::current() : nullptr;
 

arch/x64/arch-switch.hh

@@ -95,14 +95,24 @@ void thread::switch_to()
     set_fsbase(reinterpret_cast<u64>(_tcb));
     barrier();
 #if CONF_fork
-    // Address-space (CR3) switch for fork COW.  Only touch CR3 when the target
+    // Address-space (CR3) switch for fork COW.  Only switch when the target
     // thread lives in a different address space than the outgoing one, so the
-    // common single-address-space case pays nothing (and no TLB flush).  The
-    // kernel half of every AS is identically mapped, so the switch code, kernel
-    // stacks and kernel heap remain valid across the write.
-    if (_current_as != old->_current_as) {
-        processor::write_cr3(mmu::pt_root_phys(_current_as));
-    }
+    // common single-address-space case pays nothing (and no TLB flush).
+    //
+    // CRITICAL (same-VA fork stacks): the CR3 write must happen EXACTLY at the
+    // rsp/rbp swap below, NOT here.  OSv runs kernel code (this switch, the
+    // fpu-cw/mxcsr scratch saves, the register-save asm) on the app thread's
+    // stack.  A forked child resumes on the parent's EXACT stack VAs, so parent
+    // and child alias the same stack virtual addresses (backed by different
+    // physical pages per address space).  If we loaded the incoming thread's
+    // CR3 here, the subsequent scratch writes to the OUTGOING thread's stack
+    // (fnstcw/stmxcsr at -0x34(%rbp), etc., with rsp/rbp still the old thread's)
+    // would resolve through the WRONG address space and clobber the incoming
+    // thread's physical page at that shared VA -- corrupting a blocked thread's
+    // live stack.  So we defer the CR3 load into the asm, coincident with the
+    // rsp/rbp swap: old-stack writes use the old AS, new-stack writes the new.
+    bool switch_as = (_current_as != old->_current_as);
+    u64 new_cr3 = switch_as ? mmu::pt_root_phys(_current_as) : 0;
 #endif
     auto c = _detached_state->_cpu;
     old->_state.exception_stack = c->arch.get_exception_stack();
@@ -119,6 +129,34 @@ void thread::switch_to()
     c->arch._current_thread_kernel_tcb = reinterpret_cast<u64>(_tcb);
     auto fpucw = processor::fnstcw();
     auto mxcsr = processor::stmxcsr();
+#if CONF_fork
+    if (switch_as) {
+        // Same-VA-safe switch: swap rsp/rbp to the incoming thread, THEN load
+        // its CR3 (mov %rax,%cr3), THEN jump.  All old-stack writes above ran
+        // under the old AS; every access after the rsp/rbp swap is to the new
+        // thread's stack, now resolved through the new AS -- so a shared stack
+        // VA never resolves through the wrong address space.  Kernel .text and
+        // the thread_state (heap) are identically mapped in every AS, so the
+        // instruction fetch across the mov-to-cr3 and the %c[rip] load are
+        // valid before and after the CR3 write.
+        asm volatile
+            ("mov %%rbp, %c[rbp](%0) \n\t"
+             "movq $1f, %c[rip](%0) \n\t"
+             "mov %%rsp, %c[rsp](%0) \n\t"
+             "mov %c[rsp](%1), %%rsp \n\t"
+             "mov %c[rbp](%1), %%rbp \n\t"
+             "mov %2, %%cr3 \n\t"
+             "jmpq *%c[rip](%1) \n\t"
+             "1: \n\t"
+             :
+             : "a"(&old->_state), "c"(&this->_state), "d"(new_cr3),
+               [rsp]"i"(offsetof(thread_state, rsp)),
+               [rbp]"i"(offsetof(thread_state, rbp)),
+               [rip]"i"(offsetof(thread_state, rip))
+             : "rbx", "rsi", "rdi", "r8", "r9",
+               "r10", "r11", "r12", "r13", "r14", "r15", "memory");
+    } else
+#endif
     asm volatile
         ("mov %%rbp, %c[rbp](%0) \n\t"
          "movq $1f, %c[rip](%0) \n\t"

arch/x64/fork.cc

@@ -4,17 +4,26 @@
  * This work is open source software, licensed under the terms of the
  * BSD license as described in the LICENSE file in the top-level directory.
  *
- * fork_thread(): create a child thread that resumes in fork()'s CALLER, on a
- * private copy of the parent's user stack, returning 0 from fork() in the child.
- * The x86-64 arch half of the fork() emulation (see documentation/fork.md).
+ * fork_thread(): create a child thread that resumes in fork()'s CALLER on the
+ * SAME user stack virtual addresses the parent had -- no stack relocation, no
+ * SP bias -- and with the caller's FULL callee-saved register context restored,
+ * exactly as a normal `ret` from fork() would leave it.  The child's address
+ * space (built by clone_address_space) maps the parent's live stack VA range to
+ * freshly-allocated PRIVATE physical pages that byte-copy the parent's stack,
+ * so the child sees identical stack contents at identical addresses but writes
+ * to its own private frames.
  *
- * fork() (libc/process/fork.cc) passes us the caller's resume point:
- *   caller_ret  = the address fork() would return to  (__builtin_return_address)
- *   caller_sp   = the parent's SP at fork()'s return   (fork()'s frame base)
- * We copy the parent stack region [caller_sp .. stack_base) into a fresh stack,
- * bias caller_sp into the copy, and start a child thread whose trampoline sets
- * rsp=child_sp, rax=0, and jumps to caller_ret -- i.e. the child returns from
- * fork() with value 0 on its own private stack, in the caller.
+ * Why same-VA + full register restore: OSv/x86-64 code addresses saved frame
+ * pointers (rbp), return addresses, and &local pointers as absolute stack VAs,
+ * and keeps live values in callee-saved registers (rbx, r12-r15) across calls.
+ * Relocating the child stack to a different VA and biasing rsp leaves every
+ * saved rbp/&local pointing at the PARENT's stack (off by the bias), and simply
+ * jmp-ing to the return address skips fork()'s epilogue so the caller resumes
+ * with fork()'s internal register values.  Either corrupts a deep unwind.
+ * Keeping the child on the parent's exact VAs (private phys) AND restoring the
+ * full caller register context (osv::fork_resume_ctx, captured at fork() entry)
+ * makes the child continue in fork()'s caller as if fork() had returned 0 -- the
+ * only correct way to fork a deep stack (see tst-fork-deep, documentation/fork.md).
  */
 
 #include "arch.hh"
@@ -23,13 +32,16 @@
 #include <string.h>
 #include <cstdlib>
 #include <osv/sched.hh>
+#include <osv/fork.hh>
 
 // pthread_atfork child-handler chain (defined in libc/pthread.cc), run in the
 // child's context before it resumes user code.
 extern "C" void __osv_run_atfork_child();
 
-sched::thread *fork_thread(void *caller_ret, void *caller_sp, void **out_stack_to_free)
+sched::thread *fork_thread(void *caller_ret, void *caller_sp,
+                           void *resume_ctx, void **out_stack_to_free)
 {
+    auto ctx = static_cast<osv::fork_resume_ctx*>(resume_ctx);
     auto parent = sched::thread::current();
     auto parent_pinned_cpu = parent->pinned() ? sched::cpu::current() : nullptr;
 
@@ -39,57 +51,56 @@ sched::thread *fork_thread(void *caller_ret, void *caller_sp, void **out_stack_t
     if (sp < static_cast<char*>(si.begin) || sp > stack_base) {
         return nullptr;   // caller SP not within the known user stack
     }
-    size_t stack_size = si.size;
 
-    char *child_stack_mem = static_cast<char*>(malloc(stack_size));
-    if (!child_stack_mem) {
-        return nullptr;
-    }
-    // Copy ONLY the live top of the stack, [caller_sp .. stack_base), into the
-    // TOP of the child buffer.  App (pthread) stacks are demand-paged: only the
-    // used top is mapped, so copying from si.begin would fault on the first
-    // unmapped page.  Keeping the copy at the top of the child buffer preserves
-    // the base-relative bias so a biased SP resolves correctly.
-    char *child_base = child_stack_mem + stack_size;
-    ptrdiff_t bias = child_base - stack_base;
-    size_t live = static_cast<size_t>(stack_base - sp);
-    memcpy(child_base - live, sp, live);
-    char *child_sp = sp + bias;
-
-    volatile u64 resume_sp = reinterpret_cast<u64>(child_sp);
-    volatile u64 resume_pc = reinterpret_cast<u64>(caller_ret);
-    char *stack_to_free = child_stack_mem;
+    // Same-VA stack: the child resumes on the parent's EXACT register context.
+    // No copy and no bias here -- clone_address_space() privatizes the parent's
+    // live stack VA range into the child's address space (fresh private pages
+    // that byte-copy the parent's stack), so these very addresses are valid and
+    // private in the child.  Capture the resume context by value in the lambda
+    // (the parent's on-stack ctx is gone by the time the child runs).
+    (void)caller_ret;
+    osv::fork_resume_ctx rc = *ctx;
 
     // TLS handling.  The child is a real OSv sched::thread, so its constructor
-    // already ran setup_tcb() and installed a FRESH, private OSv TLS block
-    // (with its own errno and all libc __thread state).  Two cases:
-    //
-    //  (1) The app uses OSv's libc TLS (the normal dynamically-linked path,
-    //      app_tcb == 0): the child's own fresh TCB is exactly right -- do NOT
-    //      touch fsbase, let the child run on its private per-thread TLS.  This
-    //      is the clean case and fork() "just works" for TLS.
-    //  (2) The app installed its own TCB via arch_prctl(SET_FS) (app_tcb != 0,
-    //      e.g. a glibc binary's __libc_setup_tls): the child would need a
-    //      private COPY of that app TCB.  We do not duplicate it here yet;
-    //      the child inherits the parent's app_tcb (shared), which is the
-    //      documented multi-process-glibc limitation.  A musl app built against
-    //      OSv's libc takes path (1) and avoids this entirely.
+    // already ran setup_tcb() and installed a FRESH, private OSv TLS block.  If
+    // the parent installed its own app TCB via arch_prctl (app_tcb != 0), the
+    // child inherits it (shared) -- the documented multi-process-glibc limit;
+    // otherwise the child keeps its own private OSv per-thread TLS.
     u64 parent_app_tcb = parent->get_app_tcb();
 
-    auto t = sched::thread::make([resume_sp, resume_pc, parent_app_tcb] {
-        // Only override the child's own (fresh) TLS if the parent had installed
-        // an app TCB via arch_prctl; otherwise keep the child's private OSv TCB.
+    auto t = sched::thread::make([rc, parent_app_tcb] {
         if (parent_app_tcb) {
             arch::set_fsbase(parent_app_tcb);
         }
         // Run pthread_atfork child handlers in the child's context (e.g. reset
         // the malloc arena lock) before resuming user code.
         __osv_run_atfork_child();
+        // Restore the caller's full callee-saved register context and resume in
+        // fork()'s caller with return value 0, on the parent's exact stack VAs
+        // (private phys in the child AS).  We base all loads off a scratch
+        // register (rax) pointing at a LOCAL copy of the context, and restore
+        // the base-conflicting registers (rbx, rbp) LAST, so the load sequence
+        // never clobbers its own base pointer.
+        osv::fork_resume_ctx c = rc;   // local copy the asm can address stably
+        // %0 = &c pinned in a register that is NOT one of the restored regs
+        // (we let the compiler pick, but we consume it only to seed rax).  We
+        // move &c into rax first, then load every callee-saved reg from rax-
+        // relative offsets, and load rsp last -- rax (the base) is never in the
+        // restore set, so the sequence never clobbers its own base pointer.
+        // Offsets match struct fork_resume_ctx { rbx,rbp,r12,r13,r14,r15,rsp,rip }.
         asm volatile
-          ("movq %0, %%rsp \n\t"    // install the private copied stack
-           "xorq %%rax, %%rax \n\t" // fork() returns 0 in the child
-           "jmpq *%1 \n\t"          // resume in fork()'s caller
-           : : "r"(resume_sp), "r"(resume_pc) : "memory");
+          ("movq %0, %%rax        \n\t"  // rax = &c (base; not restored)
+           "movq  0(%%rax), %%rbx \n\t"  // rbx
+           "movq  8(%%rax), %%rbp \n\t"  // rbp
+           "movq 16(%%rax), %%r12 \n\t"  // r12
+           "movq 24(%%rax), %%r13 \n\t"  // r13
+           "movq 32(%%rax), %%r14 \n\t"  // r14
+           "movq 40(%%rax), %%r15 \n\t"  // r15
+           "movq 56(%%rax), %%rcx \n\t"  // rcx = caller rip (scratch)
+           "movq 48(%%rax), %%rsp \n\t"  // adopt parent's exact rsp
+           "xorq %%rax, %%rax     \n\t"  // fork() returns 0 in the child
+           "jmpq *%%rcx           \n\t"  // resume in fork()'s caller
+           : : "r"(&c) : "rax", "rcx", "memory");
     }, sched::thread::attr().
         stack(4096 * 4).
         // Detached: nobody join()s the fork child (the parent reaps it via the
@@ -106,10 +117,10 @@ sched::thread *fork_thread(void *caller_ret, void *caller_sp, void **out_stack_t
     if (parent_pinned_cpu) {
         t->pin(parent_pinned_cpu);
     }
-    // The caller (fork.cc) owns the single cleanup; hand back the copied user
-    // stack so it can be freed when the child is reaped.
+    // Same-VA: no separate user-stack buffer to free (the child's stack pages
+    // are owned by its address space and freed on destroy_address_space()).
     if (out_stack_to_free) {
-        *out_stack_to_free = stack_to_free;
+        *out_stack_to_free = nullptr;
     }
     return t;
 }

core/lfmutex.cc

@@ -16,11 +16,27 @@
 #endif
 
 #if CONF_fork
-// See wait_record.hh: true iff the current thread runs in a forked-child
-// (non-AS0) address space, in which case a wait_record queued on a shared
-// kernel object must be heap-allocated to stay coherent cross-address-space.
+namespace mmu {
+    // Live fork-child address-space count (defined in core/mmu.cc).  When > 0,
+    // more than one address space exists and a stack-resident wait_record on a
+    // shared kernel object can be dereferenced cross-AS.
+    extern std::atomic<int> live_child_address_spaces;
+}
+
+// See wait_record.hh: with per-child same-VA private stacks, a wait_record
+// queued on a SHARED kernel condvar/mutex must live in the identity-mapped
+// kernel heap (same VA->phys in every address space) whenever a cross-address-
+// space wake is possible.  That is true when EITHER the current thread runs in
+// a forked-child (non-AS0) address space (a child woken by the parent), OR any
+// child address space is currently live (the parent, in AS0, woken by a child
+// -- the child dereferences the parent's stack-resident wait_record through
+// its own page tables and would read a stale private stack copy).  Default OSv
+// (no fork, single AS) always takes the zero-overhead on-stack fast path.
 bool fork_child_needs_heap_wait_record()
 {
+    if (mmu::live_child_address_spaces.load(std::memory_order_relaxed) > 0) {
+        return true;
+    }
     auto t = sched::thread::current();
     return t && t->address_space() &&
            t->address_space() != mmu::kernel_address_space();