From f38728bb0b540c4397f297b32ff9e52d1c14503d Mon Sep 17 00:00:00 2001
From: Juliet Levesque <julietlevesque@google.com>
Date: Tue, 30 May 2023 12:02:48 -0700
Subject: [PATCH] [Nearby Connections] Do not erase CancellationFlag map entry
 in ClientProxy

To prevent invalid access of already destroyed CancellationFlag pointers, do not erase the map entry of a CancellationFlag in ClientProxy::CancelEndpoint. Crashes were caused by consumers accessing CancellationFlag::Cancelled to check if the flags were cancelled, however since the flags was destroyed by erasing the map entry in ClientProxy::CancelEndpoint, consumers are calling gabrage memory.

PiperOrigin-RevId: 536470187
---
 connections/implementation/client_proxy.cc      |  7 ++++---
 connections/implementation/client_proxy.h       |  6 ++++++
 connections/implementation/client_proxy_test.cc | 10 ++++++++++
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/connections/implementation/client_proxy.cc b/connections/implementation/client_proxy.cc
index 682426e7..2d3ecf10 100644
--- a/connections/implementation/client_proxy.cc
+++ b/connections/implementation/client_proxy.cc
@@ -638,9 +638,9 @@ CancellationFlag* ClientProxy::GetCancellationFlag(
 
 void ClientProxy::CancelEndpoint(const std::string& endpoint_id) {
   const auto item = cancellation_flags_.find(endpoint_id);
-  if (item == cancellation_flags_.end()) return;
-  item->second->Cancel();
-  cancellation_flags_.erase(item);
+  if (item != cancellation_flags_.end()) {
+    item->second->Cancel();
+  }
 }
 
 const OsInfo& ClientProxy::GetLocalOsInfo() const {
@@ -735,6 +735,7 @@ void ClientProxy::RemoveAllEndpoints() {
   // just remove without notifying.
   connections_.clear();
   cancellation_flags_.clear();
+
   OnSessionComplete();
 }
 
diff --git a/connections/implementation/client_proxy.h b/connections/implementation/client_proxy.h
index cfd60369..e3b32492 100644
--- a/connections/implementation/client_proxy.h
+++ b/connections/implementation/client_proxy.h
@@ -258,7 +258,13 @@ class ClientProxy final {
     bool IsEmpty() const { return service_id.empty(); }
   };
 
+  // `RemoveAllEndpoints` is expected to only be called during destruction of
+  // ClientProxy via `ClientProxy::Reset`, which makes destroying
+  // CancellationFlags safe here since we are destroying ClientProxy. Do not
+  // use CancellationFlags after `RemoveAllEndpoints` is called, since all
+  // flags now are referencing garbage memory.
   void RemoveAllEndpoints();
+
   void OnSessionComplete();
   bool ConnectionStatusesContains(const std::string& endpoint_id,
                                   Connection::Status status_to_match) const;
diff --git a/connections/implementation/client_proxy_test.cc b/connections/implementation/client_proxy_test.cc
index 221060c9..0abedeb4 100644
--- a/connections/implementation/client_proxy_test.cc
+++ b/connections/implementation/client_proxy_test.cc
@@ -303,6 +303,7 @@ class ClientProxyTest : public ::testing::TestWithParam<FeatureFlags::Flags> {
   DiscoveryOptions discovery_options_;
 };
 
+// Regression test for b/279962714.
 TEST_P(ClientProxyTest, CanCancelEndpoint) {
   FeatureFlags::Flags feature_flags = GetParam();
   MediumEnvironment::Instance().SetFeatureFlags(feature_flags);
@@ -313,8 +314,15 @@ TEST_P(ClientProxyTest, CanCancelEndpoint) {
   OnDiscoveryEndpointFound(&client2_, advertising_endpoint);
   OnDiscoveryConnectionInitiated(&client2_, advertising_endpoint);
 
+  // `CancellationFlag` pointers are passed to other classes in Nearby
+  // Connections, and by using the pointers directly, we test their
+  // consumption of `CancellationFlag` pointers.
+  CancellationFlag* cancellation_flag =
+      client2_.GetCancellationFlag(advertising_endpoint.id);
+
   EXPECT_FALSE(
       client2_.GetCancellationFlag(advertising_endpoint.id)->Cancelled());
+  EXPECT_FALSE(cancellation_flag->Cancelled());
 
   client2_.CancelEndpoint(advertising_endpoint.id);
 
@@ -322,10 +330,12 @@ TEST_P(ClientProxyTest, CanCancelEndpoint) {
   if (!feature_flags.enable_cancellation_flag) {
     EXPECT_FALSE(
         client2_.GetCancellationFlag(advertising_endpoint.id)->Cancelled());
+    EXPECT_FALSE(cancellation_flag->Cancelled());
   } else {
     // The Cancelled is always true as the default flag being returned.
     EXPECT_TRUE(
         client2_.GetCancellationFlag(advertising_endpoint.id)->Cancelled());
+    EXPECT_TRUE(cancellation_flag->Cancelled());
   }
 }