From 4ae41f4cdfd1d1e5d22a9cfa2d9c887d0c697625 Mon Sep 17 00:00:00 2001 From: Edwin Wu Date: Tue, 28 Apr 2026 18:13:28 -0700 Subject: [PATCH] Refactor EndpointChannel ownership to use shared_ptr. PiperOrigin-RevId: 907278448 --- .../base_endpoint_channel_test.cc | 130 ++++----- .../implementation/base_pcp_handler.cc | 249 ++++++++++++------ connections/implementation/base_pcp_handler.h | 25 +- .../implementation/base_pcp_handler_test.cc | 104 ++++---- .../connections_authentication_transport.cc | 11 +- .../connections_authentication_transport.h | 6 +- ...nnections_authentication_transport_test.cc | 35 +-- .../implementation/encryption_runner.cc | 103 +++++--- .../implementation/encryption_runner.h | 19 +- .../implementation/encryption_runner_test.cc | 91 +++---- .../endpoint_channel_manager.cc | 8 +- .../implementation/endpoint_channel_manager.h | 8 +- .../endpoint_channel_manager_test.cc | 22 +- .../implementation/endpoint_manager.cc | 168 ++++++------ connections/implementation/endpoint_manager.h | 5 +- 15 files changed, 546 insertions(+), 438 deletions(-) diff --git a/connections/implementation/base_endpoint_channel_test.cc b/connections/implementation/base_endpoint_channel_test.cc index cb05c1ff..1e3cbbde 100644 --- a/connections/implementation/base_endpoint_channel_test.cc +++ b/connections/implementation/base_endpoint_channel_test.cc @@ -108,20 +108,26 @@ std::function MakeDataMonitor(const std::string& label, std::pair, std::shared_ptr> -DoDhKeyExchange(BaseEndpointChannel* channel_a, - BaseEndpointChannel* channel_b) { +DoDhKeyExchange(std::shared_ptr channel_a, + std::shared_ptr channel_b) { std::shared_ptr context_a; std::shared_ptr context_b; EncryptionRunner crypto_a; EncryptionRunner crypto_b; ClientProxy proxy_a; ClientProxy proxy_b; - CountDownLatch latch(2); + std::shared_ptr shared_channel_a = channel_a; + std::shared_ptr shared_channel_b = channel_b; + + // Create a shared_ptr for the latch to prevent Use-After-Free if the + // negotiation times out and this function returns early. + auto latch = std::make_shared(2); + crypto_a.StartClient( - &proxy_a, "endpoint_id", channel_a, + &proxy_a, "endpoint_id", shared_channel_a, { .on_success_cb = - [&latch, &context_a]( + [latch, &context_a]( const std::string& endpoint_id, std::unique_ptr ukey2, const std::string& auth_token, @@ -131,20 +137,19 @@ DoDhKeyExchange(BaseEndpointChannel* channel_a, auto context = ukey2->ToConnectionContext(); EXPECT_NE(context, nullptr); context_a = std::move(context); - latch.CountDown(); + latch->CountDown(); }, .on_failure_cb = - [&latch](const std::string& endpoint_id, - EndpointChannel* channel) { + [latch](const std::string& endpoint_id) { LOG(INFO) << "client-A side key negotiation failed"; - latch.CountDown(); + latch->CountDown(); }, }); crypto_b.StartServer( - &proxy_b, "endpoint_id", channel_b, + &proxy_b, "endpoint_id", shared_channel_b, { .on_success_cb = - [&latch, &context_b]( + [latch, &context_b]( const std::string& endpoint_id, std::unique_ptr ukey2, const std::string& auth_token, @@ -154,16 +159,15 @@ DoDhKeyExchange(BaseEndpointChannel* channel_a, auto context = ukey2->ToConnectionContext(); EXPECT_NE(context, nullptr); context_b = std::move(context); - latch.CountDown(); + latch->CountDown(); }, .on_failure_cb = - [&latch](const std::string& endpoint_id, - EndpointChannel* channel) { + [latch](const std::string& endpoint_id) { LOG(INFO) << "client-B side key negotiation failed"; - latch.CountDown(); + latch->CountDown(); }, }); - EXPECT_TRUE(latch.Await(absl::Milliseconds(5000)).result()); + EXPECT_TRUE(latch->Await(absl::Milliseconds(5000)).result()); return std::make_pair(std::move(context_a), std::move(context_b)); } @@ -265,20 +269,22 @@ TEST_F(BaseEndpointChannelTest, TryDecrypt) { absl::string_view kMessage = "message"; auto pipe_a = CreatePipe(); // channel_a writes to pipe_a, reads from pipe_b. auto pipe_b = CreatePipe(); // channel_b writes to pipe_b, reads from pipe_a. - TestEndpointChannel channel_a(pipe_b.first.get(), pipe_a.second.get()); - TestEndpointChannel channel_b(pipe_a.first.get(), pipe_b.second.get()); - auto [context_a, context_b] = DoDhKeyExchange(&channel_a, &channel_b); + auto channel_a = std::make_shared(pipe_b.first.get(), + pipe_a.second.get()); + auto channel_b = std::make_shared(pipe_a.first.get(), + pipe_b.second.get()); + auto [context_a, context_b] = DoDhKeyExchange(channel_a, channel_b); ASSERT_NE(context_a, nullptr); ASSERT_NE(context_b, nullptr); - channel_a.EnableEncryption(context_a); - channel_b.EnableEncryption(context_b); + channel_a->EnableEncryption(context_a); + channel_b->EnableEncryption(context_b); std::unique_ptr encrypted_message = - channel_a.EncodeMessageForTests(kMessage); + channel_a->EncodeMessageForTests(kMessage); ExceptionOr decrypted_message = - channel_b.TryDecrypt(ByteArray(*encrypted_message)); + channel_b->TryDecrypt(ByteArray(*encrypted_message)); - EXPECT_TRUE(channel_b.IsEncrypted()); + EXPECT_TRUE(channel_b->IsEncrypted()); EXPECT_TRUE(decrypted_message.ok()); EXPECT_EQ(decrypted_message.result().AsStringView(), kMessage); } @@ -286,16 +292,18 @@ TEST_F(BaseEndpointChannelTest, TryDecrypt) { TEST_F(BaseEndpointChannelTest, TryDecryptFailsWhenDecryptionFails) { auto pipe_a = CreatePipe(); // channel_a writes to pipe_a, reads from pipe_b. auto pipe_b = CreatePipe(); // channel_b writes to pipe_b, reads from pipe_a. - TestEndpointChannel channel_a(pipe_b.first.get(), pipe_a.second.get()); - TestEndpointChannel channel_b(pipe_a.first.get(), pipe_b.second.get()); - auto [context_a, context_b] = DoDhKeyExchange(&channel_a, &channel_b); + auto channel_a = std::make_shared(pipe_b.first.get(), + pipe_a.second.get()); + auto channel_b = std::make_shared(pipe_a.first.get(), + pipe_b.second.get()); + auto [context_a, context_b] = DoDhKeyExchange(channel_a, channel_b); ASSERT_NE(context_a, nullptr); - channel_a.EnableEncryption(context_a); + channel_a->EnableEncryption(context_a); ExceptionOr result = - channel_a.TryDecrypt(ByteArray("invalid message")); + channel_a->TryDecrypt(ByteArray("invalid message")); - EXPECT_TRUE(channel_a.IsEncrypted()); + EXPECT_TRUE(channel_a->IsEncrypted()); EXPECT_FALSE(result.ok()); EXPECT_EQ(result.exception(), Exception::kExecution); } @@ -366,13 +374,15 @@ TEST_F(BaseEndpointChannelTest, EncryptedReadWriteCanNotBeIntercepted) { // to server "b". auto server_b = CreatePipe(); // Data pump "b" reads from client "b", writes // to server "a". - TestEndpointChannel channel_a(server_a.first.get(), client_a.second.get()); - TestEndpointChannel channel_b(server_b.first.get(), client_b.second.get()); + auto channel_a = std::make_shared(server_a.first.get(), + client_a.second.get()); + auto channel_b = std::make_shared(server_b.first.get(), + client_b.second.get()); - ON_CALL(channel_a, GetMedium).WillByDefault([]() { + ON_CALL(*channel_a, GetMedium).WillByDefault([]() { return Medium::BLUETOOTH; }); - ON_CALL(channel_b, GetMedium).WillByDefault([]() { + ON_CALL(*channel_b, GetMedium).WillByDefault([]() { return Medium::BLUETOOTH; }); @@ -385,21 +395,21 @@ TEST_F(BaseEndpointChannelTest, EncryptedReadWriteCanNotBeIntercepted) { MakeDataMonitor("monitor_b", &capture_b, &mutex))); // Run DH key exchange; setup encryption contexts for channels. - auto [context_a, context_b] = DoDhKeyExchange(&channel_a, &channel_b); + auto [context_a, context_b] = DoDhKeyExchange(channel_a, channel_b); ASSERT_NE(context_a, nullptr); ASSERT_NE(context_b, nullptr); - channel_a.EnableEncryption(context_a); - channel_b.EnableEncryption(context_b); + channel_a->EnableEncryption(context_a); + channel_b->EnableEncryption(context_b); - EXPECT_EQ(channel_a.GetType(), "ENCRYPTED_BLUETOOTH"); - EXPECT_EQ(channel_b.GetType(), "ENCRYPTED_BLUETOOTH"); - EXPECT_TRUE(channel_a.IsEncrypted()); - EXPECT_TRUE(channel_b.IsEncrypted()); + EXPECT_EQ(channel_a->GetType(), "ENCRYPTED_BLUETOOTH"); + EXPECT_EQ(channel_b->GetType(), "ENCRYPTED_BLUETOOTH"); + EXPECT_TRUE(channel_a->IsEncrypted()); + EXPECT_TRUE(channel_b->IsEncrypted()); // Start data transfer absl::string_view tx_message = "data message"; - channel_a.Write(tx_message); - ByteArray rx_message = std::move(channel_b.Read().result()); + channel_a->Write(tx_message); + ByteArray rx_message = std::move(channel_b->Read().result()); // Verify expectations. EXPECT_EQ(rx_message.AsStringView(), tx_message); @@ -411,8 +421,8 @@ TEST_F(BaseEndpointChannelTest, EncryptedReadWriteCanNotBeIntercepted) { } // Shutdown test environment. - channel_a.Close(DisconnectionReason::LOCAL_DISCONNECTION); - channel_b.Close(DisconnectionReason::REMOTE_DISCONNECTION); + channel_a->Close(DisconnectionReason::LOCAL_DISCONNECTION); + channel_b->Close(DisconnectionReason::REMOTE_DISCONNECTION); } TEST_F(BaseEndpointChannelTest, CanBesuspendedAndResumed) { @@ -486,43 +496,45 @@ TEST_F(BaseEndpointChannelTest, ReadUnencryptedFrameOnEncryptedChannel) { // Setup test communication environment. auto pipe_a = CreatePipe(); // channel_a writes to pipe_a, reads from pipe_b. auto pipe_b = CreatePipe(); // channel_b writes to pipe_b, reads from pipe_a. - TestEndpointChannel channel_a(pipe_b.first.get(), pipe_a.second.get()); - TestEndpointChannel channel_b(pipe_a.first.get(), pipe_b.second.get()); + auto channel_a = std::make_shared(pipe_b.first.get(), + pipe_a.second.get()); + auto channel_b = std::make_shared(pipe_a.first.get(), + pipe_b.second.get()); - ON_CALL(channel_a, GetMedium).WillByDefault([]() { + ON_CALL(*channel_a, GetMedium).WillByDefault([]() { return Medium::BLUETOOTH; }); - ON_CALL(channel_b, GetMedium).WillByDefault([]() { + ON_CALL(*channel_b, GetMedium).WillByDefault([]() { return Medium::BLUETOOTH; }); // Run DH key exchange; setup encryption contexts for channels. But only // encrypt |channel_b|. - auto [context_a, context_b] = DoDhKeyExchange(&channel_a, &channel_b); + auto [context_a, context_b] = DoDhKeyExchange(channel_a, channel_b); ASSERT_NE(context_a, nullptr); ASSERT_NE(context_b, nullptr); - channel_b.EnableEncryption(context_b); + channel_b->EnableEncryption(context_b); - EXPECT_EQ(channel_a.GetType(), "BLUETOOTH"); - EXPECT_EQ(channel_b.GetType(), "ENCRYPTED_BLUETOOTH"); + EXPECT_EQ(channel_a->GetType(), "BLUETOOTH"); + EXPECT_EQ(channel_b->GetType(), "ENCRYPTED_BLUETOOTH"); // An unencrypted KeepAlive should succeed. std::string keep_alive_message = parser::ForKeepAlive(); - channel_a.Write(keep_alive_message); - ExceptionOr result = channel_b.Read(); + channel_a->Write(keep_alive_message); + ExceptionOr result = channel_b->Read(); EXPECT_TRUE(result.ok()); EXPECT_EQ(result.result().AsStringView(), keep_alive_message); // An unencrypted data frame should fail. absl::string_view tx_message = "data message"; - channel_a.Write(tx_message); - result = channel_b.Read(); + channel_a->Write(tx_message); + result = channel_b->Read(); EXPECT_FALSE(result.ok()); EXPECT_EQ(result.exception(), Exception::kInvalidProtocolBuffer); // Shutdown test environment. - channel_a.Close(DisconnectionReason::LOCAL_DISCONNECTION); - channel_b.Close(DisconnectionReason::REMOTE_DISCONNECTION); + channel_a->Close(DisconnectionReason::LOCAL_DISCONNECTION); + channel_b->Close(DisconnectionReason::REMOTE_DISCONNECTION); } } // namespace diff --git a/connections/implementation/base_pcp_handler.cc b/connections/implementation/base_pcp_handler.cc index 1aab772d..0d79ceac 100644 --- a/connections/implementation/base_pcp_handler.cc +++ b/connections/implementation/base_pcp_handler.cc @@ -212,8 +212,7 @@ std::vector BasePcpHandler::GetConnectionInfoFromResult( } WifiLanConnectionInfo info( std::string(ip_address.begin(), ip_address.end()), - absl::StrCat(absl::Hex(port, absl::kZeroPad16)), "", - {}); + absl::StrCat(absl::Hex(port, absl::kZeroPad16)), "", {}); connection_infos.push_back(info); } } @@ -424,25 +423,25 @@ BooleanMediumSelector BasePcpHandler::ComputeIntersectionOfSupportedMediums( pending_connection_info.connection_options.connection_info .supported_wifi_direct_auth_types; LOG(INFO) << "Remote supported WifiDirect auth types: " - << absl::StrJoin( - remote_supported_wifi_direct_auth_types, ", ", - [](std::string* out, int auth_type) { - absl::StrAppend( - out, - WifiDirectAuthType_Name( - static_cast(auth_type))); - }); + << absl::StrJoin( + remote_supported_wifi_direct_auth_types, ", ", + [](std::string* out, int auth_type) { + absl::StrAppend( + out, + WifiDirectAuthType_Name( + static_cast(auth_type))); + }); auto local_supported_wifi_direct_auth_types = mediums_->GetWifiDirect().GetSupportedWifiDirectAuthTypes(); LOG(INFO) << "Local supported WifiDirect auth types: " - << absl::StrJoin( - local_supported_wifi_direct_auth_types, ", ", - [](std::string* out, int auth_type) { - absl::StrAppend( - out, - WifiDirectAuthType_Name( - static_cast(auth_type))); - }); + << absl::StrJoin( + local_supported_wifi_direct_auth_types, ", ", + [](std::string* out, int auth_type) { + absl::StrAppend( + out, + WifiDirectAuthType_Name( + static_cast(auth_type))); + }); bool found_common_auth_type = false; for (const auto& auth_type : local_supported_wifi_direct_auth_types) { if (auth_type == WifiDirectAuthType::WIFI_DIRECT_TYPE_UNKNOWN) { @@ -587,34 +586,50 @@ void BasePcpHandler::RunOnPcpHandlerThread(const std::string& name, serial_executor_.Execute(name, std::move(runnable)); } -EncryptionRunner::ResultListener BasePcpHandler::GetResultListener() { +EncryptionRunner::ResultListener BasePcpHandler::GetResultListener( + std::shared_ptr endpoint_channel) { + std::weak_ptr weak_channel = endpoint_channel; + return { .on_success_cb = - [this](const std::string& endpoint_id, - std::unique_ptr ukey2, - const std::string& auth_token, - const ByteArray& raw_auth_token) { + [this, weak_channel](const std::string& endpoint_id, + std::unique_ptr ukey2, + const std::string& auth_token, + const ByteArray& raw_auth_token) { + auto channel = weak_channel.lock(); + if (!channel) return; + RunOnPcpHandlerThread( "encryption-success", - [this, endpoint_id, raw_ukey2 = ukey2.release(), auth_token, - raw_auth_token]() RUN_ON_PCP_HANDLER_THREAD() mutable { - OnEncryptionSuccessRunnable( - endpoint_id, std::unique_ptr(raw_ukey2), - auth_token, raw_auth_token); - }); + [this, endpoint_id, weak_channel, raw_ukey2 = ukey2.release(), + auth_token, raw_auth_token]() + RUN_ON_PCP_HANDLER_THREAD() mutable { + std::unique_ptr ukey2(raw_ukey2); + auto channel = weak_channel.lock(); + if (!channel) return; + OnEncryptionSuccessRunnable(endpoint_id, std::move(ukey2), + auth_token, raw_auth_token, + channel); + }); }, .on_failure_cb = - [this](const std::string& endpoint_id, EndpointChannel* channel) { + [this, weak_channel](const std::string& endpoint_id) { + auto channel = weak_channel.lock(); + if (!channel) return; + RunOnPcpHandlerThread( "encryption-failure", - [this, endpoint_id, channel]() RUN_ON_PCP_HANDLER_THREAD() { - LOG(ERROR) - << "Encryption failed for endpoint_id=" << endpoint_id - << " on medium=" - << location::nearby::proto::connections::Medium_Name( - channel->GetMedium()); - OnEncryptionFailureRunnable(endpoint_id, channel); - }); + [this, endpoint_id, weak_channel]() + RUN_ON_PCP_HANDLER_THREAD() { + auto channel = weak_channel.lock(); + if (!channel) return; + LOG(ERROR) + << "Encryption failed for endpoint_id=" << endpoint_id + << " on medium=" + << location::nearby::proto::connections::Medium_Name( + channel->GetMedium()); + OnEncryptionFailureRunnable(endpoint_id, channel); + }); }, }; } @@ -622,36 +637,49 @@ EncryptionRunner::ResultListener BasePcpHandler::GetResultListener() { EncryptionRunner::ResultListener BasePcpHandler::GetResultListenerV3( const NearbyDeviceProvider& device_provider, const NearbyDevice& remote_device, - const EndpointChannel& endpoint_channel) { + std::shared_ptr endpoint_channel) { + std::weak_ptr weak_channel = endpoint_channel; + return { .on_success_cb = - [this, &device_provider, &remote_device, &endpoint_channel]( + [this, &device_provider, &remote_device, weak_channel]( const std::string& endpoint_id, std::unique_ptr ukey2, const std::string& auth_token, const ByteArray& raw_auth_token) { + auto channel = weak_channel.lock(); + if (!channel) return; + RunOnPcpHandlerThread( "encryption-success", - [this, &device_provider, &remote_device, &endpoint_channel, - raw_ukey2 = ukey2.release(), auth_token, - raw_auth_token]() RUN_ON_PCP_HANDLER_THREAD() mutable { - OnEncryptionSuccessRunnableV3( - remote_device, std::unique_ptr(raw_ukey2), - auth_token, raw_auth_token, endpoint_channel, - device_provider); - }); + [this, &device_provider, &remote_device, weak_channel, + raw_ukey2 = ukey2.release(), auth_token, raw_auth_token]() + RUN_ON_PCP_HANDLER_THREAD() mutable { + std::unique_ptr ukey2(raw_ukey2); + auto channel = weak_channel.lock(); + if (!channel) return; + OnEncryptionSuccessRunnableV3( + remote_device, std::move(ukey2), auth_token, + raw_auth_token, channel, device_provider); + }); }, .on_failure_cb = - [this](const std::string& endpoint_id, EndpointChannel* channel) { + [this, weak_channel](const std::string& endpoint_id) { + auto channel = weak_channel.lock(); + if (!channel) return; + RunOnPcpHandlerThread( "encryption-failure", - [this, endpoint_id, channel]() RUN_ON_PCP_HANDLER_THREAD() { - LOG(ERROR) - << "Encryption failed for endpoint_id=" << endpoint_id - << " on medium=" - << location::nearby::proto::connections::Medium_Name( - channel->GetMedium()); - OnEncryptionFailureRunnable(endpoint_id, channel); - }); + [this, endpoint_id, weak_channel]() + RUN_ON_PCP_HANDLER_THREAD() { + auto channel = weak_channel.lock(); + if (!channel) return; + LOG(ERROR) + << "Encryption failed for endpoint_id=" << endpoint_id + << " on medium=" + << location::nearby::proto::connections::Medium_Name( + channel->GetMedium()); + OnEncryptionFailureRunnable(endpoint_id, channel); + }); }, }; } @@ -659,7 +687,7 @@ EncryptionRunner::ResultListener BasePcpHandler::GetResultListenerV3( void BasePcpHandler::OnEncryptionSuccessRunnableV3( const NearbyDevice& remote_device, std::unique_ptr ukey2, absl::string_view auth_token, const ByteArray& raw_auth_token, - const EndpointChannel& endpoint_channel, + std::shared_ptr endpoint_channel, const NearbyDeviceProvider& device_provider) { // Quick fail if we've been removed from pending connections while we were // busy running UKEY2. @@ -674,7 +702,11 @@ void BasePcpHandler::OnEncryptionSuccessRunnableV3( } BasePcpHandler::PendingConnectionInfo& pending_connection_info = it->second; - + // Verify pointer equality to avoid accidental action on superseded + // channels. + if (endpoint_channel != pending_connection_info.channel) { + return; + } // TODO(b/300149127): Add test coverage. if (!ukey2) { // Fail early, if there is no crypto context. @@ -724,7 +756,8 @@ void BasePcpHandler::OnEncryptionSuccessRunnableV3( void BasePcpHandler::OnEncryptionSuccessRunnable( const std::string& endpoint_id, std::unique_ptr ukey2, - const std::string& auth_token, const ByteArray& raw_auth_token) { + const std::string& auth_token, const ByteArray& raw_auth_token, + std::shared_ptr endpoint_channel) { // Quick fail if we've been removed from pending connections while we were // busy running UKEY2. // TODO(b/316421187): Add test coverage @@ -738,6 +771,12 @@ void BasePcpHandler::OnEncryptionSuccessRunnable( BasePcpHandler::PendingConnectionInfo& pending_connection_info = it->second; + // Verify pointer equality to avoid accidental action on superseded + // channels. + if (endpoint_channel != pending_connection_info.channel) { + return; + } + if (!ukey2) { // Fail early, if there is no crypto context. ProcessPreConnectionInitiationFailure( @@ -801,24 +840,21 @@ void BasePcpHandler::RegisterDeviceAfterEncryptionSuccess( } void BasePcpHandler::OnEncryptionFailureRunnable( - const std::string& endpoint_id, EndpointChannel* endpoint_channel) { + const std::string& endpoint_id, + std::shared_ptr endpoint_channel) { auto it = pending_connections_.find(endpoint_id); if (it == pending_connections_.end()) { LOG(INFO) - << "Connection not found on UKEY negotination complete; endpoint_id=" + << "Connection not found on UKEY negotiation complete; endpoint_id=" << endpoint_id; return; } BasePcpHandler::PendingConnectionInfo& pending_connection_info = it->second; - // We had a bug here, caused by a race with EncryptionRunner. We now verify - // the EndpointChannel to avoid it. In a simultaneous connection, we clean - // up one of the two EndpointChannels and then update our pendingConnections - // with the winning channel's state. Closing a channel that was in the - // middle of EncryptionRunner would trigger onEncryptionFailed, and, since - // the map had already updated with the winning EndpointChannel, we closed - // it too by accident. - if (*endpoint_channel != *pending_connection_info.channel) { + + // Verify pointer equality to avoid accidental action on superseded + // channels. + if (endpoint_channel != pending_connection_info.channel) { LOG(INFO) << "Not destroying channel [mismatch]: passed=" << endpoint_channel->GetName() << "; expected=" << pending_connection_info.channel->GetName(); @@ -871,8 +907,8 @@ ConnectionInfo BasePcpHandler::FillConnectionInfo( connection_info.supported_wifi_direct_auth_types = mediums_->GetWifiDirect().GetSupportedWifiDirectAuthTypes(); VLOG(1) << "Set SupportedWifiDirectAuthTypes for WIFI_DIRECT: " - << absl::StrJoin(connection_info.supported_wifi_direct_auth_types, - ","); + << absl::StrJoin(connection_info.supported_wifi_direct_auth_types, + ","); } else { connection_info.supported_wifi_direct_auth_types = {}; } @@ -1005,17 +1041,32 @@ Status BasePcpHandler::RequestConnection( pending_connection_info.medium = channel->GetMedium(); pending_connection_info.channel = std::move(channel); - EndpointChannel* endpoint_channel = - pending_connections_ - .emplace(endpoint_id, std::move(pending_connection_info)) - .first->second.channel.get(); + std::shared_ptr channel_to_close_on_failure = + pending_connection_info.channel; + auto [it, inserted] = pending_connections_.emplace( + endpoint_id, std::move(pending_connection_info)); + if (!inserted) { + LOG(ERROR) << "Failed to add outgoing connection to pending set; " + "endpoint_id=" + << endpoint_id + << ". Likely a collision with an existing pending " + "connection."; + if (channel_to_close_on_failure) { + channel_to_close_on_failure->Close( + location::nearby::proto::connections::DisconnectionReason:: + IO_ERROR); + } + result->Set({Status::kEndpointIoError}); + return; + } + std::shared_ptr endpoint_channel = it->second.channel; LOG(INFO) << "Initiating secure connection: endpoint_id=" << endpoint_id; // Next, we'll set up encryption. When it's done, our future will return // and RequestConnection() will finish. encryption_runner_.StartClient(client, endpoint_id, endpoint_channel, - GetResultListener()); + GetResultListener(endpoint_channel)); }); LOG(INFO) << "Waiting for connection to complete: endpoint_id=" << endpoint_id; @@ -1152,10 +1203,25 @@ Status BasePcpHandler::RequestConnectionV3( pending_connection_info.medium = channel->GetMedium(); pending_connection_info.channel = std::move(channel); - EndpointChannel* endpoint_channel = - pending_connections_ - .emplace(endpoint_id, std::move(pending_connection_info)) - .first->second.channel.get(); + std::shared_ptr channel_to_close_on_failure = + pending_connection_info.channel; + auto [it, inserted] = pending_connections_.emplace( + endpoint_id, std::move(pending_connection_info)); + if (!inserted) { + LOG(ERROR) << "Failed to add outgoing connection to pending set; " + "endpoint_id=" + << endpoint_id + << ". Likely a collision with an existing pending " + "connection."; + if (channel_to_close_on_failure) { + channel_to_close_on_failure->Close( + location::nearby::proto::connections::DisconnectionReason:: + IO_ERROR); + } + result->Set({Status::kEndpointIoError}); + return; + } + std::shared_ptr endpoint_channel = it->second.channel; LOG(INFO) << "Initiating secure connection: endpoint_id=" << endpoint_id; @@ -1165,7 +1231,7 @@ Status BasePcpHandler::RequestConnectionV3( encryption_runner_.StartClient( client, endpoint_id, endpoint_channel, GetResultListenerV3(*(client->GetLocalDeviceProvider()), - remote_device, *endpoint_channel)); + remote_device, endpoint_channel)); }); LOG(INFO) << "Waiting for connection to complete: endpoint_id=" << endpoint_id; @@ -2133,14 +2199,23 @@ Exception BasePcpHandler::OnIncomingConnection( pending_connection_info.medium = channel->GetMedium(); pending_connection_info.channel = std::move(channel); - auto* owned_channel = pending_connections_ - .emplace(connection_request.endpoint_id(), - std::move(pending_connection_info)) - .first->second.channel.get(); + auto [it, inserted] = pending_connections_.emplace( + connection_request.endpoint_id(), std::move(pending_connection_info)); + // This should not happen since BreakTie() above should have checked that + // the endpoint_id is not already in pending_connections_. + if (!inserted) { + LOG(ERROR) << "Failed to add incoming connection to pending set; " + "endpoint_id=" + << connection_request.endpoint_id() + << ". Likely a collision with an existing pending connection."; + return {Exception::kIo}; + } + std::shared_ptr endpoint_channel = it->second.channel; // Next, we'll set up encryption. encryption_runner_.StartServer(client, connection_request.endpoint_id(), - owned_channel, GetResultListener()); + endpoint_channel, + GetResultListener(endpoint_channel)); return {Exception::kSuccess}; } diff --git a/connections/implementation/base_pcp_handler.h b/connections/implementation/base_pcp_handler.h index a1c3db7c..f14665f6 100644 --- a/connections/implementation/base_pcp_handler.h +++ b/connections/implementation/base_pcp_handler.h @@ -476,11 +476,11 @@ class BasePcpHandler : public PcpHandler, // Only (possibly) vector for incoming connections. std::vector supported_mediums; - // Keep track of a channel before we pass it to EndpointChannelManager. This - // is owned until the call to OnEncryptionSuccessRunnableV3 or - // OnEncryptionSuccessRunnable when ownership is transferred to the - // EndpointManager. - std::unique_ptr channel; + // Keep track of a channel before it is registered with the + // EndpointManager. This reference is held during the handshake phase and + // passed to the EndpointManager upon successful encryption + // (OnEncryptionSuccessRunnableV3 or OnEncryptionSuccessRunnable). + std::shared_ptr channel; // Crypto context; initially empty; established first thing after channel // creation by running UKey2 session. While it is in progress, we keep track @@ -509,24 +509,27 @@ class BasePcpHandler : public PcpHandler, void OnEncryptionFailureImpl(const std::string& endpoint_id, EndpointChannel* channel); - EncryptionRunner::ResultListener GetResultListener(); + EncryptionRunner::ResultListener GetResultListener( + std::shared_ptr endpoint_channel); EncryptionRunner::ResultListener GetResultListenerV3( const NearbyDeviceProvider& device_provider, const NearbyDevice& remote_device, - const EndpointChannel& endpoint_channel); + std::shared_ptr endpoint_channel); void OnEncryptionSuccessRunnable( const std::string& endpoint_id, std::unique_ptr ukey2, - const std::string& auth_token, const ByteArray& raw_auth_token); + const std::string& auth_token, const ByteArray& raw_auth_token, + std::shared_ptr endpoint_channel); void OnEncryptionSuccessRunnableV3( const NearbyDevice& remote_device, std::unique_ptr<::securegcm::UKey2Handshake> ukey2, absl::string_view auth_token, const ByteArray& raw_auth_token, - const EndpointChannel& endpoint_channel, + std::shared_ptr endpoint_channel, const NearbyDeviceProvider& device_provider); - void OnEncryptionFailureRunnable(const std::string& endpoint_id, - EndpointChannel* endpoint_channel); + void OnEncryptionFailureRunnable( + const std::string& endpoint_id, + std::shared_ptr endpoint_channel); void RegisterDeviceAfterEncryptionSuccess( std::string_view endpoint_id, std::unique_ptr<::securegcm::UKey2Handshake> ukey2, diff --git a/connections/implementation/base_pcp_handler_test.cc b/connections/implementation/base_pcp_handler_test.cc index 351e3dc3..e1fad699 100644 --- a/connections/implementation/base_pcp_handler_test.cc +++ b/connections/implementation/base_pcp_handler_test.cc @@ -460,9 +460,7 @@ class BasePcpHandlerTest MacAddress::FromString("12:34:56:78:9a:bc", remote_mac_address_); } - void TearDown() override { - env_.Stop(); - } + void TearDown() override { env_.Stop(); } void StartAdvertising(ClientProxy* client, MockPcpHandler* pcp_handler, BooleanMediumSelector allowed = GetParam()) { @@ -645,7 +643,7 @@ class BasePcpHandlerTest void RequestConnection( const std::string& endpoint_id, std::unique_ptr channel_a, - MockEndpointChannel* channel_b, ClientProxy* client, + std::shared_ptr channel_b, ClientProxy* client, MockPcpHandler* pcp_handler, location::nearby::proto::connections::Medium connect_medium, std::atomic_int* flag = nullptr, @@ -715,7 +713,7 @@ class BasePcpHandlerTest void RequestConnectionV3( const NearbyDevice& remote_device, std::unique_ptr channel_a, - MockEndpointChannel* channel_b, ClientProxy* client, + std::shared_ptr channel_b, ClientProxy* client, MockPcpHandler* pcp_handler, location::nearby::proto::connections::Medium connect_medium, FakePresenceDeviceProvider* fake_presence_device_provider, @@ -795,7 +793,7 @@ class BasePcpHandlerTest void RequestConnectionWifiLanFail( const std::string& endpoint_id, std::unique_ptr channel_a, - MockEndpointChannel* channel_b, ClientProxy* client, + std::shared_ptr channel_b, ClientProxy* client, MockPcpHandler* pcp_handler, std::atomic_int* flag = nullptr, Status expected_result = {Status::kSuccess}) { ConnectionRequestInfo info{ @@ -1136,12 +1134,13 @@ TEST_F(BasePcpHandlerTest, WifiMediumFailFallBackToBT) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); - RequestConnectionWifiLanFail(endpoint_id, std::move(channel_a), - channel_b.get(), client_.get(), &pcp_handler); + RequestConnectionWifiLanFail(endpoint_id, std::move(channel_a), channel_b, + client_.get(), &pcp_handler); LOG(INFO) << "RequestConnection complete"; channel_b->Close(); bwu.Shutdown(); @@ -1161,12 +1160,13 @@ TEST_P(BasePcpHandlerTest, RequestConnectionChangesState) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); - RequestConnection("1234", std::move(channel_a), channel_b.get(), - client_.get(), &pcp_handler, connect_medium); + RequestConnection("1234", std::move(channel_a), channel_b, client_.get(), + &pcp_handler, connect_medium); LOG(INFO) << "RequestConnection complete"; EXPECT_TRUE(pcp_handler.HasOutgoingConnections(client_.get())); EXPECT_FALSE(pcp_handler.HasIncomingConnections(client_.get())); @@ -1207,12 +1207,13 @@ TEST_P(BasePcpHandlerTest, CanRequestConnectionPresence) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); - RequestConnection("1234", std::move(channel_a), channel_b.get(), - client_.get(), &pcp_handler, connect_medium); + RequestConnection("1234", std::move(channel_a), channel_b, client_.get(), + &pcp_handler, connect_medium); LOG(INFO) << "RequestConnection complete"; channel_b->Close(); bwu.Shutdown(); @@ -1236,12 +1237,13 @@ TEST_P(BasePcpHandlerTest, CanRequestConnectionLegacy) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); - RequestConnection("1234", std::move(channel_a), channel_b.get(), - client_.get(), &pcp_handler, connect_medium); + RequestConnection("1234", std::move(channel_a), channel_b, client_.get(), + &pcp_handler, connect_medium); LOG(INFO) << "RequestConnection complete"; channel_b->Close(); bwu.Shutdown(); @@ -1266,11 +1268,12 @@ TEST_P(BasePcpHandlerTest, RequestConnectionV3) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - const auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); - RequestConnectionV3(mock_device_, std::move(channel_a), channel_b.get(), + RequestConnectionV3(mock_device_, std::move(channel_a), channel_b, client_.get(), &pcp_handler, connect_medium, &provider); LOG(INFO) << "RequestConnectionV3 complete"; channel_b->Close(); @@ -1297,12 +1300,13 @@ TEST_P(BasePcpHandlerTest, RequestConnectionV3_AuthenticationFailure) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - const auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); RequestConnectionV3( - mock_device_, std::move(channel_a), channel_b.get(), client_.get(), + mock_device_, std::move(channel_a), channel_b, client_.get(), &pcp_handler, connect_medium, &provider, /*flag=*/nullptr, /*expected_result=*/{Status::kSuccess}, /*expected_authentication_status=*/AuthenticationStatus::kFailure); @@ -1328,7 +1332,8 @@ TEST_P(BasePcpHandlerTest, RequestConnectionV3_ConnectImplFailure) { auto mediums = pcp_handler.GetDiscoveryMediums(client_.get()); auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnectionForConnectFailure(connect_medium); - const auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); ConnectionRequestInfo info{ @@ -1403,7 +1408,8 @@ TEST_P(BasePcpHandlerTest, RequestConnection_ConnectImplFailure) { auto mediums = pcp_handler.GetDiscoveryMediums(client_.get()); auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnectionForConnectFailure(connect_medium); - const auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); ConnectionRequestInfo info{ @@ -1475,12 +1481,13 @@ TEST_P(BasePcpHandlerTest, IoError_RequestConnectionV3Fails) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(AtLeast(1)); EXPECT_CALL(*channel_b, CloseImpl).Times(AtLeast(1)); channel_b->broken_write_ = true; EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); - RequestConnectionV3(mock_device_, std::move(channel_a), channel_b.get(), + RequestConnectionV3(mock_device_, std::move(channel_a), channel_b, client_.get(), &pcp_handler, connect_medium, nullptr, nullptr, {Status::kEndpointIoError}); LOG(INFO) << "RequestConnectionV3 complete"; @@ -1503,13 +1510,14 @@ TEST_P(BasePcpHandlerTest, IoError_RequestConnectionFails) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(AtLeast(1)); EXPECT_CALL(*channel_b, CloseImpl).Times(AtLeast(1)); channel_b->broken_write_ = true; EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(AtLeast(0)); - RequestConnection(endpoint_id, std::move(channel_a), channel_b.get(), - client_.get(), &pcp_handler, connect_medium, nullptr, + RequestConnection(endpoint_id, std::move(channel_a), channel_b, client_.get(), + &pcp_handler, connect_medium, nullptr, {Status::kEndpointIoError}); LOG(INFO) << "RequestConnection complete"; channel_b->Close(); @@ -1531,11 +1539,12 @@ TEST_P(BasePcpHandlerTest, AcceptConnectionChangesState) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); - RequestConnection(endpoint_id, std::move(channel_a), channel_b.get(), - client_.get(), &pcp_handler, connect_medium); + RequestConnection(endpoint_id, std::move(channel_a), channel_b, client_.get(), + &pcp_handler, connect_medium); LOG(INFO) << "Attempting to accept connection: id=" << endpoint_id; EXPECT_EQ(pcp_handler.AcceptConnection(client_.get(), endpoint_id, {}), Status{Status::kSuccess}); @@ -1559,9 +1568,10 @@ TEST_P(BasePcpHandlerTest, RejectConnectionChangesState) { auto mediums = pcp_handler.GetDiscoveryMediums(client_.get()); auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(mock_connection_listener_.rejected_cb, Call).Times(1); - RequestConnection(endpoint_id, std::move(channel_pair.first), channel_b.get(), + RequestConnection(endpoint_id, std::move(channel_pair.first), channel_b, client_.get(), &pcp_handler, connect_medium); LOG(INFO) << "Attempting to reject connection: id=" << endpoint_id; EXPECT_EQ(pcp_handler.RejectConnection(client_.get(), endpoint_id), @@ -1587,11 +1597,12 @@ TEST_P(BasePcpHandlerTest, OnIncomingFrameChangesState) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); - RequestConnection(endpoint_id, std::move(channel_a), channel_b.get(), - client_.get(), &pcp_handler, connect_medium); + RequestConnection(endpoint_id, std::move(channel_a), channel_b, client_.get(), + &pcp_handler, connect_medium); LOG(INFO) << "Attempting to accept connection: id=" << endpoint_id; EXPECT_CALL(mock_connection_listener_.accepted_cb, Call).Times(1); EXPECT_CALL(mock_connection_listener_.disconnected_cb, Call) @@ -1628,10 +1639,11 @@ TEST_P(BasePcpHandlerTest, DestructorIsCalledOnProtocolEndpoint) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); - RequestConnection(endpoint_id, std::move(channel_a), channel_b.get(), + RequestConnection(endpoint_id, std::move(channel_a), channel_b, client_.get(), &pcp_handler, connect_medium, &destroyed_flag); mediums_count = mediums.size(); @@ -1670,11 +1682,12 @@ TEST_P(BasePcpHandlerTest, MultipleMediumsProduceSingleEndpointLostEvent) { auto connect_medium = mediums[mediums.size() - 1]; auto channel_pair = SetupConnection(connect_medium); auto& channel_a = channel_pair.first; - auto& channel_b = channel_pair.second; + std::shared_ptr channel_b = + std::move(channel_pair.second); EXPECT_CALL(*channel_a, CloseImpl).Times(1); EXPECT_CALL(*channel_b, CloseImpl).Times(1); EXPECT_CALL(mock_discovery_listener_.endpoint_lost_cb, Call).Times(1); - RequestConnection(endpoint_id, std::move(channel_a), channel_b.get(), + RequestConnection(endpoint_id, std::move(channel_a), channel_b, client_.get(), &pcp_handler, connect_medium, &destroyed_flag); auto allowed_mediums = pcp_handler.GetDiscoveryMediums(client_.get()); @@ -2543,7 +2556,8 @@ TEST_F(BasePcpHandlerTest, TestDeviceFilterForConnectionsWithPresence) { } TEST_F(BasePcpHandlerTest, IncomingConnectionFailsWithEmptyEndpointId) { - env_.Start(); + env_.Start({.use_simulated_clock = true}); + client_ = std::make_unique(&mock_event_logger_); Mediums m; EndpointChannelManager ecm; EndpointManager em(&ecm); @@ -2600,7 +2614,6 @@ TEST_F(BasePcpHandlerTest, IncomingConnectionFailsWithEmptyEndpointId) { )pb"; absl::string_view client_session_log = R"pb( event_type: CLIENT_SESSION - client_session { duration_millis: 0 } version: "v1.5.0" )pb"; EXPECT_CALL(mock_event_logger_, @@ -2615,9 +2628,8 @@ TEST_F(BasePcpHandlerTest, IncomingConnectionFailsWithEmptyEndpointId) { Log(Matcher( HasEventType(EventType::START_CLIENT_SESSION)))) .Times(3); - EXPECT_CALL( - mock_event_logger_, - Log(Matcher(EqualsProto(client_session_log)))) + EXPECT_CALL(mock_event_logger_, Log(Matcher(Partially( + EqualsProto(client_session_log))))) .Times(2); EXPECT_CALL(mock_event_logger_, Log(Matcher( Partially(EqualsProto(expected_log))))); diff --git a/connections/implementation/connections_authentication_transport.cc b/connections/implementation/connections_authentication_transport.cc index ec669d9b..f4bb001d 100644 --- a/connections/implementation/connections_authentication_transport.cc +++ b/connections/implementation/connections_authentication_transport.cc @@ -14,7 +14,9 @@ #include "connections/implementation/connections_authentication_transport.h" +#include #include +#include #include "absl/strings/string_view.h" #include "connections/implementation/endpoint_channel.h" @@ -24,19 +26,18 @@ namespace nearby { namespace connections { ConnectionsAuthenticationTransport::ConnectionsAuthenticationTransport( - const EndpointChannel& channel) { - channel_ = const_cast(&channel); -} + std::shared_ptr channel) + : channel_(std::move(channel)) {} void ConnectionsAuthenticationTransport::WriteMessage( absl::string_view message) const { - // channel_ should never be null. + // channel_ is guaranteed valid by shared_ptr ownership CHECK(channel_ != nullptr); channel_->Write(message); } std::string ConnectionsAuthenticationTransport::ReadMessage() const { - // channel_ should never be null. + // channel_ is guaranteed valid by shared_ptr ownership CHECK(channel_ != nullptr); auto response = channel_->Read(); if (response.ok()) { diff --git a/connections/implementation/connections_authentication_transport.h b/connections/implementation/connections_authentication_transport.h index ab7886d4..a4bcafac 100644 --- a/connections/implementation/connections_authentication_transport.h +++ b/connections/implementation/connections_authentication_transport.h @@ -15,6 +15,7 @@ #ifndef THIRD_PARTY_NEARBY_CONNECTIONS_IMPLEMENTATION_CONNECTIONS_AUTHENTICATION_TRANSPORT_H_ #define THIRD_PARTY_NEARBY_CONNECTIONS_IMPLEMENTATION_CONNECTIONS_AUTHENTICATION_TRANSPORT_H_ +#include #include #include "absl/strings/string_view.h" @@ -30,12 +31,13 @@ namespace connections { class ConnectionsAuthenticationTransport : public nearby::AuthenticationTransport { public: - explicit ConnectionsAuthenticationTransport(const EndpointChannel& channel); + explicit ConnectionsAuthenticationTransport( + std::shared_ptr channel); void WriteMessage(absl::string_view message) const override; std::string ReadMessage() const override; private: - EndpointChannel* channel_; + std::shared_ptr channel_; }; } // namespace connections diff --git a/connections/implementation/connections_authentication_transport_test.cc b/connections/implementation/connections_authentication_transport_test.cc index 4198080f..667fbab6 100644 --- a/connections/implementation/connections_authentication_transport_test.cc +++ b/connections/implementation/connections_authentication_transport_test.cc @@ -85,35 +85,38 @@ class MockEndpointChannel : public EndpointChannel { }; TEST(ConnectionsAuthenticationTransportTest, TestWriteMessage) { - MockEndpointChannel channel; + auto channel = std::make_shared(); + auto* channel_ptr = channel.get(); ConnectionsAuthenticationTransport transport(channel); - EXPECT_CALL(channel, Write(_)).WillOnce([&channel](absl::string_view data) { - channel.messages_.push_back(std::string(data)); - return Exception{ - .value = Exception::Value::kSuccess, - }; - }); + EXPECT_CALL(*channel, Write(_)) + .WillOnce([channel_ptr](absl::string_view data) { + channel_ptr->messages_.push_back(std::string(data)); + return Exception{ + .value = Exception::Value::kSuccess, + }; + }); transport.WriteMessage("hello world"); - EXPECT_THAT(channel.messages_, testing::ElementsAre("hello world")); + EXPECT_THAT(channel_ptr->messages_, testing::ElementsAre("hello world")); } TEST(ConnectionsAuthenticationTransportTest, TestReadMessage) { - MockEndpointChannel channel; + auto channel = std::make_shared(); + auto* channel_ptr = channel.get(); ConnectionsAuthenticationTransport transport(channel); - channel.messages_.push_back("hello world"); - EXPECT_CALL(channel, Read()).WillOnce([&channel]() { - std::string ret = channel.messages_[0]; - channel.messages_.erase(channel.messages_.begin()); + channel_ptr->messages_.push_back("hello world"); + EXPECT_CALL(*channel, Read()).WillOnce([channel_ptr]() { + std::string ret = channel_ptr->messages_[0]; + channel_ptr->messages_.erase(channel_ptr->messages_.begin()); return ExceptionOr(ByteArray(ret)); }); EXPECT_EQ(transport.ReadMessage(), "hello world"); } TEST(ConnectionsAuthenticationTransportTest, TestReadMessageFail) { - MockEndpointChannel channel; + auto channel = std::make_shared(); ConnectionsAuthenticationTransport transport(channel); - channel.messages_.push_back("hello world"); - EXPECT_CALL(channel, Read()).WillOnce([]() { + channel->messages_.push_back("hello world"); + EXPECT_CALL(*channel, Read()).WillOnce([]() { return ExceptionOr(Exception::Value::kIo); }); EXPECT_EQ(transport.ReadMessage(), ""); diff --git a/connections/implementation/encryption_runner.cc b/connections/implementation/encryption_runner.cc index 7a2842db..a3ade421 100644 --- a/connections/implementation/encryption_runner.cc +++ b/connections/implementation/encryption_runner.cc @@ -68,9 +68,9 @@ bool HandleEncryptionSuccess(const std::string& endpoint_id, return true; } -void CancelableAlarmRunnable(ClientProxy* client, - const std::string& endpoint_id, - EndpointChannel* endpoint_channel) { +void CancelableAlarmRunnable( + ClientProxy* client, const std::string& endpoint_id, + std::shared_ptr endpoint_channel) { LOG(INFO) << "Timing out encryption for client " << client->GetClientId() << " to endpoint_id=" << endpoint_id << " after " << absl::FormatDuration(kTimeout); @@ -80,18 +80,31 @@ void CancelableAlarmRunnable(ClientProxy* client, class ServerRunnable final { public: ServerRunnable(ClientProxy* client, ScheduledExecutor* alarm_executor, - const std::string& endpoint_id, EndpointChannel* channel, + const std::string& endpoint_id, + std::shared_ptr channel, EncryptionRunner::ResultListener listener) : client_(client), alarm_executor_(alarm_executor), endpoint_id_(endpoint_id), - channel_(channel), + weak_channel_(channel), listener_(std::move(listener)) {} void operator()() { + // Lock the weak pointer. If it fails, the channel was freed. + auto channel = weak_channel_.lock(); + // The IsClosed() check is to provide an early exit if channel has been + // closed. Otherwise the Read() and Write() calls on the channel below will + // return error and exit. + if (!channel || channel->IsClosed()) { + return; + } CancelableAlarm timeout_alarm( "EncryptionRunner.StartServer() timeout", - [this]() { CancelableAlarmRunnable(client_, endpoint_id_, channel_); }, + [this, weak_channel = weak_channel_]() { + if (auto channel = weak_channel.lock()) { + CancelableAlarmRunnable(client_, endpoint_id_, channel); + } + }, kTimeout, alarm_executor_); std::unique_ptr server = @@ -103,7 +116,7 @@ class ServerRunnable final { } // Message 1 (Client Init) - ExceptionOr client_init = channel_->Read(); + ExceptionOr client_init = channel->Read(); if (!client_init.ok()) { LogException(); HandleHandshakeOrIoException(&timeout_alarm); @@ -117,7 +130,7 @@ class ServerRunnable final { if (!parse_result.success) { LogException(); if (parse_result.alert_to_send != nullptr) { - HandleAlertException(parse_result); + HandleAlertException(parse_result, channel); } HandleHandshakeOrIoException(&timeout_alarm); return; @@ -137,7 +150,7 @@ class ServerRunnable final { return; } - Exception write_exception = channel_->Write(*server_init); + Exception write_exception = channel->Write(*server_init); if (!write_exception.Ok()) { LogException(); HandleHandshakeOrIoException(&timeout_alarm); @@ -148,7 +161,7 @@ class ServerRunnable final { << endpoint_id_ << ")."; // Message 3 (Client Finish) - ExceptionOr client_finish = channel_->Read(); + ExceptionOr client_finish = channel->Read(); if (!client_finish.ok()) { LogException(); @@ -163,7 +176,7 @@ class ServerRunnable final { if (!parse_result.success) { LogException(); if (parse_result.alert_to_send != nullptr) { - HandleAlertException(parse_result); + HandleAlertException(parse_result, channel); } HandleHandshakeOrIoException(&timeout_alarm); return; @@ -189,13 +202,13 @@ class ServerRunnable final { void HandleHandshakeOrIoException(CancelableAlarm* timeout_alarm) { timeout_alarm->Cancel(); - listener_.CallFailureCallback(endpoint_id_, channel_); + listener_.CallFailureCallback(endpoint_id_); } void HandleAlertException( - const securegcm::UKey2Handshake::ParseResult& parse_result) const { - Exception write_exception = - channel_->Write(*parse_result.alert_to_send); + const securegcm::UKey2Handshake::ParseResult& parse_result, + std::shared_ptr channel) const { + Exception write_exception = channel->Write(*parse_result.alert_to_send); if (!write_exception.Ok()) { LOG(WARNING) << "In StartServer(), client " << client_->GetClientId() << " failed to pass the alert error message to endpoint(id=" @@ -206,25 +219,39 @@ class ServerRunnable final { ClientProxy* client_; ScheduledExecutor* alarm_executor_; const std::string endpoint_id_; - EndpointChannel* channel_; + std::weak_ptr weak_channel_; EncryptionRunner::ResultListener listener_; }; class ClientRunnable final { public: ClientRunnable(ClientProxy* client, ScheduledExecutor* alarm_executor, - const std::string& endpoint_id, EndpointChannel* channel, + const std::string& endpoint_id, + std::shared_ptr channel, EncryptionRunner::ResultListener listener) : client_(client), alarm_executor_(alarm_executor), endpoint_id_(endpoint_id), - channel_(channel), + weak_channel_(channel), listener_(std::move(listener)) {} void operator()() { + // Lock the weak pointer. If it fails, the channel was freed. + auto channel = weak_channel_.lock(); + // The IsClosed() check is to provide an early exit if channel has been + // closed. Otherwise the Read() and Write() calls on the channel below will + // return error and exit. + if (!channel || channel->IsClosed()) { + return; + } + CancelableAlarm timeout_alarm( "EncryptionRunner.StartClient() timeout", - [this]() { CancelableAlarmRunnable(client_, endpoint_id_, channel_); }, + [this, weak_channel = weak_channel_]() { + if (auto channel = weak_channel.lock()) { + CancelableAlarmRunnable(client_, endpoint_id_, channel); + } + }, kTimeout, alarm_executor_); std::unique_ptr crypto = @@ -248,7 +275,7 @@ class ClientRunnable final { return; } - Exception write_init_exception = channel_->Write(*client_init); + Exception write_init_exception = channel->Write(*client_init); if (!write_init_exception.Ok()) { LogException(); HandleHandshakeOrIoException(&timeout_alarm); @@ -259,7 +286,7 @@ class ClientRunnable final { << endpoint_id_ << ")."; // Message 2 (Server Init) - ExceptionOr server_init = channel_->Read(); + ExceptionOr server_init = channel->Read(); if (!server_init.ok()) { LogException(); @@ -274,7 +301,7 @@ class ClientRunnable final { if (!parse_result.success) { LogException(); if (parse_result.alert_to_send != nullptr) { - HandleAlertException(parse_result); + HandleAlertException(parse_result, channel); } HandleHandshakeOrIoException(&timeout_alarm); return; @@ -294,8 +321,7 @@ class ClientRunnable final { return; } - Exception write_finish_exception = - channel_->Write(*client_finish); + Exception write_finish_exception = channel->Write(*client_finish); if (!write_finish_exception.Ok()) { LogException(); HandleHandshakeOrIoException(&timeout_alarm); @@ -322,12 +348,13 @@ class ClientRunnable final { void HandleHandshakeOrIoException(CancelableAlarm* timeout_alarm) { timeout_alarm->Cancel(); - listener_.CallFailureCallback(endpoint_id_, channel_); + listener_.CallFailureCallback(endpoint_id_); } void HandleAlertException( - const securegcm::UKey2Handshake::ParseResult& parse_result) const { - Exception write_exception = channel_->Write(*parse_result.alert_to_send); + const securegcm::UKey2Handshake::ParseResult& parse_result, + std::shared_ptr channel) const { + Exception write_exception = channel->Write(*parse_result.alert_to_send); if (!write_exception.Ok()) { LOG(WARNING) << "In StartClient(), client " << client_->GetClientId() << " failed to pass the alert error message to endpoint(id=" @@ -338,7 +365,7 @@ class ClientRunnable final { ClientProxy* client_; ScheduledExecutor* alarm_executor_; const std::string endpoint_id_; - EndpointChannel* channel_; + std::weak_ptr weak_channel_; EncryptionRunner::ResultListener listener_; }; @@ -346,19 +373,19 @@ class ClientRunnable final { EncryptionRunner::~EncryptionRunner() { Shutdown(); } -void EncryptionRunner::StartServer(ClientProxy* client, - const std::string& endpoint_id, - EndpointChannel* endpoint_channel, - EncryptionRunner::ResultListener listener) { +void EncryptionRunner::StartServer( + ClientProxy* client, const std::string& endpoint_id, + std::shared_ptr endpoint_channel, + EncryptionRunner::ResultListener listener) { ServerRunnable runnable(client, &alarm_executor_, endpoint_id, endpoint_channel, std::move(listener)); server_executor_.Execute("encryption-server", std::move(runnable)); } -void EncryptionRunner::StartClient(ClientProxy* client, - const std::string& endpoint_id, - EndpointChannel* endpoint_channel, - EncryptionRunner::ResultListener listener) { +void EncryptionRunner::StartClient( + ClientProxy* client, const std::string& endpoint_id, + std::shared_ptr endpoint_channel, + EncryptionRunner::ResultListener listener) { ClientRunnable runnable(client, &alarm_executor_, endpoint_id, endpoint_channel, std::move(listener)); client_executor_.Execute("encryption-client", std::move(runnable)); @@ -387,9 +414,9 @@ void EncryptionRunner::ResultListener::CallSuccessCallback( } void EncryptionRunner::ResultListener::CallFailureCallback( - const std::string& endpoint_id, EndpointChannel* channel) { + const std::string& endpoint_id) { if (on_failure_cb) { - std::move(on_failure_cb)(endpoint_id, channel); + std::move(on_failure_cb)(endpoint_id); } Reset(); } diff --git a/connections/implementation/encryption_runner.h b/connections/implementation/encryption_runner.h index e8e8c186..6e2236d7 100644 --- a/connections/implementation/encryption_runner.h +++ b/connections/implementation/encryption_runner.h @@ -45,8 +45,7 @@ class EncryptionRunner { std::unique_ptr ukey2, const std::string& auth_token, const ByteArray& raw_auth_token); - void CallFailureCallback(const std::string& endpoint_id, - EndpointChannel* channel); + void CallFailureCallback(const std::string& endpoint_id); void Reset(); // @EncryptionRunnerThread @@ -56,27 +55,19 @@ class EncryptionRunner { const ByteArray& raw_auth_token) &&> on_success_cb; - // Encryption has failed. The remote_endpoint_id and channel are given so - // that any pending state can be cleaned up. - // - // We return the EndpointChannel because, at this stage, simultaneous - // connections are a possibility. Use this channel to verify that the state - // you're cleaning up is for this EndpointChannel, and not state for another - // channel to the same endpoint. + // Encryption has failed. // // @EncryptionRunnerThread - absl::AnyInvocable - on_failure_cb; + absl::AnyInvocable on_failure_cb; }; // @AnyThread void StartServer(ClientProxy* client, const std::string& endpoint_id, - EndpointChannel* endpoint_channel, + std::shared_ptr endpoint_channel, ResultListener result_listener); // @AnyThread void StartClient(ClientProxy* client, const std::string& endpoint_id, - EndpointChannel* endpoint_channel, + std::shared_ptr endpoint_channel, ResultListener result_listener); // @AnyThread diff --git a/connections/implementation/encryption_runner_test.cc b/connections/implementation/encryption_runner_test.cc index 5cf31250..f3cd2bd0 100644 --- a/connections/implementation/encryption_runner_test.cc +++ b/connections/implementation/encryption_runner_test.cc @@ -123,21 +123,24 @@ class FakeEndpointChannel : public EndpointChannel { }; struct User { - User(InputStream* reader, OutputStream* writer) : channel(reader, writer) {} + User(InputStream* reader, OutputStream* writer) + : channel(std::make_shared(reader, writer)) {} - FakeEndpointChannel channel; + std::shared_ptr channel; EncryptionRunner crypto; ClientProxy client; }; struct Response { + Response() : latch(2) {} + explicit Response(int count) : latch(count) {} enum class Status { kUnknown = 0, kDone = 1, kFailed = 2, }; - CountDownLatch latch{2}; + CountDownLatch latch; Status server_status = Status::kUnknown; Status client_status = Status::kUnknown; }; @@ -154,7 +157,7 @@ TEST(EncryptionRunnerTest, ReadWrite) { Response response; user_a.crypto.StartServer( - &user_a.client, "endpoint_id", &user_a.channel, + &user_a.client, "endpoint_id", user_a.channel, { .on_success_cb = [&response](const std::string& endpoint_id, @@ -165,15 +168,14 @@ TEST(EncryptionRunnerTest, ReadWrite) { response.latch.CountDown(); }, .on_failure_cb = - [&response](const std::string& endpoint_id, - EndpointChannel* channel) { - channel->Close(); + [&response, &user_a](const std::string& endpoint_id) { + user_a.channel->Close(); response.server_status = Response::Status::kFailed; response.latch.CountDown(); }, }); user_b.crypto.StartClient( - &user_b.client, "endpoint_id", &user_b.channel, + &user_b.client, "endpoint_id", user_b.channel, { .on_success_cb = [&response](const std::string& endpoint_id, @@ -184,9 +186,8 @@ TEST(EncryptionRunnerTest, ReadWrite) { response.latch.CountDown(); }, .on_failure_cb = - [&response](const std::string& endpoint_id, - EndpointChannel* channel) { - channel->Close(); + [&response, &user_b](const std::string& endpoint_id) { + user_b.channel->Close(); response.client_status = Response::Status::kFailed; response.latch.CountDown(); }, @@ -203,14 +204,13 @@ TEST(EncryptionRunnerTest, ClientWriteFails) { /*writer=*/from_a_to_b.second.get()); User user_b(/*reader=*/from_a_to_b.first.get(), /*writer=*/from_b_to_a.second.get()); - Response response; - response.latch = CountDownLatch(1); + Response response(1); // Close server's input stream, so client can't write to it. from_b_to_a.first->Close(); user_b.crypto.StartClient( - &user_b.client, "endpoint_id", &user_b.channel, + &user_b.client, "endpoint_id", user_b.channel, { .on_success_cb = [&response](const std::string& endpoint_id, @@ -221,9 +221,8 @@ TEST(EncryptionRunnerTest, ClientWriteFails) { response.latch.CountDown(); }, .on_failure_cb = - [&response](const std::string& endpoint_id, - EndpointChannel* channel) { - channel->Close(); + [&response, &user_a](const std::string& endpoint_id) { + user_a.channel->Close(); response.client_status = Response::Status::kFailed; response.latch.CountDown(); }, @@ -239,14 +238,13 @@ TEST(EncryptionRunnerTest, ServerWriteFails) { /*writer=*/from_a_to_b.second.get()); User user_b(/*reader=*/from_a_to_b.first.get(), /*writer=*/from_b_to_a.second.get()); - Response response; - response.latch = CountDownLatch(1); + Response response(1); // Close client's input stream, so server can't write to it. from_a_to_b.first->Close(); user_a.crypto.StartServer( - &user_a.client, "endpoint_id", &user_a.channel, + &user_a.client, "endpoint_id", user_a.channel, { .on_success_cb = [&response](const std::string& endpoint_id, @@ -257,24 +255,22 @@ TEST(EncryptionRunnerTest, ServerWriteFails) { response.latch.CountDown(); }, .on_failure_cb = - [&response](const std::string& endpoint_id, - EndpointChannel* channel) { - channel->Close(); + [&response, &user_a](const std::string& endpoint_id) { + user_a.channel->Close(); response.server_status = Response::Status::kFailed; response.latch.CountDown(); }, }); user_b.crypto.StartClient( - &user_b.client, "endpoint_id", &user_b.channel, + &user_b.client, "endpoint_id", user_b.channel, { - .on_success_cb = - [](const std::string& endpoint_id, - std::unique_ptr ukey2, - const std::string& auth_token, - const ByteArray& raw_auth_token) {}, + .on_success_cb = [](const std::string& endpoint_id, + std::unique_ptr ukey2, + const std::string& auth_token, + const ByteArray& raw_auth_token) {}, .on_failure_cb = - [](const std::string& endpoint_id, EndpointChannel* channel) { - channel->Close(); + [&user_b](const std::string& endpoint_id) { + user_b.channel->Close(); }, }); EXPECT_TRUE(response.latch.Await(absl::Milliseconds(5000)).result()); @@ -286,11 +282,10 @@ TEST(EncryptionRunnerTest, ClientSendsGarbageMessage1) { auto from_client_to_server = CreatePipe(); User user_a(/*reader=*/from_client_to_server.first.get(), /*writer=*/from_server_to_client.second.get()); - Response response; - response.latch = CountDownLatch(1); + Response response(1); user_a.crypto.StartServer( - &user_a.client, "endpoint_id", &user_a.channel, + &user_a.client, "endpoint_id", user_a.channel, { .on_success_cb = [&response](const std::string& endpoint_id, @@ -301,9 +296,8 @@ TEST(EncryptionRunnerTest, ClientSendsGarbageMessage1) { response.latch.CountDown(); }, .on_failure_cb = - [&response](const std::string& endpoint_id, - EndpointChannel* channel) { - channel->Close(); + [&response, &user_a](const std::string& endpoint_id) { + user_a.channel->Close(); response.server_status = Response::Status::kFailed; response.latch.CountDown(); }, @@ -327,11 +321,10 @@ TEST(EncryptionRunnerTest, ServerSendsGarbageMessage2) { auto from_client_to_server = CreatePipe(); User user_b(/*reader=*/from_server_to_client.first.get(), /*writer=*/from_client_to_server.second.get()); - Response response; - response.latch = CountDownLatch(1); + Response response(1); user_b.crypto.StartClient( - &user_b.client, "endpoint_id", &user_b.channel, + &user_b.client, "endpoint_id", user_b.channel, { .on_success_cb = [&response](const std::string& endpoint_id, @@ -342,9 +335,8 @@ TEST(EncryptionRunnerTest, ServerSendsGarbageMessage2) { response.latch.CountDown(); }, .on_failure_cb = - [&response](const std::string& endpoint_id, - EndpointChannel* channel) { - channel->Close(); + [&response, &user_b](const std::string& endpoint_id) { + user_b.channel->Close(); response.client_status = Response::Status::kFailed; response.latch.CountDown(); }, @@ -373,11 +365,10 @@ TEST(EncryptionRunnerTest, ClientSendsGarbageMessage3) { /*writer=*/from_server_to_client.second.get()); User user_b(/*reader=*/from_server_to_client.first.get(), /*writer=*/from_client_to_server.second.get()); - Response response; - response.latch = CountDownLatch(1); + Response response(1); user_a.crypto.StartServer( - &user_a.client, "endpoint_id", &user_a.channel, + &user_a.client, "endpoint_id", user_a.channel, { .on_success_cb = [&response](const std::string& endpoint_id, @@ -388,9 +379,8 @@ TEST(EncryptionRunnerTest, ClientSendsGarbageMessage3) { response.latch.CountDown(); }, .on_failure_cb = - [&response](const std::string& endpoint_id, - EndpointChannel* channel) { - channel->Close(); + [&response, &user_a](const std::string& endpoint_id) { + user_a.channel->Close(); response.server_status = Response::Status::kFailed; response.latch.CountDown(); }, @@ -410,7 +400,8 @@ TEST(EncryptionRunnerTest, ClientSendsGarbageMessage3) { EXPECT_TRUE(server_init.ok()); // Client crypto parses message 2. - client_crypto->ParseHandshakeMessage(std::string(server_init.result())); + client_crypto->ParseHandshakeMessage( + std::string(server_init.result().data(), server_init.result().size())); // Client sends garbage instead of message 3 from_client_to_server.second->Write("Garbage"); diff --git a/connections/implementation/endpoint_channel_manager.cc b/connections/implementation/endpoint_channel_manager.cc index b093b6f3..c6af41d4 100644 --- a/connections/implementation/endpoint_channel_manager.cc +++ b/connections/implementation/endpoint_channel_manager.cc @@ -46,7 +46,7 @@ EndpointChannelManager::~EndpointChannelManager() { void EndpointChannelManager::RegisterChannelForEndpoint( ClientProxy* client, const std::string& endpoint_id, - std::unique_ptr channel) { + std::shared_ptr channel) { MutexLock lock(&mutex_); LOG(INFO) << "EndpointChannelManager registered channel of type " @@ -59,7 +59,7 @@ void EndpointChannelManager::RegisterChannelForEndpoint( void EndpointChannelManager::ReplaceChannelForEndpoint( ClientProxy* client, const std::string& endpoint_id, - std::unique_ptr channel, bool enable_encryption) { + std::shared_ptr channel, bool enable_encryption) { MutexLock lock(&mutex_); if (client->IsSafeToDisconnectEnabled(endpoint_id) && channel_state_.IsWaitingForSafeToDisconnectTimeout(endpoint_id)) { @@ -106,7 +106,7 @@ std::shared_ptr EndpointChannelManager::GetChannelForEndpoint( void EndpointChannelManager::SetActiveEndpointChannel( ClientProxy* client, const std::string& endpoint_id, - std::unique_ptr channel, bool enable_encryption) { + 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); @@ -189,7 +189,7 @@ void EndpointChannelManager::ChannelState::DestroyAll() { } void EndpointChannelManager::ChannelState::UpdateChannelForEndpoint( - const std::string& endpoint_id, std::unique_ptr channel) { + const std::string& endpoint_id, std::shared_ptr channel) { // Create EndpointData instance, if necessary, and populate channel. endpoints_[endpoint_id].channel = std::move(channel); } diff --git a/connections/implementation/endpoint_channel_manager.h b/connections/implementation/endpoint_channel_manager.h index cbc5556a..19183ee0 100644 --- a/connections/implementation/endpoint_channel_manager.h +++ b/connections/implementation/endpoint_channel_manager.h @@ -59,7 +59,7 @@ class EndpointChannelManager final { // be closed before continuing the registration. void RegisterChannelForEndpoint(ClientProxy* client, const std::string& endpoint_id, - std::unique_ptr channel) + std::shared_ptr channel) ABSL_LOCKS_EXCLUDED(mutex_); // Replaces the EndpointChannel to be associated with an endpoint from here on @@ -67,7 +67,7 @@ class EndpointChannelManager final { // to the newly-provided EndpointChannel. void ReplaceChannelForEndpoint(ClientProxy* client, const std::string& endpoint_id, - std::unique_ptr channel, + std::shared_ptr channel, bool enable_encryption) ABSL_LOCKS_EXCLUDED(mutex_); @@ -168,7 +168,7 @@ class EndpointChannelManager final { // Stores a new EndpointChannel for the endpoint. // Prevoius one is destroyed, if it existed. void UpdateChannelForEndpoint(const std::string& endpoint_id, - std::unique_ptr channel); + std::shared_ptr channel); // Stores a new EncryptionContext for the endpoint. // Prevoius one is destroyed, if it existed. @@ -207,7 +207,7 @@ class EndpointChannelManager final { void SetActiveEndpointChannel(ClientProxy* client, const std::string& endpoint_id, - std::unique_ptr channel, + std::shared_ptr channel, bool enable_encryption) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); diff --git a/connections/implementation/endpoint_channel_manager_test.cc b/connections/implementation/endpoint_channel_manager_test.cc index e5ada1ec..d73b746b 100644 --- a/connections/implementation/endpoint_channel_manager_test.cc +++ b/connections/implementation/endpoint_channel_manager_test.cc @@ -110,8 +110,8 @@ std::function MakeDataMonitor(absl::string_view label, std::pair, std::unique_ptr> -DoDhKeyExchange(BaseEndpointChannel* channel_a, - BaseEndpointChannel* channel_b) { +DoDhKeyExchange(std::shared_ptr channel_a, + std::shared_ptr channel_b) { std::unique_ptr context_a; std::unique_ptr context_b; EncryptionRunner crypto_a; @@ -136,8 +136,7 @@ DoDhKeyExchange(BaseEndpointChannel* channel_a, latch.CountDown(); }, .on_failure_cb = - [&latch](const std::string& endpoint_id, - EndpointChannel* channel) { + [&latch](const std::string& endpoint_id) { LOG(INFO) << "client-A side key negotiation failed"; latch.CountDown(); }, @@ -159,8 +158,7 @@ DoDhKeyExchange(BaseEndpointChannel* channel_a, latch.CountDown(); }, .on_failure_cb = - [&latch](const std::string& endpoint_id, - EndpointChannel* channel) { + [&latch](const std::string& endpoint_id) { LOG(INFO) << "client-B side key negotiation failed"; latch.CountDown(); }, @@ -185,9 +183,9 @@ TEST(BaseEndpointChannelManagerTest, RegisterChannelEncryptedReadwrite) { // to server "b". auto server_b = CreatePipe(); // Data pump "b" reads from client "b", writes // to server "a". - auto channel_a = std::make_unique(server_a.first.get(), + auto channel_a = std::make_shared(server_a.first.get(), client_a.second.get()); - auto channel_b = std::make_unique(server_b.first.get(), + auto channel_b = std::make_shared(server_b.first.get(), client_b.second.get()); auto channel_a_raw = channel_a.get(); auto channel_b_raw = channel_b.get(); @@ -208,7 +206,7 @@ TEST(BaseEndpointChannelManagerTest, RegisterChannelEncryptedReadwrite) { MakeDataMonitor(kMonitorB, &capture_b, &mutex))); // Run DH key exchange; setup encryption contexts for channels. - auto context = DoDhKeyExchange(channel_a.get(), channel_b.get()); + auto context = DoDhKeyExchange(channel_a, channel_b); ASSERT_NE(context.first, nullptr); ASSERT_NE(context.second, nullptr); @@ -266,9 +264,9 @@ TEST(BaseEndpointChannelManagerTest, ReplaceChannelNoEncrypted) { // to server "b". auto server_b = CreatePipe(); // Data pump "b" reads from client "b", writes // to server "a". - auto channel_a = std::make_unique(server_a.first.get(), + auto channel_a = std::make_shared(server_a.first.get(), client_a.second.get()); - auto channel_b = std::make_unique(server_b.first.get(), + auto channel_b = std::make_shared(server_b.first.get(), client_b.second.get()); auto channel_a_raw = channel_a.get(); auto channel_b_raw = channel_b.get(); @@ -289,7 +287,7 @@ TEST(BaseEndpointChannelManagerTest, ReplaceChannelNoEncrypted) { MakeDataMonitor(kMonitorB, &capture_b, &mutex))); // Run DH key exchange; setup encryption contexts for channels. - auto context = DoDhKeyExchange(channel_a.get(), channel_b.get()); + auto context = DoDhKeyExchange(channel_a, channel_b); ASSERT_NE(context.first, nullptr); ASSERT_NE(context.second, nullptr); diff --git a/connections/implementation/endpoint_manager.cc b/connections/implementation/endpoint_manager.cc index a93f4688..c0e94b8b 100644 --- a/connections/implementation/endpoint_manager.cc +++ b/connections/implementation/endpoint_manager.cc @@ -539,103 +539,94 @@ void EndpointManager::RegisterEndpoint( ClientProxy* client, const std::string& endpoint_id, const ConnectionResponseInfo& info, const ConnectionOptions& connection_options, - std::unique_ptr channel, + std::shared_ptr channel, const ConnectionListener& listener, const std::string& connection_token) { CountDownLatch latch(1); - // NOTE (unique_ptr<> capture): - // std::unique_ptr<> is not copyable, so we can not pass it to - // lambda capture, because lambda eventually is converted to - // std::function<>. Instead, we release() a pointer, and pass a raw pointer, - // which is copyalbe. We ignore the risk of job not scheduled (and an - // associated risk of memory leak), because this may only happen during - // service shutdown. - RunOnEndpointManagerThread( - "register-endpoint", - [this, client, channel = channel.release(), &endpoint_id, &info, - &connection_options, &listener, &connection_token, &latch]() { - if (endpoints_.contains(endpoint_id)) { - LOG(WARNING) << "Registering duplicate endpoint " << endpoint_id; - // We must remove old endpoint state before registering a new one - // for the same endpoint_id. - RemoveEndpointState(endpoint_id); - } + RunOnEndpointManagerThread("register-endpoint", [this, client, channel, + &endpoint_id, &info, + &connection_options, + &listener, &connection_token, + &latch]() { + if (endpoints_.contains(endpoint_id)) { + LOG(WARNING) << "Registering duplicate endpoint " << endpoint_id; + // We must remove old endpoint state before registering a new one + // for the same endpoint_id. + RemoveEndpointState(endpoint_id); + } - absl::Duration keep_alive_interval = - absl::Milliseconds(connection_options.keep_alive_interval_millis); - absl::Duration keep_alive_timeout = - absl::Milliseconds(connection_options.keep_alive_timeout_millis); - LOG(INFO) << "Registering endpoint " << endpoint_id << " for client " - << client->GetClientId() - << " with keep-alive frame as interval=" - << absl::FormatDuration(keep_alive_interval) - << ", timeout=" << absl::FormatDuration(keep_alive_timeout); + absl::Duration keep_alive_interval = + absl::Milliseconds(connection_options.keep_alive_interval_millis); + absl::Duration keep_alive_timeout = + absl::Milliseconds(connection_options.keep_alive_timeout_millis); + LOG(INFO) << "Registering endpoint " << endpoint_id << " for client " + << client->GetClientId() << " with keep-alive frame as interval=" + << absl::FormatDuration(keep_alive_interval) + << ", timeout=" << absl::FormatDuration(keep_alive_timeout); - // Pass ownership of channel to EndpointChannelManager - LOG(INFO) << "Registering endpoint with channel manager: endpoint " - << endpoint_id; - channel_manager_->RegisterChannelForEndpoint( - client, endpoint_id, std::unique_ptr(channel)); + // Pass ownership of channel to EndpointChannelManager + LOG(INFO) << "Registering endpoint with channel manager: endpoint " + << endpoint_id; + channel_manager_->RegisterChannelForEndpoint(client, endpoint_id, channel); - EndpointState& endpoint_state = - endpoints_ - .emplace(endpoint_id, - EndpointState(endpoint_id, channel_manager_)) - .first->second; + EndpointState& endpoint_state = + endpoints_ + .emplace(endpoint_id, EndpointState(endpoint_id, channel_manager_)) + .first->second; - LOG(INFO) << "Starting workers: endpoint " << endpoint_id; - // For every endpoint, there's normally only one Read handler instance - // running on a dedicated thread. This instance reads data from the - // endpoint and delegates incoming frames to various FrameProcessors. - // Once the frame has been properly handled, it starts reading again - // for the next frame. If the handler fails its read and no other - // EndpointChannels are available for this endpoint, a disconnection - // will be initiated. - endpoint_state.StartEndpointReader([this, client, endpoint_id]() { + LOG(INFO) << "Starting workers: endpoint " << endpoint_id; + // For every endpoint, there's normally only one Read handler instance + // running on a dedicated thread. This instance reads data from the + // endpoint and delegates incoming frames to various FrameProcessors. + // Once the frame has been properly handled, it starts reading again + // for the next frame. If the handler fails its read and no other + // EndpointChannels are available for this endpoint, a disconnection + // will be initiated. + endpoint_state.StartEndpointReader([this, client, endpoint_id]() { + EndpointChannelLoopRunnable( + "Read", client, endpoint_id, + [this, client, endpoint_id](EndpointChannel* channel) { + return HandleData(endpoint_id, client, channel); + }); + }); + + // For every endpoint, there's only one KeepAliveManager instance + // running on a dedicated thread. This instance will periodically send + // out a ping* to the endpoint while listening for an incoming pong**. + // If it fails to send the ping, or if no pong is heard within + // keep_alive_timeout, it initiates a disconnection. + // + // (*) Bluetooth requires a constant outgoing stream of messages. If + // there's silence, Android will break the socket. This is why we + // ping. + // (**) 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. + VLOG(1) << "EndpointManager enabling KeepAlive for endpoint " + << endpoint_id; + endpoint_state.StartEndpointKeepAliveManager( + [this, client, endpoint_id, keep_alive_interval, keep_alive_timeout]( + Mutex* keep_alive_waiter_mutex, + ConditionVariable* keep_alive_waiter) { EndpointChannelLoopRunnable( - "Read", client, endpoint_id, - [this, client, endpoint_id](EndpointChannel* channel) { - return HandleData(endpoint_id, client, channel); + "KeepAliveManager", client, endpoint_id, + [this, keep_alive_interval, keep_alive_timeout, + keep_alive_waiter_mutex, + keep_alive_waiter](EndpointChannel* channel) { + return HandleKeepAlive( + channel, keep_alive_interval, keep_alive_timeout, + keep_alive_waiter_mutex, keep_alive_waiter); }); }); + LOG(INFO) << "Registering endpoint " << endpoint_id + << ", workers started and notifying client."; - // For every endpoint, there's only one KeepAliveManager instance - // running on a dedicated thread. This instance will periodically send - // out a ping* to the endpoint while listening for an incoming pong**. - // If it fails to send the ping, or if no pong is heard within - // keep_alive_timeout, it initiates a disconnection. - // - // (*) Bluetooth requires a constant outgoing stream of messages. If - // there's silence, Android will break the socket. This is why we - // ping. - // (**) 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. - VLOG(1) << "EndpointManager enabling KeepAlive for endpoint " - << endpoint_id; - endpoint_state.StartEndpointKeepAliveManager( - [this, client, endpoint_id, keep_alive_interval, - keep_alive_timeout](Mutex* keep_alive_waiter_mutex, - ConditionVariable* keep_alive_waiter) { - EndpointChannelLoopRunnable( - "KeepAliveManager", client, endpoint_id, - [this, keep_alive_interval, keep_alive_timeout, - keep_alive_waiter_mutex, - keep_alive_waiter](EndpointChannel* channel) { - return HandleKeepAlive( - channel, keep_alive_interval, keep_alive_timeout, - keep_alive_waiter_mutex, keep_alive_waiter); - }); - }); - LOG(INFO) << "Registering endpoint " << endpoint_id - << ", workers started and notifying client."; - - // It's now time to let the client know of this new connection so that - // they can accept or reject it. - client->OnConnectionInitiated(endpoint_id, info, connection_options, - listener, connection_token); - latch.CountDown(); - }); + // It's now time to let the client know of this new connection so that + // they can accept or reject it. + client->OnConnectionInitiated(endpoint_id, info, connection_options, + listener, connection_token); + latch.CountDown(); + }); latch.Await(); } @@ -722,7 +713,7 @@ void EndpointManager::DiscardEndpoint(ClientProxy* client, // of `serial_executor_` and will still have access to a valid // `is_shutdown_`. // - // TODO(b/280653613): Develop a more robost solution to prevent + // TODO(b/280653613): Develop a more robust solution to prevent // accessing an already destroyed `ClientProxy` during destruction. { MutexLock lock(&mutex_); @@ -953,8 +944,7 @@ std::vector EndpointManager::SendTransferFrameBytes( continue; } - Exception write_exception = - channel->Write(bytes, packet_meta_data); + Exception write_exception = channel->Write(bytes, packet_meta_data); if (!write_exception.Ok()) { failed_endpoint_ids.push_back(endpoint_id); LOG(INFO) << "Failed to send packet; endpoint_id=" << endpoint_id; diff --git a/connections/implementation/endpoint_manager.h b/connections/implementation/endpoint_manager.h index 68f7f4f9..4b47f9c5 100644 --- a/connections/implementation/endpoint_manager.h +++ b/connections/implementation/endpoint_manager.h @@ -25,6 +25,7 @@ #include "absl/container/flat_hash_map.h" #include "absl/functional/any_invocable.h" #include "absl/time/time.h" +#include "connections/connection_options.h" #include "connections/implementation/analytics/packet_meta_data.h" #include "connections/implementation/client_proxy.h" #include "connections/implementation/endpoint_channel.h" @@ -34,6 +35,8 @@ #include "internal/platform/byte_array.h" #include "internal/platform/condition_variable.h" #include "internal/platform/count_down_latch.h" +#include "internal/platform/exception.h" +#include "internal/platform/mutex.h" #include "internal/platform/runnable.h" #include "internal/platform/single_thread_executor.h" @@ -112,7 +115,7 @@ class EndpointManager { void RegisterEndpoint(ClientProxy* client, const std::string& endpoint_id, const ConnectionResponseInfo& info, const ConnectionOptions& connection_options, - std::unique_ptr channel, + std::shared_ptr channel, const ConnectionListener& listener, const std::string& connection_token); // Called when a client explicitly asks to disconnect from this endpoint. In