Track whether a thread is currently vforked. am: 5378351432

Change-Id: I70d09823175a35cee7136594061d6c05c66cead9

6 files changed, 97 insertions, 13 deletions

diff --git a/libc/arch-arm/bionic/vfork.S b/libc/arch-arm/bionic/vfork.S
index 6855db74b..a964be53a 100644
--- a/libc/arch-arm/bionic/vfork.S
+++ b/libc/arch-arm/bionic/vfork.S
@@ -31,17 +31,27 @@
 
 ENTRY(vfork)
 __BIONIC_WEAK_ASM_FOR_NATIVE_BRIDGE(vfork)
-    // __get_tls()[TLS_SLOT_THREAD_ID]->cached_pid_ = 0
+    // r3 = &__get_tls()[TLS_SLOT_THREAD_ID]
     mrc     p15, 0, r3, c13, c0, 3
     ldr     r3, [r3, #(TLS_SLOT_THREAD_ID * 4)]
-    mov     r0, #0
+
+    // Set cached_pid_ to 0, vforked_ to 1, and stash the previous value.
+    mov     r0, #0x80000000
+    ldr     r1, [r3, #12]
     str     r0, [r3, #12]
 
     mov     ip, r7
     ldr     r7, =__NR_vfork
     swi     #0
     mov     r7, ip
+
+    teq     r0, #0
+    bxeq    lr
+
+    // rc != 0: reset cached_pid_ and vforked_.
+    str     r1, [r3, #12]
     cmn     r0, #(MAX_ERRNO + 1)
+
     bxls    lr
     neg     r0, r0
     b       __set_errno_internal
diff --git a/libc/arch-arm64/bionic/vfork.S b/libc/arch-arm64/bionic/vfork.S
index a76e9ca39..5cfb8b0cc 100644
--- a/libc/arch-arm64/bionic/vfork.S
+++ b/libc/arch-arm64/bionic/vfork.S
@@ -36,10 +36,14 @@
 
 ENTRY(vfork)
 __BIONIC_WEAK_ASM_FOR_NATIVE_BRIDGE(vfork)
-    // __get_tls()[TLS_SLOT_THREAD_ID]->cached_pid_ = 0
-    mrs     x0, tpidr_el0
-    ldr     x0, [x0, #(TLS_SLOT_THREAD_ID * 8)]
-    str     wzr, [x0, #20]
+    // x9 = __get_tls()[TLS_SLOT_THREAD_ID]
+    mrs     x9, tpidr_el0
+    ldr     x9, [x9, #(TLS_SLOT_THREAD_ID * 8)]
+
+    // Set cached_pid_ to 0, vforked_ to 1, and stash the previous value.
+    mov     w0, #0x80000000
+    ldr     w10, [x9, #20]
+    str     w0, [x9, #20]
 
     mov     x0, #(CLONE_VM | CLONE_VFORK | SIGCHLD)
     mov     x1, xzr
@@ -50,6 +54,10 @@ __BIONIC_WEAK_ASM_FOR_NATIVE_BRIDGE(vfork)
     mov     x8, __NR_clone
     svc     #0
 
+    cbz     x0, .L_exit
+
+    // rc != 0: reset cached_pid_ and vforked_.
+    str     w10, [x9, #20]
     cmn     x0, #(MAX_ERRNO + 1)
     cneg    x0, x0, hi
     b.hi    __set_errno_internal
diff --git a/libc/arch-x86/bionic/vfork.S b/libc/arch-x86/bionic/vfork.S
index 663169c9b..231a36ec7 100644
--- a/libc/arch-x86/bionic/vfork.S
+++ b/libc/arch-x86/bionic/vfork.S
@@ -37,13 +37,25 @@ __BIONIC_WEAK_ASM_FOR_NATIVE_BRIDGE(vfork)
   .cfi_adjust_cfa_offset 4
   .cfi_rel_offset ecx, 0
 
-  // __get_tls()[TLS_SLOT_THREAD_ID]->cached_pid_ = 0
+  // Set cached_pid_ to 0, vforked_ to 1, and stash the previous value.
   movl    %gs:0, %eax
   movl    (TLS_SLOT_THREAD_ID * 4)(%eax), %eax
-  movl    $0, 12(%eax)
+  movl    12(%eax), %edx
+  movl    $0x80000000, 12(%eax)
 
   movl    $__NR_vfork, %eax
   int     $0x80
+
+  test    %eax, %eax
+  jz      1f
+
+  // rc != 0: restore the previous cached_pid_/vforked_ values.
+  pushl   %ecx
+  movl    %gs:0, %ecx
+  movl    (TLS_SLOT_THREAD_ID * 4)(%ecx), %ecx
+  movl    %edx, 12(%ecx)
+  popl    %ecx
+
   cmpl    $-MAX_ERRNO, %eax
   jb      1f
   negl    %eax
diff --git a/libc/arch-x86_64/bionic/vfork.S b/libc/arch-x86_64/bionic/vfork.S
index 86c5db2a1..8cfcc3658 100644
--- a/libc/arch-x86_64/bionic/vfork.S
+++ b/libc/arch-x86_64/bionic/vfork.S
@@ -35,14 +35,22 @@ ENTRY(vfork)
 __BIONIC_WEAK_ASM_FOR_NATIVE_BRIDGE(vfork)
   popq    %rdi  // Grab the return address.
 
-  // __get_tls()[TLS_SLOT_THREAD_ID]->cached_pid_ = 0
-  mov    %fs:0, %rax
-  mov    (TLS_SLOT_THREAD_ID * 8)(%rax), %rax
-  movl   $0, 20(%rax)
+  // Set cached_pid_ to 0, vforked_ to 1, and stash the previous value.
+  mov     %fs:0, %r8
+  mov     (TLS_SLOT_THREAD_ID * 8)(%r8), %r8
+  movl    20(%r8), %r9d
+  movl    $0x80000000, 20(%r8)
 
   movl    $__NR_vfork, %eax
   syscall
   pushq   %rdi  // Restore the return address.
+
+  test    %eax, %eax
+  jz      1f
+
+  // rc != 0: restore the previous cached_pid_/vforked_ values.
+  movl    %r9d, 20(%r8)
+
   cmpq    $-MAX_ERRNO, %rax
   jb      1f
   negl    %eax
diff --git a/libc/bionic/pthread_internal.h b/libc/bionic/pthread_internal.h
index ab8b9556d..1f055f59e 100644
--- a/libc/bionic/pthread_internal.h
+++ b/libc/bionic/pthread_internal.h
@@ -70,9 +70,12 @@ class pthread_internal_t {
   pid_t tid;
 
  private:
-  pid_t cached_pid_;
+  uint32_t cached_pid_ : 31;
+  uint32_t vforked_ : 1;
 
  public:
+  bool is_vforked() { return vforked_; }
+
   pid_t invalidate_cached_pid() {
     pid_t old_value;
     get_cached_pid(&old_value);
diff --git a/tests/unistd_test.cpp b/tests/unistd_test.cpp
index f3b08c3f4..6b2856169 100644
--- a/tests/unistd_test.cpp
+++ b/tests/unistd_test.cpp
@@ -41,6 +41,10 @@
 
 #include "private/get_cpu_count_from_string.h"
 
+#if defined(__BIONIC__)
+#include "bionic/pthread_internal.h"
+#endif
+
 #if defined(NOFORTIFY)
 #define UNISTD_TEST unistd_nofortify
 #define UNISTD_DEATHTEST unistd_nofortify_DeathTest
@@ -431,6 +435,45 @@ TEST(UNISTD_TEST, syncfs) {
   TestSyncFunction(syncfs);
 }
 
+TEST(UNISTD_TEST, vfork) {
+#if defined(__BIONIC__)
+  pthread_internal_t* self = __get_thread();
+
+  pid_t cached_pid;
+  ASSERT_TRUE(self->get_cached_pid(&cached_pid));
+  ASSERT_EQ(syscall(__NR_getpid), cached_pid);
+  ASSERT_FALSE(self->is_vforked());
+
+  pid_t rc = vfork();
+  ASSERT_NE(-1, rc);
+  if (rc == 0) {
+    if (self->get_cached_pid(&cached_pid)) {
+      const char* error = "__get_thread()->cached_pid_ set after vfork\n";
+      write(STDERR_FILENO, error, strlen(error));
+      _exit(1);
+    }
+
+    if (!self->is_vforked()) {
+      const char* error = "__get_thread()->vforked_ not set after vfork\n";
+      write(STDERR_FILENO, error, strlen(error));
+      _exit(1);
+    }
+
+    _exit(0);
+  } else {
+    ASSERT_TRUE(self->get_cached_pid(&cached_pid));
+    ASSERT_EQ(syscall(__NR_getpid), cached_pid);
+    ASSERT_FALSE(self->is_vforked());
+
+    int status;
+    pid_t wait_result = waitpid(rc, &status, 0);
+    ASSERT_EQ(wait_result, rc);
+    ASSERT_TRUE(WIFEXITED(status));
+    ASSERT_EQ(0, WEXITSTATUS(status));
+  }
+#endif
+}
+
 static void AssertGetPidCorrect() {
   // The loop is just to make manual testing/debugging with strace easier.
   pid_t getpid_syscall_result = syscall(__NR_getpid);