Fix use-after-free in EndpointChannelManager

PiperOrigin-RevId: 945706442
2026-09-14 14:46:12 -04:00 · 2026-07-10 07:42:18 -07:00
parent 8eedf05f2c
commit 9e6f47f6f0
3 changed files with 358 additions and 214 deletions
@@ -17,7 +17,9 @@
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>

+#include "absl/strings/string_view.h"
 #include "absl/time/time.h"
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel.h"
@@ -48,7 +50,9 @@ void EndpointChannelManager::RegisterChannelForEndpoint(

  LOG(INFO) << "EndpointChannelManager registered channel of type "
            << channel->GetType() << " to endpoint " << endpoint_id;
-  SetActiveEndpointChannel(client, endpoint_id, std::move(channel),
+  std::shared_ptr<ChannelState::EndpointData> endpoint =
+      channel_state_.RegisterEndpoint(endpoint_id);
+  SetActiveEndpointChannel(client, endpoint_id, endpoint, std::move(channel),
                           true /* enable_encryption */);

  LOG(INFO) << "Registered channel: id=" << endpoint_id;
@@ -58,8 +62,17 @@ void EndpointChannelManager::ReplaceChannelForEndpoint(
    ClientProxy* client, const std::string& endpoint_id,
    std::shared_ptr<EndpointChannel> channel, bool enable_encryption) {
  MutexLock lock(&mutex_);
+
+  std::shared_ptr<ChannelState::EndpointData> endpoint =
+      channel_state_.GetEndpointData(endpoint_id);
+  if (endpoint == nullptr) {
+    LOG(WARNING) << "EndpointChannelManager failed to replace channel because "
+                    "endpoint "
+                 << endpoint_id << " is not registered.";
+    return;
+  }
  if (client->IsSafeToDisconnectEnabled(endpoint_id) &&
-      channel_state_.IsWaitingForSafeToDisconnectTimeout(endpoint_id)) {
+      endpoint->IsWaitingForSafeToDisconnectTimeout()) {
    LOG(WARNING)
        << "EndpointChannelManager failed to replace endpoint " << endpoint_id
        << "'s channel with type " << channel->GetType()
@@ -67,13 +80,7 @@ void EndpointChannelManager::ReplaceChannelForEndpoint(
    return;
  }

-  auto* endpoint = channel_state_.LookupEndpointData(endpoint_id);
-  if (endpoint != nullptr && endpoint->channel == nullptr) {
-    LOG(INFO) << "EndpointChannelManager is missing channel while "
-                 "trying to update: endpoint "
-              << endpoint_id;
-  }
-  SetActiveEndpointChannel(client, endpoint_id, std::move(channel),
+  SetActiveEndpointChannel(client, endpoint_id, endpoint, std::move(channel),
                           enable_encryption);
 }

@@ -82,38 +89,47 @@ bool EndpointChannelManager::EncryptChannelForEndpoint(
    std::unique_ptr<EncryptionContext> context) {
  MutexLock lock(&mutex_);

-  channel_state_.UpdateEncryptionContextForEndpoint(endpoint_id,
-                                                    std::move(context));
-  auto* endpoint = channel_state_.LookupEndpointData(endpoint_id);
-  return channel_state_.EncryptChannel(endpoint);
+  std::shared_ptr<ChannelState::EndpointData> endpoint =
+      channel_state_.GetEndpointData(endpoint_id);
+  if (endpoint == nullptr) {
+    LOG(WARNING) << "EncryptChannelForEndpoint failed "
+                 << "because endpoint is not registered: " << endpoint_id;
+    return false;
+  }
+  endpoint->set_context(std::move(context));
+  return endpoint->EncryptChannel();
 }

 std::shared_ptr<EndpointChannel> EndpointChannelManager::GetChannelForEndpoint(
-    const std::string& endpoint_id) {
+    absl::string_view endpoint_id) {
  MutexLock lock(&mutex_);

-  auto* endpoint = channel_state_.LookupEndpointData(endpoint_id);
+  std::shared_ptr<ChannelState::EndpointData> endpoint =
+      channel_state_.GetEndpointData(endpoint_id);
  if (endpoint == nullptr) {
    LOG(INFO) << "No channel info for endpoint " << endpoint_id;
    return {};
  }

-  return endpoint->channel;
+  return endpoint->channel();
 }

 void EndpointChannelManager::SetActiveEndpointChannel(
    ClientProxy* client, const std::string& endpoint_id,
+    std::shared_ptr<ChannelState::EndpointData> endpoint,
    std::shared_ptr<EndpointChannel> channel, bool enable_encryption) {
  // Update the channel first, then encrypt this new channel, if
  // crypto context is present.
  channel->SetAnalyticsRecorder(&client->GetAnalyticsRecorder(), endpoint_id);
  channel->SetLocalEndpointId(client->GetLocalEndpointId());
-  channel_state_.UpdateChannelForEndpoint(endpoint_id, std::move(channel));
-  channel_state_.UpdateSafeToDisconnectForEndpoint(
-      endpoint_id, client->IsSafeToDisconnectEnabled(endpoint_id));
-  auto* endpoint = channel_state_.LookupEndpointData(endpoint_id);
-  if (endpoint->IsEncrypted() && enable_encryption)
-    channel_state_.EncryptChannel(endpoint);
+
+  endpoint->set_channel(std::move(channel));
+  endpoint->set_safe_to_disconnect_enabled(
+      client->IsSafeToDisconnectEnabled(endpoint_id));
+
+  if (endpoint->IsEncrypted() && enable_encryption) {
+    endpoint->EncryptChannel();
+  }
 }

 int EndpointChannelManager::GetConnectedEndpointsCount() const {
@@ -134,7 +150,7 @@ void EndpointChannelManager::UpdateSafeToDisconnectForEndpoint(
 }

 void EndpointChannelManager::MarkEndpointStopWaitToDisconnect(
-    const std::string& endpoint_id, bool is_safe_to_disconnect,
+    absl::string_view endpoint_id, bool is_safe_to_disconnect,
    bool notify_stop_waiting) {
  MutexLock lock(&mutex_);
  channel_state_.MarkEndpointStopWaitToDisconnect(
@@ -142,61 +158,128 @@ void EndpointChannelManager::MarkEndpointStopWaitToDisconnect(
 }

 bool EndpointChannelManager::CreateNewTimeoutDisconnectedState(
-    const std::string& endpoint_id, absl::Duration timeout_millis) {
-  return channel_state_.CreateNewTimeoutDisconnectedState(endpoint_id,
-                                                          timeout_millis);
+    absl::string_view endpoint_id, absl::Duration timeout_millis) {
+  std::shared_ptr<ChannelState::EndpointData> endpoint_data;
+  {
+    MutexLock lock(&mutex_);
+    endpoint_data = channel_state_.GetEndpointData(endpoint_id);
+  }
+  if (!endpoint_data) return false;
+
+  LOG(INFO) << "[safe-to-disconnect] "
+               "Create TimeoutDisconnectedState for endpoint: "
+            << endpoint_id;
+  endpoint_data->CreateNewTimeoutDisconnectedState(timeout_millis);
+  return true;
 }

-bool EndpointChannelManager::IsSafeToDisconnect(
-    const std::string& endpoint_id) {
+bool EndpointChannelManager::IsSafeToDisconnect(absl::string_view endpoint_id) {
+  MutexLock lock(&mutex_);
  return channel_state_.IsSafeToDisconnect(endpoint_id);
 }
+
+bool EndpointChannelManager::IsWaitingForSafeToDisconnectTimeoutForTesting(
+    absl::string_view endpoint_id) {
+  MutexLock lock(&mutex_);
+  return channel_state_.IsWaitingForSafeToDisconnectTimeoutForTesting(
+      endpoint_id);
+}
+
 void EndpointChannelManager::RemoveTimeoutDisconnectedState(
-    const std::string& endpoint_id) {
+    absl::string_view endpoint_id) {
  MutexLock lock(&mutex_);
  channel_state_.RemoveTimeoutDisconnectedState(endpoint_id);
 }

 ///////////////////////////////// ChannelState /////////////////////////////////

-// endpoint - channel endpoint to encrypt
-bool EndpointChannelManager::ChannelState::EncryptChannel(
-    EndpointChannelManager::ChannelState::EndpointData* endpoint) {
-  if (endpoint != nullptr && endpoint->channel != nullptr &&
-      endpoint->context != nullptr) {
-    endpoint->channel->EnableEncryption(endpoint->context);
+void EndpointChannelManager::ChannelState::EndpointData::
+    CreateNewTimeoutDisconnectedState(absl::Duration timeout_millis) {
+  MutexLock lock(&timeout_to_disconnected_mutex_);
+  timeout_to_disconnected_enabled_ = true;
+  timeout_to_disconnected_notified_ = false;
+  timeout_to_disconnected_.Wait(timeout_millis);
+  LOG(INFO) << "[safe-to-disconnect] Wait is done with "
+            << (timeout_to_disconnected_notified_ ? "notification" : "timeout");
+  if (!timeout_to_disconnected_notified_) {
+    is_safe_to_disconnect_ = true;
+  }
+  timeout_to_disconnected_notified_ = false;
+  timeout_to_disconnected_enabled_ = false;
+}
+
+void EndpointChannelManager::ChannelState::EndpointData::
+    MarkEndpointStopWaitToDisconnect(bool is_safe_to_disconnect,
+                                     bool notify_stop_waiting) {
+  MutexLock lock(&timeout_to_disconnected_mutex_);
+  this->is_safe_to_disconnect_ = is_safe_to_disconnect;
+  if (!timeout_to_disconnected_enabled_) return;
+  if (notify_stop_waiting) {
+    LOG(INFO) << "[safe-to-disconnect] Notify stop waiting before timeout.";
+    timeout_to_disconnected_.Notify();
+    timeout_to_disconnected_notified_ = true;
+  }
+}
+
+bool EndpointChannelManager::ChannelState::EndpointData::
+    IsWaitingForSafeToDisconnectTimeout() const {
+  MutexLock lock(&timeout_to_disconnected_mutex_);
+  return timeout_to_disconnected_enabled_;
+}
+
+bool EndpointChannelManager::ChannelState::EndpointData::IsSafeToDisconnect()
+    const {
+  MutexLock lock(&timeout_to_disconnected_mutex_);
+  return is_safe_to_disconnect_;
+}
+
+void EndpointChannelManager::ChannelState::EndpointData::
+    RemoveTimeoutDisconnectedState() {
+  MutexLock lock(&timeout_to_disconnected_mutex_);
+  timeout_to_disconnected_notified_ = false;
+  timeout_to_disconnected_enabled_ = false;
+}
+
+bool EndpointChannelManager::ChannelState::EndpointData::EncryptChannel() {
+  if (context_ != nullptr) {
+    channel_->EnableEncryption(context_);
    return true;
  }
  return false;
 }

-EndpointChannelManager::ChannelState::EndpointData*
-EndpointChannelManager::ChannelState::LookupEndpointData(
-    const std::string& endpoint_id) {
-  auto item = endpoints_.find(endpoint_id);
-  return item != endpoints_.end() ? &item->second : nullptr;
+std::shared_ptr<EndpointChannelManager::ChannelState::EndpointData>
+EndpointChannelManager::ChannelState::GetEndpointData(
+    absl::string_view endpoint_id) {
+  auto it = endpoints_.find(endpoint_id);
+  return it != endpoints_.end() ? it->second : nullptr;
+}
+
+std::shared_ptr<EndpointChannelManager::ChannelState::EndpointData>
+EndpointChannelManager::ChannelState::RegisterEndpoint(
+    absl::string_view endpoint_id) {
+  std::shared_ptr<EndpointData>& endpoint = endpoints_[endpoint_id];
+  if (endpoint == nullptr) {
+    endpoint = std::make_shared<EndpointData>();
+  } else {
+    LOG(DFATAL) << "Endpoint " << endpoint_id
+                << " is already registered. It might not have been cleaned up "
+                   "properly.";
+  }
+  return endpoint;
 }

 void EndpointChannelManager::ChannelState::DestroyAll() {
-  for (auto& item : endpoints_) {
-    RemoveEndpoint(item.first, DisconnectionReason::SHUTDOWN,
-                   /* safe_to_disconnect_enabled */ false,
-                   SafeDisconnectionResult::kSafeDisconnection);
+  // Collect all endpoint IDs to avoid iterator invalidation.
+  std::vector<std::string> endpoint_ids;
+  endpoint_ids.reserve(endpoints_.size());
+  for (const auto& [endpoint_id, endpoint_data] : endpoints_) {
+    endpoint_ids.push_back(endpoint_id);
  }
-  endpoints_.clear();
-}

-void EndpointChannelManager::ChannelState::UpdateChannelForEndpoint(
-    const std::string& endpoint_id, std::shared_ptr<EndpointChannel> channel) {
-  // Create EndpointData instance, if necessary, and populate channel.
-  endpoints_[endpoint_id].channel = std::move(channel);
-}
-
-void EndpointChannelManager::ChannelState::UpdateEncryptionContextForEndpoint(
-    const std::string& endpoint_id,
-    std::unique_ptr<EncryptionContext> context) {
-  // Create EndpointData instance, if necessary, and populate crypto context.
-  endpoints_[endpoint_id].context = std::move(context);
+  for (const auto& endpoint_id : endpoint_ids) {
+    RemoveEndpoint(endpoint_id, DisconnectionReason::SHUTDOWN);
+  }
 }

 void EndpointChannelManager::ChannelState::UpdateSafeToDisconnectForEndpoint(
@@ -205,32 +288,28 @@ void EndpointChannelManager::ChannelState::UpdateSafeToDisconnectForEndpoint(
               "UpdateSafeToDisconnectForEndpoint for: "
            << endpoint_id << " " << safe_to_disconnect_enabled;

-  endpoints_[endpoint_id].safe_to_disconnect_enabled =
-      safe_to_disconnect_enabled;
-}
-
-bool EndpointChannelManager::ChannelState::GetSafeToDisconnectForEndpoint(
-    const std::string& endpoint_id) {
-  auto item = endpoints_.find(endpoint_id);
-  if (item == endpoints_.end()) return false;
-  LOG(INFO) << "[safe-to-disconnect] GetSafeToDisconnectForEndpoint: "
-            << item->second.safe_to_disconnect_enabled;
-  return item->second.safe_to_disconnect_enabled;
+  std::shared_ptr<EndpointData> endpoint = GetEndpointData(endpoint_id);
+  if (endpoint == nullptr) {
+    LOG(WARNING) << "UpdateSafeToDisconnectForEndpoint failed because endpoint "
+                 << endpoint_id << " is not registered.";
+    return;
+  }
+  endpoint->set_safe_to_disconnect_enabled(safe_to_disconnect_enabled);
 }

 bool EndpointChannelManager::ChannelState::RemoveEndpoint(
-    const std::string& endpoint_id, DisconnectionReason reason,
-    bool safe_to_disconnect_enabled, SafeDisconnectionResult result) {
-  auto item = endpoints_.find(endpoint_id);
-  if (item == endpoints_.end()) return false;
+    absl::string_view endpoint_id, DisconnectionReason reason) {
+  auto it = endpoints_.find(endpoint_id);
+  if (it == endpoints_.end()) return false;

  MarkEndpointStopWaitToDisconnect(endpoint_id,
                                   /* is_safe_to_disconnect */ true,
                                   /* notify_stop_waiting */ true);
-  item->second.disconnect_reason = reason;
-  auto channel = item->second.channel;
+  it->second->set_disconnect_reason(reason);
+  std::shared_ptr<EndpointChannel> channel = it->second->channel();
+  bool safe_to_disconnect_enabled = it->second->safe_to_disconnect_enabled();

-  if (channel && !channel->IsClosed() && !safe_to_disconnect_enabled) {
+  if (!channel->IsClosed() && !safe_to_disconnect_enabled) {
    // If the channel was paused (i.e. during a bandwidth upgrade negotiation)
    // we resume to ensure the thread won't hang when trying to write to it.
    channel->Resume();
@@ -247,18 +326,16 @@ bool EndpointChannelManager::ChannelState::RemoveEndpoint(
  }

  LOG(INFO) << "Remove Endpoint: " << endpoint_id;
-  endpoints_.erase(item);
+  endpoints_.erase(it);
  return true;
 }

 bool EndpointChannelManager::ChannelState::isWifiLanConnected() const {
-  for (auto& endpoint : endpoints_) {
-    auto channel = endpoint.second.channel;
-    if (channel) {
-      if (channel->GetMedium() == Medium::WIFI_LAN) {
-        LOG(INFO) << "Found WIFI_LAN Medium for endpoint:" << endpoint.first;
-        return true;
-      }
+  for (const auto& [endpoint_id, endpoint_data] : endpoints_) {
+    std::shared_ptr<EndpointChannel> channel = endpoint_data->channel();
+    if (channel->GetMedium() == Medium::WIFI_LAN) {
+      LOG(INFO) << "Found WIFI_LAN Medium for endpoint:" << endpoint_id;
+      return true;
    }
  }

@@ -266,96 +343,53 @@ bool EndpointChannelManager::ChannelState::isWifiLanConnected() const {
 }

 void EndpointChannelManager::ChannelState::MarkEndpointStopWaitToDisconnect(
-    const std::string& endpoint_id, bool is_safe_to_disconnect,
+    absl::string_view endpoint_id, bool is_safe_to_disconnect,
    bool notify_stop_waiting) {
-  auto item = endpoints_.find(endpoint_id);
-  if (item == endpoints_.end()) return;
+  std::shared_ptr<EndpointData> endpoint = GetEndpointData(endpoint_id);
+  if (endpoint == nullptr) return;
  LOG(INFO) << "[safe-to-disconnect] is_safe_to_disconnect= "
            << is_safe_to_disconnect
            << ", notify_stop_waiting= " << notify_stop_waiting
            << " for endpoint: " << endpoint_id;
-  {
-    MutexLock lock(&item->second.timeout_to_disconnected_mutex);
-    item->second.is_safe_to_disconnect = is_safe_to_disconnect;
-    if (!item->second.timeout_to_disconnected_enabled) return;
-    if (notify_stop_waiting) {
-      LOG(INFO) << "[safe-to-disconnect] Notify stop "
-                   "waiting before timeout.";
-      item->second.timeout_to_disconnected.Notify();
-      item->second.timeout_to_disconnected_notified = true;
-    }
-  }
+  endpoint->MarkEndpointStopWaitToDisconnect(is_safe_to_disconnect,
+                                             notify_stop_waiting);
 }

-bool EndpointChannelManager::ChannelState::CreateNewTimeoutDisconnectedState(
-    const std::string& endpoint_id, absl::Duration timeout_millis) {
-  auto item = endpoints_.find(endpoint_id);
-  if (item == endpoints_.end()) return false;
+bool EndpointChannelManager::ChannelState::
+    IsWaitingForSafeToDisconnectTimeoutForTesting(
+        absl::string_view endpoint_id) {
+  std::shared_ptr<EndpointData> endpoint = GetEndpointData(endpoint_id);
+  if (endpoint == nullptr) return false;
+  bool enabled = endpoint->IsWaitingForSafeToDisconnectTimeout();
  LOG(INFO) << "[safe-to-disconnect] "
-               "Create TimeoutDisconnectedState for endpoint: "
-            << endpoint_id;
-  {
-    MutexLock lock(&item->second.timeout_to_disconnected_mutex);
-    item->second.timeout_to_disconnected_enabled = true;
-    item->second.timeout_to_disconnected_notified = false;
-    item->second.timeout_to_disconnected.Wait(timeout_millis);
-    LOG(INFO) << "[safe-to-disconnect] Wait is done with "
-              << (item->second.timeout_to_disconnected_notified ? "notification"
-                                                                : "timeout");
-    if (!item->second.timeout_to_disconnected_notified)
-      item->second.is_safe_to_disconnect = true;
-    item->second.timeout_to_disconnected_notified = false;
-    item->second.timeout_to_disconnected_enabled = false;
-  }
-  return true;
-}
-bool EndpointChannelManager::ChannelState::IsWaitingForSafeToDisconnectTimeout(
-    const std::string& endpoint_id) {
-  auto item = endpoints_.find(endpoint_id);
-  if (item == endpoints_.end()) return false;
-  {
-    MutexLock lock(&item->second.timeout_to_disconnected_mutex);
-    LOG(INFO) << "[safe-to-disconnect] "
-                 "IsWaitingForSafeToDisconnectTimeout for endpoint: "
-              << endpoint_id << ": "
-              << item->second.timeout_to_disconnected_enabled;
-    return (item->second.timeout_to_disconnected_enabled);
-  }
+               "IsWaitingForSafeToDisconnectTimeout for endpoint: "
+            << endpoint_id << ": " << enabled;
+  return enabled;
 }

 bool EndpointChannelManager::ChannelState::IsSafeToDisconnect(
-    const std::string& endpoint_id) {
-  auto item = endpoints_.find(endpoint_id);
-  if (item == endpoints_.end()) return true;
-  {
-    MutexLock lock(&item->second.timeout_to_disconnected_mutex);
-    LOG(INFO)
-        << "[safe-to-disconnect] Get SafeToDisconnect status for endpoint: "
-        << endpoint_id << ": " << item->second.is_safe_to_disconnect;
-    return (item->second.is_safe_to_disconnect);
-  }
+    absl::string_view endpoint_id) {
+  std::shared_ptr<EndpointData> endpoint = GetEndpointData(endpoint_id);
+  if (endpoint == nullptr) return true;
+  bool is_safe = endpoint->IsSafeToDisconnect();
+  LOG(INFO) << "[safe-to-disconnect] Get SafeToDisconnect status for endpoint: "
+            << endpoint_id << ": " << is_safe;
+  return is_safe;
 }

 void EndpointChannelManager::ChannelState::RemoveTimeoutDisconnectedState(
-    const std::string& endpoint_id) {
-  auto item = endpoints_.find(endpoint_id);
-  if (item == endpoints_.end()) return;
-  {
-    MutexLock lock(&item->second.timeout_to_disconnected_mutex);
-    item->second.timeout_to_disconnected_notified = false;
-    item->second.timeout_to_disconnected_enabled = false;
-  }
+    absl::string_view endpoint_id) {
+  std::shared_ptr<EndpointData> endpoint = GetEndpointData(endpoint_id);
+  if (endpoint == nullptr) return;
+  endpoint->RemoveTimeoutDisconnectedState();
 }

 bool EndpointChannelManager::UnregisterChannelForEndpoint(
-    const std::string& endpoint_id, DisconnectionReason reason,
+    absl::string_view endpoint_id, DisconnectionReason reason,
    SafeDisconnectionResult result) {
  MutexLock lock(&mutex_);

-  auto safe_to_disconnect_enabled =
-      channel_state_.GetSafeToDisconnectForEndpoint(endpoint_id);
-  if (!channel_state_.RemoveEndpoint(endpoint_id, reason,
-                                     safe_to_disconnect_enabled, result)) {
+  if (!channel_state_.RemoveEndpoint(endpoint_id, reason)) {
    return false;
  }
  LOG(INFO) << "EndpointChannelManager unregistered channel for endpoint "
@@ -17,9 +17,11 @@

 #include <memory>
 #include <string>
+#include <utility>

 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/strings/string_view.h"
 #include "absl/time/time.h"
 #include "connections/implementation/analytics/analytics_recorder.h"
 #include "connections/implementation/client_proxy.h"
@@ -91,11 +93,11 @@ class EndpointChannelManager final {
  // EndpointManager methods that use a channel are running, it is better to
  // have a shared ownership.
  std::shared_ptr<EndpointChannel> GetChannelForEndpoint(
-      const std::string& endpoint_id) ABSL_LOCKS_EXCLUDED(mutex_);
+      absl::string_view endpoint_id) ABSL_LOCKS_EXCLUDED(mutex_);

  // Returns true if 'endpoint_id' actually had a registered EndpointChannel.
  // IOW, a return of false signifies a no-op.
-  bool UnregisterChannelForEndpoint(const std::string& endpoint_id,
+  bool UnregisterChannelForEndpoint(absl::string_view endpoint_id,
                                    DisconnectionReason reason,
                                    SafeDisconnectionResult result)
      ABSL_LOCKS_EXCLUDED(mutex_);
@@ -107,16 +109,19 @@ class EndpointChannelManager final {
  void UpdateSafeToDisconnectForEndpoint(const std::string& endpoint_id,
                                         bool safe_to_disconnect_enabled)
      ABSL_LOCKS_EXCLUDED(mutex_);
-  void MarkEndpointStopWaitToDisconnect(const std::string& endpoint_id,
+  bool CreateNewTimeoutDisconnectedState(absl::string_view endpoint_id,
+                                         absl::Duration timeout_millis)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+  void MarkEndpointStopWaitToDisconnect(absl::string_view endpoint_id,
                                        bool is_safe_to_disconnect,
                                        bool notify_stop_waiting)
      ABSL_LOCKS_EXCLUDED(mutex_);
-  bool CreateNewTimeoutDisconnectedState(const std::string& endpoint_id,
-                                         absl::Duration timeout_millis)
+
+  bool IsSafeToDisconnect(absl::string_view endpoint_id)
      ABSL_LOCKS_EXCLUDED(mutex_);
-  bool IsSafeToDisconnect(const std::string& endpoint_id)
-      ABSL_LOCKS_EXCLUDED(mutex_);
-  void RemoveTimeoutDisconnectedState(const std::string& endpoint_id)
+  bool IsWaitingForSafeToDisconnectTimeoutForTesting(
+      absl::string_view endpoint_id) ABSL_LOCKS_EXCLUDED(mutex_);
+  void RemoveTimeoutDisconnectedState(absl::string_view endpoint_id)
      ABSL_LOCKS_EXCLUDED(mutex_);

 private:
@@ -125,32 +130,69 @@ class EndpointChannelManager final {
  // been encrypted yet.
  class ChannelState {
   public:
-    struct EndpointData {
+    class EndpointData {
+     public:
      EndpointData() = default;
      EndpointData(EndpointData&&) = default;
      EndpointData& operator=(EndpointData&&) = default;
      ~EndpointData() {
-        if (channel != nullptr) {
-          channel->Close(disconnect_reason);
+        if (channel_ != nullptr) {
+          channel_->Close(disconnect_reason_);
        }
      }

      // True if we have a 'context' for the endpoint.
-      bool IsEncrypted() const { return context != nullptr; }
+      bool IsEncrypted() const { return context_ != nullptr; }

-      std::shared_ptr<EndpointChannel> channel;
-      std::shared_ptr<EncryptionContext> context;
-      DisconnectionReason disconnect_reason =
+      void CreateNewTimeoutDisconnectedState(absl::Duration timeout_millis)
+          ABSL_LOCKS_EXCLUDED(timeout_to_disconnected_mutex_);
+      void MarkEndpointStopWaitToDisconnect(bool is_safe_to_disconnect,
+                                            bool notify_stop_waiting)
+          ABSL_LOCKS_EXCLUDED(timeout_to_disconnected_mutex_);
+      bool IsWaitingForSafeToDisconnectTimeout() const
+          ABSL_LOCKS_EXCLUDED(timeout_to_disconnected_mutex_);
+      bool IsSafeToDisconnect() const
+          ABSL_LOCKS_EXCLUDED(timeout_to_disconnected_mutex_);
+      void RemoveTimeoutDisconnectedState()
+          ABSL_LOCKS_EXCLUDED(timeout_to_disconnected_mutex_);
+      bool EncryptChannel();
+
+      std::shared_ptr<EndpointChannel> channel() const { return channel_; }
+      void set_channel(std::shared_ptr<EndpointChannel> channel) {
+        channel_ = std::move(channel);
+      }
+
+      std::shared_ptr<EncryptionContext> context() const { return context_; }
+      void set_context(std::shared_ptr<EncryptionContext> context) {
+        context_ = std::move(context);
+      }
+
+      void set_disconnect_reason(DisconnectionReason disconnect_reason) {
+        disconnect_reason_ = disconnect_reason;
+      }
+
+      bool safe_to_disconnect_enabled() const {
+        return safe_to_disconnect_enabled_;
+      }
+      void set_safe_to_disconnect_enabled(bool safe_to_disconnect_enabled) {
+        safe_to_disconnect_enabled_ = safe_to_disconnect_enabled;
+      }
+
+     private:
+      std::shared_ptr<EndpointChannel> channel_;
+      std::shared_ptr<EncryptionContext> context_;
+      DisconnectionReason disconnect_reason_ =
          DisconnectionReason::UNKNOWN_DISCONNECTION_REASON;
-      bool safe_to_disconnect_enabled = false;
-      mutable Mutex timeout_to_disconnected_mutex;
-      ConditionVariable timeout_to_disconnected{&timeout_to_disconnected_mutex};
-      bool timeout_to_disconnected_enabled
-          ABSL_GUARDED_BY(timeout_to_disconnected_mutex) = false;
-      bool timeout_to_disconnected_notified
-          ABSL_GUARDED_BY(timeout_to_disconnected_mutex) = false;
-      bool is_safe_to_disconnect
-          ABSL_GUARDED_BY(timeout_to_disconnected_mutex) = false;
+      bool safe_to_disconnect_enabled_ = false;
+      mutable Mutex timeout_to_disconnected_mutex_;
+      ConditionVariable timeout_to_disconnected_{
+          &timeout_to_disconnected_mutex_};
+      bool timeout_to_disconnected_enabled_
+          ABSL_GUARDED_BY(timeout_to_disconnected_mutex_) = false;
+      bool timeout_to_disconnected_notified_
+          ABSL_GUARDED_BY(timeout_to_disconnected_mutex_) = false;
+      bool is_safe_to_disconnect_
+          ABSL_GUARDED_BY(timeout_to_disconnected_mutex_) = false;
    };

    ChannelState() = default;
@@ -160,57 +202,48 @@ class EndpointChannelManager final {

    // Provides a way to destroy contents of a container, while holding a lock.
    void DestroyAll();
+
    // Return pointer to endpoint data, or nullptr, it not found.
-    EndpointData* LookupEndpointData(const std::string& endpoint_id);
+    std::shared_ptr<EndpointData> GetEndpointData(
+        absl::string_view endpoint_id);

-    // Stores a new EndpointChannel for the endpoint.
-    // Prevoius one is destroyed, if it existed.
-    void UpdateChannelForEndpoint(const std::string& endpoint_id,
-                                  std::shared_ptr<EndpointChannel> channel);
-
-    // Stores a new EncryptionContext for the endpoint.
-    // Prevoius one is destroyed, if it existed.
-    void UpdateEncryptionContextForEndpoint(
-        const std::string& endpoint_id,
-        std::unique_ptr<EncryptionContext> context);
+    // Registers a new endpoint id. This is the only spot EndpointData is
+    // created.
+    std::shared_ptr<EndpointData> RegisterEndpoint(
+        absl::string_view endpoint_id);

    void UpdateSafeToDisconnectForEndpoint(const std::string& endpoint_id,
                                           bool safe_to_disconnect_enabled);
-    bool GetSafeToDisconnectForEndpoint(const std::string& endpoint_id);

    // Removes all knowledge of this endpoint, cleaning up as necessary.
    // Returns false if the endpoint was not found.
-    bool RemoveEndpoint(const std::string& endpoint_id,
-                        DisconnectionReason reason,
-                        bool safe_to_disconnect_enabled,
-                        SafeDisconnectionResult result);
+    bool RemoveEndpoint(absl::string_view endpoint_id,
+                        DisconnectionReason reason);

-    bool EncryptChannel(EndpointData* endpoint);
    int GetConnectedEndpointsCount() const { return endpoints_.size(); }
    bool isWifiLanConnected() const;
-    void MarkEndpointStopWaitToDisconnect(const std::string& endpoint_id,
+    void MarkEndpointStopWaitToDisconnect(absl::string_view endpoint_id,
                                          bool is_safe_to_disconnect,
                                          bool notify_stop_waiting);
-    bool CreateNewTimeoutDisconnectedState(const std::string& endpoint_id,
-                                           absl::Duration timeout_millis);
-    bool IsWaitingForSafeToDisconnectTimeout(const std::string& endpoint_id);
-    bool IsSafeToDisconnect(const std::string& endpoint_id);
-    void RemoveTimeoutDisconnectedState(const std::string& endpoint_id);
+    bool IsWaitingForSafeToDisconnectTimeoutForTesting(
+        absl::string_view endpoint_id);
+    bool IsSafeToDisconnect(absl::string_view endpoint_id);
+    void RemoveTimeoutDisconnectedState(absl::string_view endpoint_id);

   private:
    // Endpoint ID -> EndpointData. Contains everything we know about the
    // endpoint.
-    absl::flat_hash_map<std::string, EndpointData> endpoints_;
+    absl::flat_hash_map<std::string, std::shared_ptr<EndpointData>> endpoints_;
  };

-  void SetActiveEndpointChannel(ClientProxy* client,
-                                const std::string& endpoint_id,
-                                std::shared_ptr<EndpointChannel> channel,
-                                bool enable_encryption)
+  void SetActiveEndpointChannel(
+      ClientProxy* client, const std::string& endpoint_id,
+      std::shared_ptr<ChannelState::EndpointData> endpoint,
+      std::shared_ptr<EndpointChannel> channel, bool enable_encryption)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  mutable Mutex mutex_;
-  ChannelState channel_state_;
+  ChannelState channel_state_ ABSL_GUARDED_BY(mutex_);
 };

 }  // namespace nearby::connections
@@ -26,6 +26,7 @@
 #include "gtest/gtest.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
+#include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "connections/implementation/analytics/analytics_recorder.h"
 #include "connections/implementation/base_endpoint_channel.h"
@@ -210,16 +211,16 @@ TEST(BaseEndpointChannelManagerTest, RegisterChannelEncryptedReadwrite) {
  ASSERT_NE(context.second, nullptr);

  EndpointChannelManager ecm_a;
-  ecm_a.EncryptChannelForEndpoint(std::string(kEndpointId),
-                                  std::move(context.first));
  ecm_a.RegisterChannelForEndpoint(&proxy_a, std::string(kEndpointId),
                                   std::move(channel_a));
+  ecm_a.EncryptChannelForEndpoint(std::string(kEndpointId),
+                                  std::move(context.first));

  EndpointChannelManager ecm_b;
-  ecm_b.EncryptChannelForEndpoint(std::string(kEndpointId),
-                                  std::move(context.second));
  ecm_b.RegisterChannelForEndpoint(&proxy_b, std::string(kEndpointId),
                                   std::move(channel_b));
+  ecm_b.EncryptChannelForEndpoint(std::string(kEndpointId),
+                                  std::move(context.second));

  EXPECT_EQ(channel_a_raw->GetType(), "ENCRYPTED_BLUETOOTH");
  EXPECT_EQ(channel_b_raw->GetType(), "ENCRYPTED_BLUETOOTH");
@@ -241,10 +242,10 @@ TEST(BaseEndpointChannelManagerTest, RegisterChannelEncryptedReadwrite) {
  channel_a_raw->Close(DisconnectionReason::LOCAL_DISCONNECTION);
  channel_b_raw->Close(DisconnectionReason::REMOTE_DISCONNECTION);
  ecm_a.UnregisterChannelForEndpoint(
-      std::string(kEndpointId), DisconnectionReason::LOCAL_DISCONNECTION,
+      kEndpointId, DisconnectionReason::LOCAL_DISCONNECTION,
      SafeDisconnectionResult::kSafeDisconnection);
  ecm_b.UnregisterChannelForEndpoint(
-      std::string(kEndpointId), DisconnectionReason::REMOTE_DISCONNECTION,
+      kEndpointId, DisconnectionReason::REMOTE_DISCONNECTION,
      SafeDisconnectionResult::kSafeDisconnection);
 }

@@ -290,13 +291,26 @@ TEST(BaseEndpointChannelManagerTest, ReplaceChannelNoEncrypted) {
  ASSERT_NE(context.first, nullptr);
  ASSERT_NE(context.second, nullptr);

+  auto client_a_dummy = CreatePipe();
+  auto server_a_dummy = CreatePipe();
+  auto channel_a_init = std::make_shared<MockEndpointChannel>(
+      server_a_dummy.first.get(), client_a_dummy.second.get());
+  auto client_b_dummy = CreatePipe();
+  auto server_b_dummy = CreatePipe();
+  auto channel_b_init = std::make_shared<MockEndpointChannel>(
+      server_b_dummy.first.get(), client_b_dummy.second.get());
+
  EndpointChannelManager ecm_a;
+  ecm_a.RegisterChannelForEndpoint(&proxy_a, std::string(kEndpointId),
+                                   std::move(channel_a_init));
  ecm_a.EncryptChannelForEndpoint(std::string(kEndpointId),
                                  std::move(context.first));
  ecm_a.ReplaceChannelForEndpoint(&proxy_a, std::string(kEndpointId),
                                  std::move(channel_a), false);

  EndpointChannelManager ecm_b;
+  ecm_b.RegisterChannelForEndpoint(&proxy_b, std::string(kEndpointId),
+                                   std::move(channel_b_init));
  ecm_b.EncryptChannelForEndpoint(std::string(kEndpointId),
                                  std::move(context.second));
  ecm_b.ReplaceChannelForEndpoint(&proxy_b, std::string(kEndpointId),
@@ -309,12 +323,75 @@ TEST(BaseEndpointChannelManagerTest, ReplaceChannelNoEncrypted) {
  channel_a_raw->Close(DisconnectionReason::LOCAL_DISCONNECTION);
  channel_b_raw->Close(DisconnectionReason::REMOTE_DISCONNECTION);
  ecm_a.UnregisterChannelForEndpoint(
-      std::string(kEndpointId), DisconnectionReason::LOCAL_DISCONNECTION,
+      kEndpointId, DisconnectionReason::LOCAL_DISCONNECTION,
      SafeDisconnectionResult::kSafeDisconnection);
  ecm_b.UnregisterChannelForEndpoint(
-      std::string(kEndpointId), DisconnectionReason::REMOTE_DISCONNECTION,
+      kEndpointId, DisconnectionReason::REMOTE_DISCONNECTION,
      SafeDisconnectionResult::kSafeDisconnection);
 }

+TEST(BaseEndpointChannelManagerTest,
+     CreateNewTimeoutDisconnectedStateUnregisterDuringWait) {
+  ClientProxy proxy;
+  EndpointChannelManager ecm;
+  auto client = CreatePipe();
+  auto server = CreatePipe();
+  auto channel = std::make_shared<MockEndpointChannel>(server.first.get(),
+                                                       client.second.get());
+  auto channel_raw = channel.get();
+
+  ON_CALL(*channel_raw, GetMedium).WillByDefault([]() {
+    return Medium::BLUETOOTH;
+  });
+
+  ecm.RegisterChannelForEndpoint(&proxy, std::string(kEndpointId),
+                                 std::move(channel));
+
+  EXPECT_EQ(ecm.GetConnectedEndpointsCount(), 1);
+
+  MultiThreadExecutor executor(1);
+  CountDownLatch start_latch(1);
+  CountDownLatch finish_latch(1);
+  bool wait_result = false;
+
+  executor.Execute([&]() {
+    start_latch.CountDown();
+    wait_result =
+        ecm.CreateNewTimeoutDisconnectedState(kEndpointId, absl::Seconds(5));
+    finish_latch.CountDown();
+  });
+
+  ASSERT_TRUE(start_latch.Await(absl::Seconds(1)).result());
+
+  // Wait for the endpoint to enter the waiting state.
+  absl::Time deadline = absl::Now() + absl::Seconds(1);
+  while (!ecm.IsWaitingForSafeToDisconnectTimeoutForTesting(kEndpointId)) {
+    ASSERT_TRUE(absl::Now() < deadline)
+        << "Timed out waiting for endpoint to enter wait state.";
+    absl::SleepFor(absl::Milliseconds(10));
+  }
+
+  // Close the channel first to prevent UnregisterChannelForEndpoint from
+  // attempting to write disconnection frames to it, bypassing the 500ms data
+  // transfer delay and potential segfaults.
+  channel_raw->Close(DisconnectionReason::LOCAL_DISCONNECTION);
+
+  bool unregister_result = ecm.UnregisterChannelForEndpoint(
+      kEndpointId, DisconnectionReason::LOCAL_DISCONNECTION,
+      SafeDisconnectionResult::kSafeDisconnection);
+  EXPECT_TRUE(unregister_result);
+
+  EXPECT_TRUE(finish_latch.Await(absl::Seconds(2)).result());
+  EXPECT_TRUE(wait_result);
+  EXPECT_EQ(ecm.GetConnectedEndpointsCount(), 0);
+}
+
+TEST(BaseEndpointChannelManagerTest,
+     CreateNewTimeoutDisconnectedStateReturnsFalseForNonexistentEndpoint) {
+  EndpointChannelManager ecm;
+  EXPECT_FALSE(ecm.CreateNewTimeoutDisconnectedState("NonexistentEndpoint",
+                                                     absl::Seconds(1)));
+}
+
 }  // namespace
 }  // namespace nearby::connections