From 053833d1039e46a0315036d3940918772e10469e Mon Sep 17 00:00:00 2001
From: hai007 <hais@google.com>
Date: Mon, 15 Mar 2021 16:35:14 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 363058247
---
 cpp/core/internal/endpoint_manager.cc | 41 +++++++++++++++++++++++----
 cpp/core/internal/endpoint_manager.h  |  5 ++--
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/cpp/core/internal/endpoint_manager.cc b/cpp/core/internal/endpoint_manager.cc
index f2ff0833..a5d04132 100644
--- a/cpp/core/internal/endpoint_manager.cc
+++ b/cpp/core/internal/endpoint_manager.cc
@@ -47,7 +47,7 @@ constexpr absl::Time EndpointManager::kInvalidTimestamp;
 //           the loop.
 void EndpointManager::EndpointChannelLoopRunnable(
     const std::string& runnable_name, ClientProxy* client,
-    const std::string& endpoint_id, CountDownLatch* barrier,
+    const std::string& endpoint_id, std::weak_ptr<CountDownLatch> barrier,
     std::function<ExceptionOr<bool>(EndpointChannel*)> handler) {
   // EndpointChannelManager will not let multiple channels exist simultaneously
   // for the same endpoint_id; it will be closing "old" channels as new ones
@@ -114,7 +114,13 @@ void EndpointManager::EndpointChannelLoopRunnable(
   // if needed.
   NEARBY_LOG(INFO, "Worker going down; name=%s; id=%s", runnable_name.c_str(),
              endpoint_id.c_str());
-  barrier->CountDown();
+  if (auto latch = barrier.lock()) {
+    latch->CountDown();
+  } else {
+    NEARBY_LOG(WARNING,
+               "Barrier already expired in worker name=%s, for endpoint %s",
+               runnable_name.c_str(), endpoint_id.c_str());
+  }
 
   // Always clear out all state related to this endpoint before terminating
   // this thread.
@@ -239,7 +245,14 @@ EndpointManager::~EndpointManager() {
       // This will close the channel; all workers will sense that and
       // terminate.
       channel_manager_->UnregisterChannelForEndpoint(endpoint_id);
-      state.barrier.Await();
+      if (state.barrier) {
+        state.barrier->Await();
+      } else {
+        NEARBY_LOG(
+            WARNING,
+            "State barrier already freed before EM destructor for endpoint %s",
+            endpoint_id.c_str());
+      }
     }
     latch.CountDown();
   });
@@ -331,7 +344,15 @@ void EndpointManager::EnsureWorkersTerminated(const std::string& endpoint_id) {
     EndpointState& endpoint_state = item->second;
     NEARBY_LOGS(INFO) << "Waiting for workers to terminate for id: "
                       << endpoint_id;
-    endpoint_state.barrier.Await();
+    if (endpoint_state.barrier) {
+      endpoint_state.barrier->Await();
+    } else {
+      NEARBY_LOG(
+          WARNING,
+          "State barrier already freed before EnsureWorkersTerminated for "
+          "endpoint %s",
+          endpoint_id.c_str());
+    }
     endpoints_.erase(item);
     NEARBY_LOGS(INFO) << "Workers terminated for id: " << endpoint_id;
   } else {
@@ -383,8 +404,12 @@ void EndpointManager::RegisterEndpoint(ClientProxy* client,
     // the next frame. If the handler fails its read and no other
     // EndpointChannels are available for this endpoint, a disconnection
     // will be initiated.
+    //
+    // Using weak_ptr just in case the barrier is freed, to save the UAF crash
+    // in b/179800119.
     StartEndpointReader(
-        [this, client, endpoint_id, barrier = &endpoint_state.barrier]() {
+        [this, client, endpoint_id,
+         barrier = std::weak_ptr<CountDownLatch>(endpoint_state.barrier)]() {
           EndpointChannelLoopRunnable(
               "Read", client, endpoint_id, barrier,
               [this, client, endpoint_id](EndpointChannel* channel) {
@@ -404,8 +429,12 @@ void EndpointManager::RegisterEndpoint(ClientProxy* client,
     // (**) Wifi Hotspots can fail to notice a connection has been lost, and
     // they will happily keep writing to /dev/null. This is why we listen
     // for the pong.
+    //
+    // Using weak_ptr just in case the barrier is freed, to save the UAF crash
+    // in b/179800119.
     StartEndpointKeepAliveManager([this, client, endpoint_id,
-                                   barrier = &endpoint_state.barrier]() {
+                                   barrier = std::weak_ptr<CountDownLatch>(
+                                       endpoint_state.barrier)]() {
       EndpointChannelLoopRunnable("KeepAliveManager", client, endpoint_id,
                                   barrier, [this](EndpointChannel* channel) {
                                     return HandleKeepAlive(channel);
diff --git a/cpp/core/internal/endpoint_manager.h b/cpp/core/internal/endpoint_manager.h
index 476b0c96..34bb39e3 100644
--- a/cpp/core/internal/endpoint_manager.h
+++ b/cpp/core/internal/endpoint_manager.h
@@ -149,7 +149,8 @@ class EndpointManager {
     ClientProxy* client;
     // Execution barrier, used to ensure that all workers associated with an
     // endpoint on handlers_executor_ and keep_alive_executor_ are terminated.
-    CountDownLatch barrier{2};
+    std::shared_ptr<CountDownLatch> barrier =
+        std::make_shared<CountDownLatch>(2);
   };
 
   FrameProcessor* GetFrameProcessor(V1Frame::FrameType frame_type);
@@ -169,7 +170,7 @@ class EndpointManager {
 
   void EndpointChannelLoopRunnable(
       const std::string& runnable_name, ClientProxy* client_proxy,
-      const std::string& endpoint_id, CountDownLatch* barrier,
+      const std::string& endpoint_id, std::weak_ptr<CountDownLatch> barrier,
       std::function<ExceptionOr<bool>(EndpointChannel*)> handler);
 
   static void WaitForLatch(const std::string& method_name,