Allow collecting stacktraces on normal process exit

This mainly a debugging facility.
It makes diagnosing problems where sandboxed process just randomly exits whereas unsandboxed one runs to completion due to differences in the setup/environment much easier.

PiperOrigin-RevId: 391005548
Change-Id: Ia19fe6632748da93c1f4291bb55e895f50a4e2b0
diff --git a/sandboxed_api/sandbox2/monitor.cc b/sandboxed_api/sandbox2/monitor.cc
index 26528dc..942e7c2 100644
--- a/sandboxed_api/sandbox2/monitor.cc
+++ b/sandboxed_api/sandbox2/monitor.cc
@@ -360,6 +360,8 @@
       return policy_->collect_stacktrace_on_signal_;
     case Result::VIOLATION:
       return policy_->collect_stacktrace_on_violation_;
+    case Result::OK:
+      return policy_->collect_stacktrace_on_exit_;
     default:
       return false;
   }
@@ -861,7 +863,9 @@
 
 void Monitor::EventPtraceExit(pid_t pid, int event_msg) {
   // A regular exit, let it continue (fast-path).
-  if (WIFEXITED(event_msg)) {
+  if (ABSL_PREDICT_TRUE(
+          WIFEXITED(event_msg) &&
+          (!policy_->collect_stacktrace_on_exit_ || pid != pid_))) {
     ContinueProcess(pid, 0);
     return;
   }
@@ -883,10 +887,11 @@
     return;
   }
 
-  // This can be reached in three cases:
+  // This can be reached in four cases:
   // 1) Process was killed from the sandbox.
   // 2) Process was killed because it hit a timeout.
   // 3) Regular signal/other exit cause.
+  // 4) Normal exit for which we want to obtain stack trace.
   if (pid == pid_) {
     VLOG(1) << "PID: " << pid << " main special exit";
     if (network_violation_) {
@@ -896,6 +901,8 @@
       SetExitStatusCode(Result::EXTERNAL_KILL, 0);
     } else if (timed_out_) {
       SetExitStatusCode(Result::TIMEOUT, 0);
+    } else if (WIFEXITED(event_msg)) {
+      SetExitStatusCode(Result::OK, WEXITSTATUS(event_msg));
     } else {
       SetExitStatusCode(Result::SIGNALED, WTERMSIG(event_msg));
     }
diff --git a/sandboxed_api/sandbox2/policy.h b/sandboxed_api/sandbox2/policy.h
index edd5cb8..1ac7a96 100644
--- a/sandboxed_api/sandbox2/policy.h
+++ b/sandboxed_api/sandbox2/policy.h
@@ -96,6 +96,7 @@
   bool collect_stacktrace_on_signal_ = true;
   bool collect_stacktrace_on_timeout_ = true;
   bool collect_stacktrace_on_kill_ = true;
+  bool collect_stacktrace_on_exit_ = false;
 
   // The capabilities to keep in the sandboxee.
   std::vector<int> capabilities_;
diff --git a/sandboxed_api/sandbox2/policybuilder.cc b/sandboxed_api/sandbox2/policybuilder.cc
index 54df4d0..5b9c9e1 100644
--- a/sandboxed_api/sandbox2/policybuilder.cc
+++ b/sandboxed_api/sandbox2/policybuilder.cc
@@ -807,6 +807,7 @@
   output->collect_stacktrace_on_violation_ = collect_stacktrace_on_violation_;
   output->collect_stacktrace_on_timeout_ = collect_stacktrace_on_timeout_;
   output->collect_stacktrace_on_kill_ = collect_stacktrace_on_kill_;
+  output->collect_stacktrace_on_exit_ = collect_stacktrace_on_exit_;
   output->user_policy_ = std::move(user_policy_);
   output->user_policy_handles_bpf_ = user_policy_handles_bpf_;
 
@@ -959,6 +960,11 @@
   return *this;
 }
 
+PolicyBuilder& PolicyBuilder::CollectStacktracesOnExit(bool enable) {
+  collect_stacktrace_on_exit_ = enable;
+  return *this;
+}
+
 PolicyBuilder& PolicyBuilder::AddNetworkProxyPolicy() {
   if (allowed_hosts_) {
     SetError(absl::FailedPreconditionError(
diff --git a/sandboxed_api/sandbox2/policybuilder.h b/sandboxed_api/sandbox2/policybuilder.h
index c4aa2ab..db2fea9 100644
--- a/sandboxed_api/sandbox2/policybuilder.h
+++ b/sandboxed_api/sandbox2/policybuilder.h
@@ -540,6 +540,9 @@
   // monitor / the user.
   PolicyBuilder& CollectStacktracesOnKill(bool enable);
 
+  // Enables/disables stack trace collection on normal process exit.
+  PolicyBuilder& CollectStacktracesOnExit(bool enable);
+
   // Appends an unconditional ALLOW action for all syscalls.
   // Do not use in environment with untrusted code and/or data, ask
   // sandbox-team@ first if unsure.
@@ -590,6 +593,7 @@
   bool collect_stacktrace_on_signal_ = true;
   bool collect_stacktrace_on_timeout_ = true;
   bool collect_stacktrace_on_kill_ = false;
+  bool collect_stacktrace_on_exit_ = false;
 
   // Seccomp fields
   std::vector<sock_filter> user_policy_;
diff --git a/sandboxed_api/sandbox2/sandbox2_test.cc b/sandboxed_api/sandbox2/sandbox2_test.cc
index 7601be8..b74545f 100644
--- a/sandboxed_api/sandbox2/sandbox2_test.cc
+++ b/sandboxed_api/sandbox2/sandbox2_test.cc
@@ -118,6 +118,26 @@
   ASSERT_EQ(result.final_status(), Result::OK);
 }
 
+TEST(StackTraceTest, StackTraceOnExitWorks) {
+  SKIP_SANITIZERS_AND_COVERAGE;
+
+  const std::string path = GetTestSourcePath("sandbox2/testcases/minimal");
+  std::vector<std::string> args = {path};
+  auto executor = absl::make_unique<Executor>(path, args);
+
+  SAPI_ASSERT_OK_AND_ASSIGN(auto policy,
+                            PolicyBuilder()
+                                // Don't restrict the syscalls at all.
+                                .DangerDefaultAllowAll()
+                                .CollectStacktracesOnExit(true)
+                                .TryBuild());
+  Sandbox2 sandbox(std::move(executor), std::move(policy));
+  auto result = sandbox.Run();
+
+  ASSERT_EQ(result.final_status(), Result::OK);
+  ASSERT_THAT(result.stack_trace(), Not(IsEmpty()));
+}
+
 // Tests that we return the correct state when the sandboxee was killed by an
 // external signal. Also make sure that we do not have the stack trace.
 TEST(RunAsyncTest, SandboxeeExternalKill) {