diff --git a/connections/implementation/endpoint_channel_manager.cc b/connections/implementation/endpoint_channel_manager.cc index 51090c57..099aab27 100644 --- a/connections/implementation/endpoint_channel_manager.cc +++ b/connections/implementation/endpoint_channel_manager.cc @@ -17,7 +17,9 @@ #include #include #include +#include +#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 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 channel, bool enable_encryption) { MutexLock lock(&mutex_); + + std::shared_ptr 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 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 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 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 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 endpoint, std::shared_ptr 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 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::GetEndpointData( + absl::string_view endpoint_id) { + auto it = endpoints_.find(endpoint_id); + return it != endpoints_.end() ? it->second : nullptr; +} + +std::shared_ptr +EndpointChannelManager::ChannelState::RegisterEndpoint( + absl::string_view endpoint_id) { + std::shared_ptr& endpoint = endpoints_[endpoint_id]; + if (endpoint == nullptr) { + endpoint = std::make_shared(); + } 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 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 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 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 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 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 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 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 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 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 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 " diff --git a/connections/implementation/endpoint_channel_manager.h b/connections/implementation/endpoint_channel_manager.h index 47878023..eeb87df6 100644 --- a/connections/implementation/endpoint_channel_manager.h +++ b/connections/implementation/endpoint_channel_manager.h @@ -17,9 +17,11 @@ #include #include +#include #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 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 channel; - std::shared_ptr 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 channel() const { return channel_; } + void set_channel(std::shared_ptr channel) { + channel_ = std::move(channel); + } + + std::shared_ptr context() const { return context_; } + void set_context(std::shared_ptr 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 channel_; + std::shared_ptr 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 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 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 context); + // Registers a new endpoint id. This is the only spot EndpointData is + // created. + std::shared_ptr 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 endpoints_; + absl::flat_hash_map> endpoints_; }; - void SetActiveEndpointChannel(ClientProxy* client, - const std::string& endpoint_id, - std::shared_ptr channel, - bool enable_encryption) + void SetActiveEndpointChannel( + ClientProxy* client, const std::string& endpoint_id, + std::shared_ptr endpoint, + std::shared_ptr 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 diff --git a/connections/implementation/endpoint_channel_manager_test.cc b/connections/implementation/endpoint_channel_manager_test.cc index 0bdbb6d1..546fea04 100644 --- a/connections/implementation/endpoint_channel_manager_test.cc +++ b/connections/implementation/endpoint_channel_manager_test.cc @@ -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( + 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( + 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(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