From fb690062cf4d97b0a154a9e93b8d0eef5f8f2bc9 Mon Sep 17 00:00:00 2001
From: Wiktor Garbacz <wiktorg@google.com>
Date: Wed, 23 Mar 2022 03:34:56 -0700
Subject: [PATCH] Fix a race when terminating sandbox from another thread

PiperOrigin-RevId: 436695251
Change-Id: I50599cefb346813f594982641c78dc902e10ccb5
---
 sandboxed_api/sandbox.cc   |  5 +++--
 sandboxed_api/sandbox.h    |  4 ++++
 sandboxed_api/sapi_test.cc | 17 +++++++++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/sandboxed_api/sandbox.cc b/sandboxed_api/sandbox.cc
index d49f3ed..d114087 100644
--- a/sandboxed_api/sandbox.cc
+++ b/sandboxed_api/sandbox.cc
@@ -200,6 +200,7 @@ absl::Status Sandbox::Init() {
 
   s2_ = absl::make_unique<sandbox2::Sandbox2>(std::move(executor),
                                               std::move(s2p), CreateNotifier());
+  s2_awaited_ = false;
   auto res = s2_->RunAsync();
 
   comms_ = s2_->comms();
@@ -432,9 +433,9 @@ absl::StatusOr<std::string> Sandbox::GetCString(const v::RemotePtr& str,
 }
 
 const sandbox2::Result& Sandbox::AwaitResult() {
-  if (s2_) {
+  if (s2_ && !s2_awaited_) {
     result_ = s2_->AwaitResult();
-    s2_.reset(nullptr);
+    s2_awaited_ = true;
   }
   return result_;
 }
diff --git a/sandboxed_api/sandbox.h b/sandboxed_api/sandbox.h
index 80d0534..7c3b8fc 100644
--- a/sandboxed_api/sandbox.h
+++ b/sandboxed_api/sandbox.h
@@ -149,6 +149,10 @@ class Sandbox {
 
   // The main sandbox2::Sandbox2 object.
   std::unique_ptr<sandbox2::Sandbox2> s2_;
+  // Marks whether Sandbox2 result was already fetched.
+  // We cannot just delete s2_ as Terminate might be called from another thread
+  // and comms object can be still in use then.
+  bool s2_awaited_ = false;
 
   // Result of the most recent sandbox execution
   sandbox2::Result result_;
diff --git a/sandboxed_api/sapi_test.cc b/sandboxed_api/sapi_test.cc
index 460609d..807f7e3 100644
--- a/sandboxed_api/sapi_test.cc
+++ b/sandboxed_api/sapi_test.cc
@@ -14,6 +14,8 @@
 
 #include <fcntl.h>
 
+#include <thread>  // NOLINT(build/c++11)
+
 #include "benchmark/benchmark.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -262,5 +264,20 @@ TEST(SandboxTest, NoRaceInAwaitResult) {
   EXPECT_THAT(result.final_status(), Eq(sandbox2::Result::VIOLATION));
 }
 
+TEST(SandboxTest, NoRaceInConcurrentTerminate) {
+  SumSandbox sandbox;
+  ASSERT_THAT(sandbox.Init(), IsOk());
+  SumApi api(&sandbox);
+  std::thread th([&sandbox] {
+    // Sleep so that the call already starts
+    absl::SleepFor(absl::Seconds(1));
+    sandbox.Terminate(/*attempt_graceful_exit=*/false);
+  });
+  EXPECT_THAT(api.sleep_for_sec(10), StatusIs(absl::StatusCode::kUnavailable));
+  th.join();
+  const auto& result = sandbox.AwaitResult();
+  EXPECT_THAT(result.final_status(), Eq(sandbox2::Result::EXTERNAL_KILL));
+}
+
 }  // namespace
 }  // namespace sapi