From 7c7277416787986d10220b37e6f014e6d686088d Mon Sep 17 00:00:00 2001 From: hai007 Date: Fri, 10 May 2024 16:09:17 -0700 Subject: [PATCH] Multiplex implementation - Create Multiplex Socket PiperOrigin-RevId: 632628955 --- Package.swift | 1 + connections/implementation/BUILD | 1 + .../implementation/base_pcp_handler.cc | 6 +- .../implementation/base_pcp_handler_test.cc | 4 +- connections/implementation/client_proxy.cc | 53 ++ connections/implementation/client_proxy.h | 26 + .../flags/nearby_connections_feature_flags.h | 4 + .../implementation/mediums/multiplex/BUILD | 8 + .../multiplex/multiplex_output_stream.cc | 37 +- .../multiplex/multiplex_output_stream.h | 11 +- .../multiplex/multiplex_output_stream_test.cc | 4 +- .../mediums/multiplex/multiplex_socket.cc | 743 ++++++++++++++++++ .../mediums/multiplex/multiplex_socket.h | 214 +++++ .../multiplex/multiplex_socket_test.cc | 361 +++++++++ connections/implementation/mediums/utils.cc | 20 +- connections/implementation/mediums/utils.h | 4 +- connections/implementation/offline_frames.cc | 5 +- connections/implementation/offline_frames.h | 3 +- .../implementation/offline_frames_test.cc | 9 +- .../offline_frames_validator_test.cc | 9 +- internal/platform/BUILD | 1 + internal/platform/feature_flags.h | 3 + internal/platform/socket.h | 75 ++ 23 files changed, 1560 insertions(+), 42 deletions(-) create mode 100644 connections/implementation/mediums/multiplex/multiplex_socket.cc create mode 100644 connections/implementation/mediums/multiplex/multiplex_socket.h create mode 100644 connections/implementation/mediums/multiplex/multiplex_socket_test.cc diff --git a/Package.swift b/Package.swift index 67012bc4..8e018c46 100644 --- a/Package.swift +++ b/Package.swift @@ -463,6 +463,7 @@ let package = Package( "connections/implementation/mediums/ble_v2/discovered_peripheral_tracker_test.cc", "connections/implementation/mediums/ble_v2/instant_on_lost_advertisement_test.cc", "connections/implementation/mediums/multiplex/multiplex_frames_test.cc", + "connections/implementation/mediums/multiplex/multiplex_socket_test.cc", "connections/implementation/mediums/multiplex/multiplex_output_stream_test.cc", "connections/implementation/mediums/webrtc_peer_id_test.cc", "connections/implementation/mediums/wifi_lan_test.cc", diff --git a/connections/implementation/BUILD b/connections/implementation/BUILD index 3ba0201a..512790f0 100644 --- a/connections/implementation/BUILD +++ b/connections/implementation/BUILD @@ -109,6 +109,7 @@ cc_library( "//chrome/chromeos/assistant/data_migration/lib:__pkg__", "//connections:__pkg__", "//connections/implementation/fuzzers:__pkg__", + "//connections/implementation/mediums/multiplex:__pkg__", "//sharing:__subpackages__", ], deps = [ diff --git a/connections/implementation/base_pcp_handler.cc b/connections/implementation/base_pcp_handler.cc index 711f208f..338f0617 100644 --- a/connections/implementation/base_pcp_handler.cc +++ b/connections/implementation/base_pcp_handler.cc @@ -1338,7 +1338,8 @@ Status BasePcpHandler::AcceptConnection(ClientProxy* client, Exception write_exception = channel->Write(parser::ForConnectionResponse( - Status::kSuccess, client->GetLocalOsInfo())); + Status::kSuccess, client->GetLocalOsInfo(), + client->GetLocalMultiplexSocketBitmask())); if (!write_exception.Ok()) { NEARBY_LOGS(INFO) << "AcceptConnection: failed to send response: endpoint_id=" @@ -1400,7 +1401,8 @@ Status BasePcpHandler::RejectConnection(ClientProxy* client, Exception write_exception = channel->Write(parser::ForConnectionResponse( - Status::kConnectionRejected, client->GetLocalOsInfo())); + Status::kConnectionRejected, client->GetLocalOsInfo(), + client->GetLocalMultiplexSocketBitmask())); if (!write_exception.Ok()) { NEARBY_LOGS(INFO) << "RejectConnection: failed to send response: endpoint_id=" diff --git a/connections/implementation/base_pcp_handler_test.cc b/connections/implementation/base_pcp_handler_test.cc index 314f8ad5..5c108795 100644 --- a/connections/implementation/base_pcp_handler_test.cc +++ b/connections/implementation/base_pcp_handler_test.cc @@ -1497,8 +1497,8 @@ TEST_P(BasePcpHandlerTest, OnIncomingFrameChangesState) { Status{Status::kSuccess}); NEARBY_LOG(INFO, "Simulating remote accept: id=%s", endpoint_id.c_str()); OsInfo os_info; - auto frame = parser::FromBytes( - parser::ForConnectionResponse(Status::kSuccess, os_info)); + auto frame = parser::FromBytes(parser::ForConnectionResponse( + Status::kSuccess, os_info, /*multiplex_socket_bitmask=*/0)); EXPECT_CALL(mock_connection_listener_.bandwidth_changed_cb, Call).Times(1); pcp_handler.OnIncomingFrame(frame.result(), endpoint_id, &client, connect_medium, packet_meta_data); diff --git a/connections/implementation/client_proxy.cc b/connections/implementation/client_proxy.cc index a6fd5c6e..75a93f5d 100644 --- a/connections/implementation/client_proxy.cc +++ b/connections/implementation/client_proxy.cc @@ -30,13 +30,17 @@ #include "absl/functional/any_invocable.h" #include "absl/strings/escaping.h" #include "absl/strings/string_view.h" +#include "connections/advertising_options.h" +#include "connections/discovery_options.h" #include "connections/implementation/flags/nearby_connections_feature_flags.h" #include "connections/listeners.h" +#include "connections/medium_selector.h" #include "connections/v3/bandwidth_info.h" #include "connections/v3/connection_listening_options.h" #include "connections/v3/connections_device_provider.h" #include "internal/analytics/event_logger.h" #include "internal/flags/nearby_flags.h" +#include "internal/platform/cancelable_alarm.h" #include "internal/platform/error_code_recorder.h" #include "internal/platform/feature_flags.h" #include "internal/platform/implementation/platform.h" @@ -1102,6 +1106,55 @@ OsInfo::OsType ClientProxy::OSNameToOsInfoType(api::OSName osName) { } } +std::int32_t ClientProxy::GetLocalMultiplexSocketBitmask() const { + if (NearbyFlags::GetInstance().GetBoolFlag( + config_package_nearby::nearby_connections_feature:: + kEnableMultiplex)) { + return kBtMultiplexEnabled; + } + return 0; +} + +void ClientProxy::SetRemoteMultiplexSocketBitmask( + absl::string_view endpoint_id, int remote_multiplex_socket_bitmask) { + MutexLock lock(&mutex_); + ConnectionPair* item = LookupConnection(endpoint_id); + if (item != nullptr) { + item->first.remote_multiplex_socket_bitmask = + remote_multiplex_socket_bitmask; + } +} + +std::optional ClientProxy::GetRemoteMultiplexSocketBitmask( + absl::string_view endpoint_id) const { + MutexLock lock(&mutex_); + const ConnectionPair* item = LookupConnection(endpoint_id); + if (item != nullptr) { + return item->first.remote_multiplex_socket_bitmask; + } + return std::nullopt; +} +bool ClientProxy::IsMultiplexSocketSupported(absl::string_view endpoint_id, + Medium medium) { + MutexLock lock(&mutex_); + ConnectionPair* item = LookupConnection(endpoint_id); + if (item == nullptr) { + return false; + } + + int combined_result = GetLocalMultiplexSocketBitmask() & + item->first.remote_multiplex_socket_bitmask; + switch (medium) { + case Medium::BLUETOOTH: + return (combined_result & kBtMultiplexEnabled) != 0; + case Medium::WIFI_LAN: + return (combined_result & kWifiLanMultiplexEnabled) != 0; + default: + return false; + } +} + + std::string ClientProxy::ToString(PayloadProgressInfo::Status status) const { switch (status) { case PayloadProgressInfo::Status::kSuccess: diff --git a/connections/implementation/client_proxy.h b/connections/implementation/client_proxy.h index 00067c0a..4d8e409b 100644 --- a/connections/implementation/client_proxy.h +++ b/connections/implementation/client_proxy.h @@ -25,10 +25,12 @@ #include "absl/functional/any_invocable.h" #include "absl/strings/string_view.h" #include "connections/advertising_options.h" +#include "connections/connection_options.h" #include "connections/discovery_options.h" #include "connections/implementation/analytics/analytics_recorder.h" #include "connections/implementation/proto/offline_wire_formats.pb.h" #include "connections/listeners.h" +#include "connections/medium_selector.h" #include "connections/status.h" #include "connections/strategy.h" #include "connections/v3/connection_listening_options.h" @@ -296,6 +298,17 @@ class ClientProxy final { bool IsAutoReconnectEnabled(absl::string_view endpoint_id); bool IsPayloadReceivedAckEnabled(absl::string_view endpoint_id); + // Returns the multiplex socket supports status for local device. + std::int32_t GetLocalMultiplexSocketBitmask() const; + // Sets the multiplex socket supports status for remote device. + void SetRemoteMultiplexSocketBitmask(absl::string_view endpoint_id, + int remote_multiplex_socket_bitmask); + // Gets the multiplex socket supports status for remote device. + std::optional GetRemoteMultiplexSocketBitmask( + absl::string_view endpoint_id) const; + // Returns true if the multiplex socket is supported for the given medium. + bool IsMultiplexSocketSupported(absl::string_view endpoint_id, Medium medium); + private: struct Connection { // Status: may be either: @@ -326,6 +339,7 @@ class ClientProxy final { std::string connection_token; std::optional os_info; std::int32_t safe_to_disconnect_version; + std::int32_t remote_multiplex_socket_bitmask; }; using ConnectionPair = std::pair; @@ -463,6 +477,18 @@ class ClientProxy final { bool supports_safe_to_disconnect_; bool support_auto_reconnect_; std::int32_t local_safe_to_disconnect_version_; + + /** Bitmask for bt multiplex connection support. */ + // Note. Deprecates the first and second bit of BT_MULTIPLEX_ENABLED and + // WIFI_LAN_MULTIPLEX_ENABLED and shift them to the third and the forth bit. + // The reason is we need to escape the (0, 1) bit which has been set in some + // devices without salt enabled. If accompany with the devices with salted + // enabled, the frames passed cannot be decrypted and the connection shall be + // failed. Please refer to b/295925531#comment#14 for the details. + enum MultiplexSocketBitmask : uint32_t { + kBtMultiplexEnabled = 1 << 2, + kWifiLanMultiplexEnabled = 1 << 3, + }; }; } // namespace connections diff --git a/connections/implementation/flags/nearby_connections_feature_flags.h b/connections/implementation/flags/nearby_connections_feature_flags.h index b91ec394..580ee9d3 100644 --- a/connections/implementation/flags/nearby_connections_feature_flags.h +++ b/connections/implementation/flags/nearby_connections_feature_flags.h @@ -70,6 +70,10 @@ constexpr auto kProcessBwuFrameAfterPcpConnected = constexpr auto kCheckIllegalCharacters = flags::Flag(kConfigPackage, "45632028", false); +// When true, allows to enable Multiplex feature. +constexpr auto kEnableMultiplex = + flags::Flag(kConfigPackage, "45627836", false); + } // namespace nearby_connections_feature } // namespace config_package_nearby } // namespace connections diff --git a/connections/implementation/mediums/multiplex/BUILD b/connections/implementation/mediums/multiplex/BUILD index 903fd433..e7209e7e 100644 --- a/connections/implementation/mediums/multiplex/BUILD +++ b/connections/implementation/mediums/multiplex/BUILD @@ -18,10 +18,12 @@ cc_library( srcs = [ "multiplex_frames.cc", "multiplex_output_stream.cc", + "multiplex_socket.cc", ], hdrs = [ "multiplex_frames.h", "multiplex_output_stream.h", + "multiplex_socket.h", ], copts = ["-DCORE_ADAPTER_DLL"], visibility = [ @@ -38,6 +40,7 @@ cc_library( "//internal/platform:uuid", "//internal/platform/implementation:comm", "//internal/platform/implementation:types", + "//proto:connections_enums_cc_proto", "//proto/mediums:multiplex_frames_cc_proto", "@aappleby_smhasher//:libmurmur3", "@com_google_absl//absl/base:core_headers", @@ -59,16 +62,21 @@ cc_test( srcs = [ "multiplex_frames_test.cc", "multiplex_output_stream_test.cc", + "multiplex_socket_test.cc", ], deps = [ ":multiplex", + "//connections/implementation:internal", "//internal/platform:base", "//internal/platform:comm", "//internal/platform:test_util", "//internal/platform:types", + "//internal/platform/implementation:comm", "//internal/platform/implementation/g3", # buildcleaner: keep + "//proto:connections_enums_cc_proto", "//proto/mediums:multiplex_frames_cc_proto", "@com_github_protobuf_matchers//protobuf-matchers", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/hash:hash_testing", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/connections/implementation/mediums/multiplex/multiplex_output_stream.cc b/connections/implementation/mediums/multiplex/multiplex_output_stream.cc index ba65ce5e..89acacec 100644 --- a/connections/implementation/mediums/multiplex/multiplex_output_stream.cc +++ b/connections/implementation/mediums/multiplex/multiplex_output_stream.cc @@ -161,18 +161,17 @@ bool MultiplexOutputStream::Close(const std::string& service_id) { return false; } - auto service_id_hash_salt = item->second->GetserviceIdHashSalt(); + auto service_id_hash_salt = item->second->GetServiceIdHashSalt(); item->second->Close(); if (is_enabled_.Get()) { Future future; multiplex_writer_.EnqueueToSend( &future, - ForDisconnection(service_id, item->second->GetserviceIdHashSalt()), + ForDisconnection(service_id, item->second->GetServiceIdHashSalt()), "MultiplexFrame::DISCONNECTION"); WaitForResult("MultiplexFrame::DISCONNECTION", &future); } virtual_output_streams_.erase(service_id); - if (virtual_output_streams_.empty()) { physical_writer_->Close(); multiplex_writer_.Close(); @@ -203,11 +202,11 @@ OutputStream* MultiplexOutputStream::CreateVirtualOutputStream( .first->second.get(); } -std::string MultiplexOutputStream::GetserviceIdHashSalt( +std::string MultiplexOutputStream::GetServiceIdHashSalt( const std::string& service_id) { auto item = virtual_output_streams_.find(service_id); if (item != virtual_output_streams_.end()) { - return item->second->GetserviceIdHashSalt(); + return item->second->GetServiceIdHashSalt(); } return {}; } @@ -231,7 +230,7 @@ MultiplexOutputStream::MultiplexWriter::~MultiplexWriter() { void MultiplexOutputStream::MultiplexWriter::EnqueueToSend( Future* future, const ByteArray& data, const std::string& frame_name) { - MutexLock lock(&mutex_); + MutexLock lock(&writing_mutex_); data_queue_.Put(EnqueuedFrame(future, data)); if (is_writing_) { @@ -253,18 +252,20 @@ void MultiplexOutputStream::MultiplexWriter::StartWriting() { Write(enqueued_frame.value()); continue; } - - MutexLock lock(&mutex_); - if (data_queue_.Empty() && is_writing_) { - is_writing_ = false; - Exception wait_succeeded = is_writing_cond_.Wait(); - if (!wait_succeeded.Ok()) { - NEARBY_LOGS(WARNING) - << TAG << __func__ - << ": Failure waiting to wait: " << wait_succeeded.value; - return; + { + MutexLock lock(&writing_mutex_); + if (data_queue_.Empty() && is_writing_) { + is_writing_ = false; + NEARBY_LOGS(INFO) << TAG << "Waiting for data_queue_ has data."; + Exception wait_succeeded = is_writing_cond_.Wait(); + if (!wait_succeeded.Ok()) { + NEARBY_LOGS(WARNING) + << TAG << __func__ + << ": Failure waiting to wait: " << wait_succeeded.value; + return; + } + if (is_closed_) break; } - if (is_closed_) break; } } NEARBY_LOGS(INFO) << TAG << "Writing loop stopped."; @@ -291,7 +292,7 @@ void MultiplexOutputStream::MultiplexWriter::Write( } void MultiplexOutputStream::MultiplexWriter::Close() { - MutexLock lock(&mutex_); + MutexLock lock(&writing_mutex_); is_closed_ = true; if (is_write_loop_running_) { NEARBY_LOGS(INFO) << TAG << "Stop writing loop and Shutdown writer thread."; diff --git a/connections/implementation/mediums/multiplex/multiplex_output_stream.h b/connections/implementation/mediums/multiplex/multiplex_output_stream.h index f777af7a..ea16394f 100644 --- a/connections/implementation/mediums/multiplex/multiplex_output_stream.h +++ b/connections/implementation/mediums/multiplex/multiplex_output_stream.h @@ -129,7 +129,7 @@ class MultiplexOutputStream { const std::string& service_id, const std::string& service_id_hash_salt); // Gets the service id hash salt. - std::string GetserviceIdHashSalt(const std::string& service_id); + std::string GetServiceIdHashSalt(const std::string& service_id); // Shuts down the multiplex output stream. void Shutdown(); @@ -169,10 +169,9 @@ class MultiplexOutputStream { FeatureFlags::GetInstance() .GetFlags() .multiplex_socket_middle_priority_queue_capacity}; - - mutable Mutex mutex_; - ConditionVariable is_writing_cond_{&mutex_}; - bool is_writing_ = false; + mutable Mutex writing_mutex_; + ConditionVariable is_writing_cond_{&writing_mutex_}; + bool is_writing_ ABSL_GUARDED_BY(writing_mutex_) = false; bool is_closed_ = false; // The single thread to write all enqueued frames. @@ -198,7 +197,7 @@ class MultiplexOutputStream { } // Returns the service id hash salt. - std::string GetserviceIdHashSalt() { return service_id_hash_salt_; } + std::string GetServiceIdHashSalt() { return service_id_hash_salt_; } // Sets the service id hash salt. void SetserviceIdHashSalt(std::string service_id_hash_salt) { diff --git a/connections/implementation/mediums/multiplex/multiplex_output_stream_test.cc b/connections/implementation/mediums/multiplex/multiplex_output_stream_test.cc index af6bbda8..071ad322 100644 --- a/connections/implementation/mediums/multiplex/multiplex_output_stream_test.cc +++ b/connections/implementation/mediums/multiplex/multiplex_output_stream_test.cc @@ -204,7 +204,7 @@ TEST_F(MultiplexOutputStreamTest, CreateTwoVirtualStreams_SendData) { MultiThreadExecutor executor(2); CountDownLatch latch(2); executor.Execute([&virtual_output_stream_1, &latch, &data_1]() { - absl::SleepFor(absl::Milliseconds(500)); + absl::SleepFor(absl::Milliseconds(100)); virtual_output_stream_1->Write(data_1); virtual_output_stream_1->Flush(); latch.CountDown(); @@ -214,7 +214,7 @@ TEST_F(MultiplexOutputStreamTest, CreateTwoVirtualStreams_SendData) { virtual_output_stream_2->Flush(); latch.CountDown(); }); - EXPECT_TRUE(latch.Await(absl::Milliseconds(2000)).result()); + EXPECT_TRUE(latch.Await(absl::Milliseconds(5000)).result()); auto frame_data = ReadFrame(); ASSERT_TRUE(frame_data.ok()); diff --git a/connections/implementation/mediums/multiplex/multiplex_socket.cc b/connections/implementation/mediums/multiplex/multiplex_socket.cc new file mode 100644 index 00000000..1801900c --- /dev/null +++ b/connections/implementation/mediums/multiplex/multiplex_socket.cc @@ -0,0 +1,743 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "connections/implementation/mediums/multiplex/multiplex_socket.h" + +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/functional/any_invocable.h" +#include "absl/strings/string_view.h" +#include "connections/implementation/mediums/multiplex/multiplex_frames.h" +#include "connections/implementation/mediums/multiplex/multiplex_output_stream.h" +#include "connections/implementation/mediums/utils.h" +#include "internal/platform/base64_utils.h" +#include "internal/platform/byte_array.h" +#include "internal/platform/exception.h" +#include "internal/platform/feature_flags.h" +#include "internal/platform/future.h" +#include "internal/platform/logging.h" +#include "internal/platform/mutex.h" +#include "internal/platform/mutex_lock.h" +#include "internal/platform/socket.h" +#include "proto/connections_enums.pb.h" + +namespace nearby { +namespace connections { +namespace mediums { +namespace multiplex { + +namespace { +// It is defined for the receiver which send the first packet to the sender +// without getting salt from it yet. The fake salt reminds sender to get the +// correct socket from `virtualSockets` without remapping it. +constexpr absl::string_view kFakeSalt = "RECEIVER_CONDIMENT"; +} // namespace + +using ::location::nearby::mediums::ConnectionResponseFrame; +using ::location::nearby::mediums::MultiplexControlFrame; +using ::location::nearby::mediums::MultiplexDataFrame; +using ::location::nearby::mediums::MultiplexFrame; +using ::location::nearby::proto::connections::Medium; +using ::location::nearby::proto::connections::Medium_Name; +using ConnectionResponseCode = ConnectionResponseFrame::ConnectionResponseCode; + +void MultiplexSocket::ListenForIncomingConnection( + const std::string& service_id, Medium type, + MultiplexIncomingConnectionCb incoming_connection_cb) { + GetIncomingConnectionCallbacks().emplace( + std::pair(service_id, type), + std::move(incoming_connection_cb)); +} + +void MultiplexSocket::StopListeningForIncomingConnection( + const std::string& service_id, Medium type) { + GetIncomingConnectionCallbacks().erase( + std::pair(service_id, type)); +} + +MultiplexSocket::MultiplexSocket(MediumSocket* physical_socket) + : physical_socket_(physical_socket), + multiplex_output_stream_{&physical_socket->GetOutputStream(), enabled_}, + physical_reader_(&physical_socket->GetInputStream()) {} + +absl::flat_hash_map, + MultiplexIncomingConnectionCb>& +MultiplexSocket::GetIncomingConnectionCallbacks() { + static std::aligned_storage_t< + sizeof(absl::flat_hash_map, + MultiplexIncomingConnectionCb>), + alignof(absl::flat_hash_map, + MultiplexIncomingConnectionCb>)> + storage; + static absl::flat_hash_map, + MultiplexIncomingConnectionCb>* + incoming_connection_callbacks = + new (&storage) absl::flat_hash_map, + MultiplexIncomingConnectionCb>(); + + return *incoming_connection_callbacks; +} + +MultiplexSocket* MultiplexSocket::CreateIncomingSocket( + MediumSocket* physical_socket, const std::string& service_id) { + static MultiplexSocket* multiplex_incoming_socket = nullptr; + switch (physical_socket->GetMedium()) { + case Medium::BLUETOOTH: + static std::aligned_storage_t + storage_bt; + multiplex_incoming_socket = + new (&storage_bt) MultiplexSocket(physical_socket); + break; + case Medium::BLE: + static std::aligned_storage_t + storage_ble; + multiplex_incoming_socket = + new (&storage_ble) MultiplexSocket(physical_socket); + break; + case Medium::WIFI_LAN: + static std::aligned_storage_t + storage_wlan; + multiplex_incoming_socket = + new (&storage_wlan) MultiplexSocket(physical_socket); + break; + default: + NEARBY_LOGS(ERROR) << __func__ << "Unsupported medium: " + << physical_socket->GetMedium(); + multiplex_incoming_socket = nullptr; + return multiplex_incoming_socket; + } + + auto on_physical_socket_closed_listener = + std::make_unique>( + []() { multiplex_incoming_socket->OnPhysicalSocketClosed(); }); + physical_socket->AddOnSocketClosedListener( + std::move(on_physical_socket_closed_listener)); + NEARBY_LOGS(INFO) << __func__ + << "CreateIncomingSocket with serviceId=" << service_id + << ", serviceIdHashSalt=" << kFakeSalt; + multiplex_incoming_socket->CreateFirstVirtualSocket(service_id, + (std::string)kFakeSalt); + multiplex_incoming_socket->StartReaderThread(); + + return multiplex_incoming_socket; +} + +MultiplexSocket* MultiplexSocket::CreateOutgoingSocket( + MediumSocket* physical_socket, const std::string& service_id, + const std::string& service_id_hash_salt) { + static MultiplexSocket* multiplex_outgoing_socket = nullptr; + switch (physical_socket->GetMedium()) { + case Medium::BLUETOOTH: + static std::aligned_storage_t + storage_bt; + multiplex_outgoing_socket = + new (&storage_bt) MultiplexSocket(physical_socket); + break; + case Medium::BLE: + static std::aligned_storage_t + storage_ble; + multiplex_outgoing_socket = + new (&storage_ble) MultiplexSocket(physical_socket); + break; + case Medium::WIFI_LAN: + static std::aligned_storage_t + storage_wlan; + multiplex_outgoing_socket = + new (&storage_wlan) MultiplexSocket(physical_socket); + break; + default: + NEARBY_LOGS(ERROR) << __func__ << "Unsupported medium: " + << physical_socket->GetMedium(); + return multiplex_outgoing_socket; + } + + auto on_physical_socket_closed_listener = + std::make_unique>( + []() { multiplex_outgoing_socket->OnPhysicalSocketClosed(); }); + physical_socket->AddOnSocketClosedListener( + std::move(on_physical_socket_closed_listener)); + NEARBY_LOGS(INFO) << __func__ + << "CreateOutgoingSocket with serviceId=" << service_id + << ", serviceIdHashSalt=" << service_id_hash_salt; + NEARBY_LOGS(INFO) << __func__ << "multiplex_outgoing_socket:" + << multiplex_outgoing_socket; + multiplex_outgoing_socket->CreateFirstVirtualSocket(service_id, + service_id_hash_salt); + multiplex_outgoing_socket->StartReaderThread(); + return multiplex_outgoing_socket; +} + +MultiplexSocket* MultiplexSocket::CreateOutgoingSocket( + MediumSocket* physical_socket, const std::string& service_id) { + return CreateOutgoingSocket(physical_socket, service_id, + Utils::GenerateSalt()); +} + +MediumSocket* MultiplexSocket::CreateFirstVirtualSocket( + const std::string& service_id, const std::string& service_id_hash_salt) { + auto output_stream = + multiplex_output_stream_.CreateVirtualOutputStreamForFirstVirtualSocket( + service_id, service_id_hash_salt); + + MutexLock lock(&virtual_socket_mutex_); + std::string salted_service_id_hash_key = + GenerateServiceIdHashKeyWithSalt(service_id, service_id_hash_salt); + + MediumSocket* virtual_socket = physical_socket_->CreateVirtualSocket( + salted_service_id_hash_key, output_stream, physical_socket_->GetMedium(), + &virtual_sockets_); + + virtual_socket->AddOnSocketClosedListener( + std::make_unique>( + [this, &service_id]() { OnVirtualSocketClosed(service_id); })); + + if (!IsEnabled()) { + virtual_socket->RegisterMultiplexEnabledCallback(enable_cb_); + } + + return virtual_socket; +} + +MediumSocket* MultiplexSocket::CreateVirtualSocket( + const std::string& service_id, const std::string& service_id_hash_salt) { + auto output_stream = multiplex_output_stream_.CreateVirtualOutputStream( + service_id, service_id_hash_salt); + MutexLock lock(&virtual_socket_mutex_); + std::string salted_service_id_hash_key = + GenerateServiceIdHashKeyWithSalt(service_id, service_id_hash_salt); + + MediumSocket* virtual_socket = physical_socket_->CreateVirtualSocket( + salted_service_id_hash_key, output_stream, physical_socket_->GetMedium(), + &virtual_sockets_); + + virtual_socket->AddOnSocketClosedListener( + std::make_unique>( + [this, &service_id]() { OnVirtualSocketClosed(service_id); })); + + return virtual_socket; +} + +MediumSocket* MultiplexSocket::GetVirtualSocket(const std::string& service_id) { + MutexLock lock(&virtual_socket_mutex_); + auto item = virtual_sockets_.find(GenerateServiceIdHashKeyWithSalt( + service_id, multiplex_output_stream_.GetServiceIdHashSalt(service_id))); + if (item == virtual_sockets_.end()) { + return nullptr; + } + return item->second.get(); +} + +int MultiplexSocket::GetVirtualSocketCount() { + MutexLock lock(&virtual_socket_mutex_); + return virtual_sockets_.size(); +} + +std::shared_ptr> +MultiplexSocket::RegisterConnectionResponse(const std::string& service_id) { + auto future = std::make_shared>(); + connection_response_futures_.emplace(service_id, future); + + return future; +} + +void MultiplexSocket::UnRegisterConnectionResponse( + const std::string& service_id) { + connection_response_futures_.erase(service_id); +} + +MediumSocket* MultiplexSocket::EstablishVirtualSocket( + const std::string& service_id) { + if (!IsEnabled()) { + NEARBY_LOGS(ERROR) << __func__ << "EstablishVirtualSocket disabled"; + return nullptr; + } + + std::string service_id_hash_salt = Utils::GenerateSalt(); + auto future = RegisterConnectionResponse(service_id); + + multiplex_output_stream_.WriteConnectionRequestFrame(service_id, + service_id_hash_salt); + auto result = + future->Get(FeatureFlags::GetInstance() + .GetFlags() + .multiplex_socket_connection_response_timeout_millis); + if (!result.ok()) { + NEARBY_LOGS(ERROR) << __func__ + << "EstablishVirtualSocket failed with response code=" + << result.exception(); + return nullptr; + } + + ConnectionResponseCode response_code = result.GetResult(); + switch (response_code) { + case ConnectionResponseFrame::CONNECTION_ACCEPTED: + NEARBY_LOGS(INFO) << __func__ + << "EstablishVirtualSocket after remote response to" + " accept the connection with service_id=" + << service_id + << ", service_id_hash_salt=" << service_id_hash_salt; + return CreateVirtualSocket(service_id, service_id_hash_salt); + case ConnectionResponseFrame::NOT_LISTENING: + NEARBY_LOGS(ERROR) << __func__ + << "EstablishVirtualSocket failed for service_id=" + << service_id + << ", service_id_hash_salt=" << service_id_hash_salt + << " with response code=NOT_LISTENING"; + break; + default: + NEARBY_LOGS(ERROR) << __func__ + << "EstablishVirtualSocket failed for service_id=" + << service_id + << ", service_id_hash_salt=" << service_id_hash_salt + << " with response code=UNKNOWN_RESPONSE_CODE"; + break; + } + return nullptr; +} + +void MultiplexSocket::StartReaderThread() { + if (is_shutdown_) { + NEARBY_LOGS(WARNING) << __func__ + << "Stop to start reader thread since socket is " + "shutdown."; + return; + } + physical_reader_thread_.Execute([this]() { + NEARBY_LOGS(INFO) << __func__ << " Reader thread starts."; + while (!is_shutdown_) { + bool fail = false; + ExceptionOr bytes; + ExceptionOr read_int = + Base64Utils::ReadInt(physical_reader_); + if (!read_int.ok()) { + NEARBY_LOGS(WARNING) + << __func__ << "Failed to read. Exception:" << read_int.exception(); + fail = true; + } else { + auto length = read_int.result(); + NEARBY_LOGS(VERBOSE) << __func__ << " length:" << length; + + if (length < 0 || length > FeatureFlags::GetInstance() + .GetFlags() + .connection_max_frame_length) { + // Ignore the failure because not only one client use this + // connection. + NEARBY_LOGS(WARNING) + << __func__ << "Failed to read because received a invalid length " + << length << ", but continue to read."; + continue; + } + + bytes = physical_reader_->ReadExactly(length); + if (!bytes.ok()) { + NEARBY_LOGS(WARNING) + << __func__ << "Read data exception:" << bytes.exception(); + fail = true; + } + } + if (fail) { + { + MutexLock lock(&virtual_socket_mutex_); + if (virtual_sockets_.empty()) { + NEARBY_LOGS(INFO) + << __func__ + << "The reader thread stopped because all virtual socket " + "closed."; + } else { + NEARBY_LOGS(ERROR) << __func__ + << "The reader thread stopped because " + "unexpected IOException"; + } + } + return; + } + ExceptionOr frame_exc = + multiplex::FromBytes(bytes.result()); + if (!frame_exc.ok()) { + HandleOfflineFrame(bytes.result()); + continue; + } + + if (!IsEnabled()) { + // The reader thread will only be enabled when local device + // supports multiplex if we received a multiplex frame from + // the remote, it means that the remote and the local both + // support multiplex as well. So it is safe to just turn on + // the feature at this point. + NEARBY_LOGS(INFO) + << __func__ + << "Received a multiplex frame while not enabled, enable " + "multiplex."; + Enable(); + } + auto frame = frame_exc.result(); + auto salted_service_id_hash = + ByteArray{std::move(frame.header().salted_service_id_hash())}; + auto service_id_hash_salt = frame.header().has_service_id_hash_salt() + ? frame.header().service_id_hash_salt() + : ""; + switch (frame.frame_type()) { + case MultiplexFrame::CONTROL_FRAME: + HandleControlFrame(salted_service_id_hash, service_id_hash_salt, + frame.control_frame()); + break; + case MultiplexFrame::DATA_FRAME: + HandleDataFrame(salted_service_id_hash, service_id_hash_salt, + frame.data_frame()); + break; + default: + NEARBY_LOGS(WARNING) + << __func__ << "Received MultiplexFrame with unknown frame type " + << frame.frame_type(); + } + } + }); +} + +void MultiplexSocket::HandleOfflineFrame(const ByteArray& bytes) { + // Only pass the data when there's only 1 VirtualSocket. + MutexLock lock(&virtual_socket_mutex_); + NEARBY_LOGS(INFO) << __func__ + << " Virtual_socket num:" << virtual_sockets_.size(); + if (virtual_sockets_.size() == 1) { + auto item = virtual_sockets_.begin(); + if (item->second == nullptr) { + NEARBY_LOGS(WARNING) << "Expected one live socket, but found null."; + return; + } + item->second->FeedIncomingData(Base64Utils::IntToBytes(bytes.size())); + item->second->FeedIncomingData(bytes); + } +} + +void MultiplexSocket::HandleControlFrame( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt, + const MultiplexControlFrame& frame) { + switch (frame.control_frame_type()) { + case MultiplexControlFrame::CONNECTION_REQUEST: + RunOffloadThread("CONNECTION_REQUEST", [this, &salted_service_id_hash, + &service_id_hash_salt] { + HandleConnectionRequest(salted_service_id_hash, service_id_hash_salt); + }); + break; + case MultiplexControlFrame::CONNECTION_RESPONSE: + NEARBY_LOGS(INFO) + << __func__ << "Received an CONNECTION_RESPONSE frame." + << " salted_service_id_hash: " << std::string(salted_service_id_hash) + << ", service_id_hash_salt: " << service_id_hash_salt + << ", ConnectionResponseCode: " + << frame.connection_response_frame().connection_response_code(); + + RunOffloadThread("CONNECTION_RESPONSE", [this, &salted_service_id_hash, + &service_id_hash_salt, + frame = frame] { + HandleConnectionResponse(salted_service_id_hash, service_id_hash_salt, + frame.connection_response_frame()); + }); + break; + case MultiplexControlFrame::DISCONNECTION: + RunOffloadThread("DISCONNECTION", [this, &salted_service_id_hash] { + HandleDisconnection(salted_service_id_hash); + }); + break; + default: + NEARBY_LOGS(WARNING) << __func__ << "Received an unknown frame type " + << frame.control_frame_type(); + break; + } +} + +void MultiplexSocket::HandleConnectionRequest( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt) { + if (!IsEnabled()) { + NEARBY_LOGS(WARNING) << __func__ + << "Received a CONNECTION_REQUEST frame on medium " + << Medium_Name(physical_socket_->GetMedium()) + << " but status is disabled, ignore it."; + return; + } + + std::string salted_service_id_hash_key = + GenerateServiceIdHashKey(salted_service_id_hash); + MultiplexIncomingConnectionCb* incoming_connection_callback = nullptr; + std::string listening_service_id = ""; + for (auto& [service_id_medium_pair, callback] : + GetIncomingConnectionCallbacks()) { + if (GenerateServiceIdHashWithSalt(service_id_medium_pair.first, + service_id_hash_salt) == + salted_service_id_hash) { + incoming_connection_callback = &callback; + listening_service_id = service_id_medium_pair.first; + } + } + + if (incoming_connection_callback == nullptr || listening_service_id.empty()) { + NEARBY_LOGS(INFO) << __func__ + << "There's no client listening for hash salt : " + << service_id_hash_salt + << ", hash key : " << salted_service_id_hash_key + << " on medium " + << Medium_Name(physical_socket_->GetMedium()); + + NEARBY_LOGS(INFO) << __func__ << "Dump incomingConnectionCallbacks : " + << GetIncomingConnectionCallbacks().size(); + if (!multiplex_output_stream_.WriteConnectionResponseFrame( + salted_service_id_hash, service_id_hash_salt, + ConnectionResponseFrame::NOT_LISTENING)) { + NEARBY_LOGS(INFO) << __func__ << "Failed to write NOT_LISTENING frame."; + } + return; + } + NEARBY_LOGS(INFO) << __func__ + << "Accept new virtual socket request service ID : " + << listening_service_id + << ", hash salt : " << service_id_hash_salt + << ", hash key : " << salted_service_id_hash_key + << " on medium " + << Medium_Name(physical_socket_->GetMedium()); + + if (!multiplex_output_stream_.WriteConnectionResponseFrame( + salted_service_id_hash, service_id_hash_salt, + ConnectionResponseFrame::CONNECTION_ACCEPTED)) { + NEARBY_LOGS(INFO) << __func__ + << "Failed to write CONNECTION_ACCEPTED frame."; + return; + } + + NEARBY_LOGS(VERBOSE) + << __func__ + << "establishVirtualSocket after local device accept the connection " + "with serviceId=" + << listening_service_id << ", serviceIdHashSalt=" << service_id_hash_salt; + MediumSocket* virtual_socket = + CreateVirtualSocket(listening_service_id, service_id_hash_salt); + (*incoming_connection_callback)(std::move(listening_service_id), + virtual_socket); +} + +void MultiplexSocket::HandleConnectionResponse( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt, + const ConnectionResponseFrame& frame) { + NEARBY_LOGS(INFO) << __func__ << "connection_response_code: " + << frame.connection_response_code(); + for (auto& [service_id, future] : connection_response_futures_) { + if (GenerateServiceIdHashWithSalt(service_id, service_id_hash_salt) == + salted_service_id_hash) { + if (future != nullptr) { + future->Set(frame.connection_response_code()); + NEARBY_LOGS(INFO) << __func__ + << "Set the future for serviceId=" << service_id + << ", serviceIdHashSalt=" << service_id_hash_salt + << " with response code=" + << frame.connection_response_code(); + return; + } + } + } + + NEARBY_LOGS(WARNING) + << __func__ + << "Received a CONNECTION_RESPONSE frame but no client waiting for " + "service ID Hash Key" + << GenerateServiceIdHashKey(salted_service_id_hash); +} + +void MultiplexSocket::HandleDisconnection( + const ByteArray& salted_service_id_hash) { + std::string salted_service_id_hash_key = + GenerateServiceIdHashKey(salted_service_id_hash); + { + MutexLock lock(&virtual_socket_mutex_); + auto item = virtual_sockets_.find(salted_service_id_hash_key); + if (item != virtual_sockets_.end()) { + NEARBY_LOGS(INFO) + << __func__ + << "Received a DISCONNECTION frame to disconnect virtual socket for " + "salted service ID Hash Key " + << salted_service_id_hash_key; + if (item->second != nullptr) { + item->second->Close(); + } + virtual_sockets_.erase(item); + // physical_socket_->RemoveVirtualSocket(salted_service_id_hash_key); + + if (virtual_sockets_.empty()) { + NEARBY_LOGS(INFO) << __func__ + << "Close the physical socket because all services " + "disconnected."; + physical_socket_->Close(); + } + } else { + NEARBY_LOGS(WARNING) + << __func__ + << "Received a DISCONNECTION frame but there's no alive socket to " + "disconnect for service ID Hash Key " + << salted_service_id_hash_key; + } + } +} + +void MultiplexSocket::HandleDataFrame(const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt, + const MultiplexDataFrame& frame) { + std::string salted_service_id_hash_key = + GenerateServiceIdHashKey(salted_service_id_hash); + MediumSocket* virtual_socket = nullptr; + if (service_id_hash_salt.empty()) { + { + MutexLock lock(&virtual_socket_mutex_); + auto item = virtual_sockets_.find(salted_service_id_hash_key); + if (item != virtual_sockets_.end()) { + virtual_socket = item->second.get(); + } + } + } else { + virtual_socket = + ReMapAndGetVirtualSocket(salted_service_id_hash, service_id_hash_salt); + } + + if (virtual_socket != nullptr) { + NEARBY_LOGS(INFO) + << __func__ + << "Received a DATA frame to feed virtual socket for salted service ID " + "Hash Key " + << salted_service_id_hash_key; + virtual_socket->FeedIncomingData(ByteArray(frame.data())); + } else { + NEARBY_LOGS(WARNING) + << __func__ + << "Received a DATA frame but there's no alive socket to feed for " + "salted service ID Hash Key " + << salted_service_id_hash_key; + } +} + +void MultiplexSocket::OnPhysicalSocketClosed() { + RunOffloadThread("Shutdown", [this]() { Shutdown(); }); +} + +void MultiplexSocket::OnVirtualSocketClosed(const std::string& service_id) { + RunOffloadThread("VirtualSocketClosed", [this, &service_id]() { + { + MutexLock lock(&virtual_socket_mutex_); + MediumSocket* virtual_socket = GetVirtualSocket(service_id); + if (virtual_socket != nullptr) { + virtual_sockets_.erase(GenerateServiceIdHashKeyWithSalt( + service_id, + multiplex_output_stream_.GetServiceIdHashSalt(service_id))); + NEARBY_LOGS(INFO) << __func__ << "Virtual socket(" << service_id + << ") disconnected"; + multiplex_output_stream_.Close(service_id); + virtual_socket->Close(); + + if (virtual_sockets_.empty()) { + NEARBY_LOGS(INFO) << __func__ + << "Close the physical socket because all virtual " + "sockets disconnected."; + physical_socket_->Close(); + } + return; + } + + NEARBY_LOGS(INFO) << __func__ << "Virtual socket(" << service_id + << ") not found"; + } + }); +} + +MediumSocket* MultiplexSocket::ReMapAndGetVirtualSocket( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt) { + std::string salted_service_id_hash_key = + GenerateServiceIdHashKey(salted_service_id_hash); + NEARBY_LOGS(VERBOSE) << __func__ + << "reMapAndGetVirtualSocket with serviceIdHashSalt=" + << service_id_hash_salt << ", saltedServiceIdHashKey=" + << salted_service_id_hash_key; + { + MutexLock lock(&virtual_socket_mutex_); + for (auto& [hash_key, virtual_socket] : virtual_sockets_) { + auto output_stream = + dynamic_cast( + &(virtual_socket->GetOutputStream())); + if (output_stream == nullptr) { + continue; + } + if (!output_stream->IsFirstVirtualOutputStream()) { + continue; + } + if ((service_id_hash_salt == kFakeSalt) || + (hash_key == salted_service_id_hash_key)) { + return virtual_socket.get(); + } else { + NEARBY_LOGS(INFO) << __func__ << "Remap the virtualSockets."; + virtual_sockets_.erase(hash_key); + output_stream->SetserviceIdHashSalt(service_id_hash_salt); + virtual_sockets_.emplace(salted_service_id_hash_key, virtual_socket); + + return virtual_socket.get(); + } + } + } + + NEARBY_LOGS(INFO) << __func__ << "Failed to remap the virtualSockets."; + return nullptr; +} + +void MultiplexSocket::RunOffloadThread(const std::string& name, + absl::AnyInvocable runnable) { + single_thread_offloader_.Execute(name, std::move(runnable)); +} + +void MultiplexSocket::Shutdown() { + NEARBY_LOGS(INFO) << __func__ << " shutdown"; + { + MutexLock lock(&virtual_socket_mutex_); + for (auto& [hash_key, virtual_socket] : virtual_sockets_) { + if (virtual_socket != nullptr) { + virtual_socket->Close(); + } + } + virtual_sockets_.clear(); + } + + multiplex_output_stream_.Shutdown(); + physical_socket_->Close(); + + GetIncomingConnectionCallbacks().clear(); + connection_response_futures_.clear(); + + physical_reader_thread_.Shutdown(); + single_thread_offloader_.Shutdown(); + is_shutdown_ = true; +} + +} // namespace multiplex +} // namespace mediums +} // namespace connections +} // namespace nearby diff --git a/connections/implementation/mediums/multiplex/multiplex_socket.h b/connections/implementation/mediums/multiplex/multiplex_socket.h new file mode 100644 index 00000000..2740ce50 --- /dev/null +++ b/connections/implementation/mediums/multiplex/multiplex_socket.h @@ -0,0 +1,214 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef CORE_INTERNAL_MEDIUMS_MULTIPLEX_MULTIPLEX_SOCKET_H_ +#define CORE_INTERNAL_MEDIUMS_MULTIPLEX_MULTIPLEX_SOCKET_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/flat_hash_map.h" +#include "absl/functional/any_invocable.h" +#include "connections/implementation/mediums/multiplex/multiplex_output_stream.h" +#include "internal/platform/atomic_boolean.h" +#include "internal/platform/byte_array.h" +#include "internal/platform/count_down_latch.h" +#include "internal/platform/exception.h" +#include "internal/platform/feature_flags.h" +#include "internal/platform/future.h" +#include "internal/platform/input_stream.h" +#include "internal/platform/logging.h" +#include "internal/platform/mutex.h" +#include "internal/platform/output_stream.h" +#include "internal/platform/settable_future.h" +#include "internal/platform/single_thread_executor.h" +#include "internal/platform/socket.h" +#include "proto/connections_enums.pb.h" +#include "proto/mediums/multiplex_frames.pb.h" + +namespace nearby { +namespace connections { +namespace mediums { +namespace multiplex { + +using MultiplexEnbaleCb = absl::AnyInvocable; +using MultiplexIncomingConnectionCb = absl::AnyInvocable; + +class MultiplexSocket { + public: + MultiplexSocket(const MultiplexSocket&) = delete; + MultiplexSocket& operator=(const MultiplexSocket&) = delete; + + // Creates a new incoming MultiplexSocket. + static MultiplexSocket* CreateIncomingSocket(MediumSocket* physical_socket, + const std::string& service_id); + // Creates a new outgoing MultiplexSocket. + static MultiplexSocket* CreateOutgoingSocket( + MediumSocket* physical_socket, const std::string& service_id, + const std::string& service_id_hash_salt); + + // Creates a new outgoing MultiplexSocket with default service_id_hash_salt. + static MultiplexSocket* CreateOutgoingSocket(MediumSocket* physical_socket, + const std::string& service_id); + + // A Table of service Id as row key, medium type as column key, and {@link + // IncomingConnectionCallback} as value. Non-empty while the client starts + // listening for incoming virtual socket. + static absl::flat_hash_map< + std::pair, + MultiplexIncomingConnectionCb>& + GetIncomingConnectionCallbacks(); + + // Listens for incoming connection through multiplex for specified {@code + // service_id} on medium + // {@code type}. Should register the callback before new the MultiplexSocket. + static void ListenForIncomingConnection( + const std::string& service_id, + ::location::nearby::proto::connections::Medium type, + absl::AnyInvocable + incoming_connection_cb); + + // Stops listening for incoming multiplex connection for {@code service_id} on + // medium {@code type}. + static void StopListeningForIncomingConnection( + const std::string& service_id, + ::location::nearby::proto::connections::Medium type); + + bool IsEnabled() { return enabled_.Get(); } + void Enable() { + NEARBY_LOGS(INFO) << "Enable the Multiplex MediumSocket."; + enabled_.Set(true); + } + + // Gets the physical socket. + MediumSocket* GetPhysicalSocket() { return physical_socket_; } + // Gets the virtual socket by service id. + MediumSocket* GetVirtualSocket(const std::string& service_id); + // Gets the virtual socket count. + int GetVirtualSocketCount(); + + // Establishes the virtual socket by service id. + MediumSocket* EstablishVirtualSocket(const std::string& service_id); + // Shuts down the multiplex socket. + void Shutdown(); + + private: + explicit MultiplexSocket(MediumSocket* physical_socket); + ~MultiplexSocket() = default; + + // Creates the first virtual socket for the service id. The first virtual + // socket is created by the sender. + MediumSocket* CreateFirstVirtualSocket( + const std::string& service_id, const std::string& service_id_hash_salt); + // Creates the virtual socket for the service id. + MediumSocket* CreateVirtualSocket(const std::string& service_id, + const std::string& service_id_hash_salt); + // Registers the connection response future for the service id. + std::shared_ptr> + RegisterConnectionResponse(const std::string& service_id); + // Unregisters the connection response future for the service id. + void UnRegisterConnectionResponse(const std::string& service_id); + // Starts the reader thread to read the incoming MultiplexFrame from the + // physical socket. + void StartReaderThread(); + // Handles the offline frame from the physical socket. + void HandleOfflineFrame(const ByteArray& bytes); + // Handles the control frame from the physical socket. + void HandleControlFrame( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt, + const ::location::nearby::mediums::MultiplexControlFrame& frame); + // Handles the connection request frame from the physical socket. + void HandleConnectionRequest(const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt); + // Handles the connection response frame from the physical socket. + void HandleConnectionResponse( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt, + const ::location::nearby::mediums::ConnectionResponseFrame& frame); + // Handles the disconnection frame from the physical socket. + void HandleDisconnection(const ByteArray& salted_service_id_hash); + // Handles the data frame from the physical socket. + void HandleDataFrame( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt, + const ::location::nearby::mediums::MultiplexDataFrame& frame); + // Handles the physical socket closed. + void OnPhysicalSocketClosed(); + // Remaps and gets the virtual socket by service id hash. + MediumSocket* ReMapAndGetVirtualSocket( + const ByteArray& salted_service_id_hash, + const std::string& service_id_hash_salt); + // Handles the virtual socket closed. + void OnVirtualSocketClosed(const std::string& service_id); + // Runs the offload thread. + void RunOffloadThread(const std::string& name, + absl::AnyInvocable runnable); + + // The physical socket connect to the remote device. + MediumSocket* physical_socket_; + // The output stream to manage all outgoing frames from all clients. + MultiplexOutputStream multiplex_output_stream_; + // The {@link InputStream} of the physical socket. + InputStream* physical_reader_; + + // The callback to enable the MultiplexSocket. + std::shared_ptr> enable_cb_ = + std::make_shared>([this]() { Enable(); }); + + // A map of service Id -> {@link SettableFuture} for waiting the + // ConnectionResponse. Non-empty while requesting the virtual socket. + absl::flat_hash_map>> + connection_response_futures_; + + // A map of service Id hash key -> virtual socket. Non-empty while at least + // one virtual socket alive. Class derived from "MediumSocket" should define a + // pointer to the virtual sockets map. When here's any virtual socket + // operation, it will be reflected in both derived MediumSocket class and + // MultiplexSocket object + mutable Mutex virtual_socket_mutex_; + absl::flat_hash_map> + virtual_sockets_ ABSL_GUARDED_BY(virtual_socket_mutex_); + + // The thread to receive incoming MultiplexFrame from the physical socket. + SingleThreadExecutor physical_reader_thread_; + // The single thread we throw the potentially blocking work on to. + SingleThreadExecutor single_thread_offloader_; + + // The status of the MultiplexSocket enabled or disabled, it depends on both + // Sender and Receiver supports MultiplexSocket or not. Default disabled and + // enable it once two devices negotiated finished. + AtomicBoolean enabled_{false}; + + // If the socket is already shutdown and no longer in use. + bool is_shutdown_ = false; +}; + +} // namespace multiplex +} // namespace mediums +} // namespace connections +} // namespace nearby +#endif // CORE_INTERNAL_MEDIUMS_MULTIPLEX_MULTIPLEX_SOCKET_H_ diff --git a/connections/implementation/mediums/multiplex/multiplex_socket_test.cc b/connections/implementation/mediums/multiplex/multiplex_socket_test.cc new file mode 100644 index 00000000..e114e9a6 --- /dev/null +++ b/connections/implementation/mediums/multiplex/multiplex_socket_test.cc @@ -0,0 +1,361 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "connections/implementation/mediums/multiplex/multiplex_socket.h" + +#include +#include +#include +#include + +#include "gmock/gmock.h" +#include "protobuf-matchers/protocol-buffer-matchers.h" +#include "gtest/gtest.h" +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "connections/implementation/mediums/multiplex/multiplex_frames.h" +#include "connections/implementation/offline_frames.h" +#include "internal/platform/base64_utils.h" +#include "internal/platform/byte_array.h" +#include "internal/platform/count_down_latch.h" +#include "internal/platform/exception.h" +#include "internal/platform/feature_flags.h" +#include "internal/platform/future.h" +#include "internal/platform/input_stream.h" +#include "internal/platform/logging.h" +#include "internal/platform/output_stream.h" +#include "internal/platform/pipe.h" +#include "internal/platform/single_thread_executor.h" +#include "internal/platform/socket.h" +#include "proto/connections_enums.proto.h" + +namespace nearby { +namespace connections { +namespace mediums { +namespace multiplex { + +constexpr absl::string_view SERVICE_ID_1 = "serviceId_1"; +constexpr absl::string_view SERVICE_ID_2 = "serviceId_2"; + +using location::nearby::mediums::MultiplexFrame; +using location::nearby::mediums::MultiplexControlFrame; +using location::nearby::mediums::ConnectionResponseFrame; +using location::nearby::proto::connections::Medium; +using location::nearby::proto::connections::Medium_Name; + +// A fake socket for testing. +class FakeSocket : public MediumSocket { + public: + explicit FakeSocket(Medium medium) : MediumSocket(medium) { + pipe_1_ = CreatePipe(); + reader_1_ = std::move(pipe_1_.first); + writer_1_ = std::move(pipe_1_.second); + pipe_2_ = CreatePipe(); + reader_2_ = std::move(pipe_2_.first); + writer_2_ = std::move(pipe_2_.second); + NEARBY_LOGS(WARNING) << "Physical Socket Medium:" + << Medium_Name(GetMedium()); + }; + ~FakeSocket() override = default; + + FakeSocket(const FakeSocket&) = default; + FakeSocket& operator=(const FakeSocket&) = default; + + /** + * The constructor for a virtual socket which own the virtual {@link + * OutputStream} and {@link InputStream}. + */ + explicit FakeSocket(Medium medium, OutputStream* virtualOutputStream) + : MediumSocket(medium), is_virtual_socket_(true) { + pipe_1_ = CreatePipe(); + reader_1_ = std::move(pipe_1_.first); + writer_1_ = std::move(pipe_1_.second); + pipe_2_ = CreatePipe(); + reader_2_ = std::move(pipe_2_.first); + writer_2_ = std::move(pipe_2_.second); + } + + InputStream& GetInputStream() override { return *reader_1_; } + OutputStream& GetOutputStream() override { return *writer_2_; } + void Close() override { + reader_1_->Close(); + reader_2_->Close(); + writer_1_->Close(); + writer_2_->Close(); + } + + MediumSocket* CreateVirtualSocket( + const std::string& salted_service_id_hash_key, OutputStream* outputstream, + Medium medium, + absl::flat_hash_map>* + virtual_sockets_ptr) override { + if (IsVirtualSocket()) { + NEARBY_LOGS(WARNING) + << "Creating the virtual socket on a virtual socket is not allowed."; + return nullptr; + } + + auto virtual_socket = std::make_shared(medium, outputstream); + NEARBY_LOGS(WARNING) << "Created the virtual socket for Medium: " + << Medium_Name(virtual_socket->GetMedium()); + + if (virtual_sockets_ptr_ == nullptr) { + virtual_sockets_ptr_ = virtual_sockets_ptr; + } + + (*virtual_sockets_ptr_)[salted_service_id_hash_key] = virtual_socket; + NEARBY_LOGS(INFO) << "virtual_sockets_ size: " + << virtual_sockets_ptr_->size(); + return virtual_socket.get(); + } + + void FeedIncomingData(ByteArray data) override { + bytes_read_future_.Set(data); + NEARBY_LOGS(INFO) << "FeedIncomingData. Size of receive data: " + << data.size() << ", bytes content:" << std::string(data); + } + + bool IsVirtualSocket() override { return is_virtual_socket_; } + Future& GetByteReadFuture() { return bytes_read_future_; } + + std::pair, std::unique_ptr> + pipe_1_; + std::unique_ptr reader_1_; + std::unique_ptr writer_1_; + std::pair, std::unique_ptr> + pipe_2_; + std::unique_ptr reader_2_; + std::unique_ptr writer_2_; + + private: + bool is_virtual_socket_ = false; + Future bytes_read_future_; + absl::flat_hash_map>* + virtual_sockets_ptr_ = nullptr; +}; + +TEST(MultiplexSocketTest, CreateSuccessAndReaderThreadStarted) { + testing::NiceMock fake_socket{Medium::WIFI_LAN}; + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_1), + Medium::WIFI_LAN); + MultiplexSocket* multiplex_socket_incoming = + MultiplexSocket::CreateIncomingSocket(&fake_socket, + std::string(SERVICE_ID_1)); + + ASSERT_NE(multiplex_socket_incoming, nullptr); + FakeSocket* virtual_socket = + (FakeSocket*)multiplex_socket_incoming->GetVirtualSocket( + std::string(SERVICE_ID_1)); + if (virtual_socket == nullptr) { + NEARBY_LOGS(INFO) << "Virtual socket not found for " << SERVICE_ID_1; + return; + } + + SingleThreadExecutor executor; + CountDownLatch latch(1); + executor.Execute([&fake_socket, &latch]() { + ByteArray connection_req_frame = parser::ForConnectionRequestConnections( + {}, { + .local_endpoint_id = "endpoint1", + .local_endpoint_info = ByteArray("endpoint1 info"), + }); + auto& writer = fake_socket.writer_1_; + NEARBY_LOGS(INFO) << "writer_1_ Write start"; + writer->Write(Base64Utils::IntToBytes(connection_req_frame.size())); + writer->Write(connection_req_frame); + writer->Flush(); + NEARBY_LOGS(INFO) << "writer_1_ Write end"; + latch.CountDown(); + }); + + latch.Await(absl::Milliseconds(100)); + ExceptionOr result = virtual_socket->GetByteReadFuture().Get(); + if (!result.ok()) { + ADD_FAILURE() << "Read error: " << result.GetException().value; + } + ByteArray data = result.result(); + NEARBY_LOGS(INFO) << "Received " << data.size() << " bytes of data."; + EXPECT_NE(data.size(), 0); + fake_socket.reader_1_->Close(); + EXPECT_EQ(multiplex_socket_incoming->GetVirtualSocketCount(), 1); + multiplex_socket_incoming->Shutdown(); + EXPECT_EQ(multiplex_socket_incoming->GetVirtualSocketCount(), 0); +} + +TEST(MultiplexSocketTest, CreateFail_MediumNotSupport) { + testing::NiceMock fake_socket{Medium::WEB_RTC}; + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_1), + Medium::WEB_RTC); + MultiplexSocket* multiplex_socket_incoming = + MultiplexSocket::CreateIncomingSocket(&fake_socket, + std::string(SERVICE_ID_1)); + + ASSERT_EQ(multiplex_socket_incoming, nullptr); +} + +TEST(MultiplexSocketTest, + EstablishVirtualSocket_ReturnNullWhenMultiplexSocketDisabled) { + testing::NiceMock fake_socket{Medium::WIFI_LAN}; + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_1), + Medium::WIFI_LAN); + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_2), + Medium::WIFI_LAN); + MultiplexSocket* multiplex_socket = MultiplexSocket::CreateOutgoingSocket( + &fake_socket, std::string(SERVICE_ID_1)); + ASSERT_NE(multiplex_socket, nullptr); + + MediumSocket* socket = + multiplex_socket->EstablishVirtualSocket(std::string(SERVICE_ID_2)); + EXPECT_EQ(socket, nullptr); + absl::SleepFor(absl::Milliseconds(100)); + + fake_socket.reader_1_->Close(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 1); + multiplex_socket->Shutdown(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0); +} + +TEST(MultiplexSocketTest, + EstablishVirtualSocket_TimeoutBecauseNoConnectionResponse) { + testing::NiceMock fake_socket{Medium::BLUETOOTH}; + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_1), + Medium::BLUETOOTH); + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_2), + Medium::BLUETOOTH); + MultiplexSocket* multiplex_socket = MultiplexSocket::CreateOutgoingSocket( + &fake_socket, std::string(SERVICE_ID_1)); + ASSERT_NE(multiplex_socket, nullptr); + multiplex_socket->Enable(); + + SingleThreadExecutor executor; + CountDownLatch latch(1); + executor.Execute([&multiplex_socket, &latch]() { + NEARBY_LOGS(INFO) << "EstablishVirtualSocket"; + MediumSocket* socket = + multiplex_socket->EstablishVirtualSocket(std::string(SERVICE_ID_2)); + NEARBY_LOGS(INFO) << "EstablishVirtualSocket finished"; + EXPECT_EQ(socket, nullptr); + latch.CountDown(); + }); + latch.Await(absl::Milliseconds(3000)); + + auto reader = fake_socket.reader_2_.get(); + NEARBY_LOGS(INFO) << "reader_2_ Read start"; + ExceptionOr read_int = Base64Utils::ReadInt(reader); + if (!read_int.ok()) { + ADD_FAILURE() << "Failed to read. Exception:" + << read_int.exception(); + } + auto length = read_int.result(); + NEARBY_LOGS(INFO) << " length:" << length; + EXPECT_GT(length, 0); + EXPECT_EQ(multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_2)), + nullptr); + + absl::SleepFor(absl::Milliseconds(100)); + fake_socket.reader_1_->Close(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 1); + multiplex_socket->Shutdown(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0); +} + +TEST(MultiplexSocketTest, + EstablishVirtualSocket_RemoteAccepted) { + testing::NiceMock fake_socket{Medium::BLE}; + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_1), + Medium::BLE); + MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_2), + Medium::BLE); + + MultiplexSocket* multiplex_socket = MultiplexSocket::CreateOutgoingSocket( + &fake_socket, std::string(SERVICE_ID_1)); + ASSERT_NE(multiplex_socket, nullptr); + multiplex_socket->Enable(); + + SingleThreadExecutor executor; + executor.Execute([&multiplex_socket]() { + NEARBY_LOGS(INFO) << "EstablishVirtualSocket"; + MediumSocket* socket = + multiplex_socket->EstablishVirtualSocket(std::string(SERVICE_ID_2)); + EXPECT_NE(socket, nullptr); + }); + + auto reader = fake_socket.reader_2_.get(); + NEARBY_LOGS(INFO) << "reader_2_ Waiting for CONNECTION_REQUEST frame."; + ExceptionOr read_int = Base64Utils::ReadInt(reader); + if (!read_int.ok()) { + ADD_FAILURE() << "Failed to read length.Exception:" << read_int.exception(); + } + auto length = read_int.result(); + if (length < 0 || + length > + FeatureFlags::GetInstance().GetFlags().connection_max_frame_length) { + ADD_FAILURE() << "Invalid length:" << length; + } + + auto bytes = reader->ReadExactly(length); + if (!bytes.ok()) { + ADD_FAILURE() << "Failed to read frame. Exception:" << bytes.exception(); + } + length = read_int.result(); + if (length < 0 || + length > + FeatureFlags::GetInstance().GetFlags().connection_max_frame_length) { + ADD_FAILURE() << "Invalid frame length:" << length; + } + + ExceptionOr frame_exc = + multiplex::FromBytes(bytes.result()); + if (!frame_exc.ok()) { + ADD_FAILURE() << "Failed to parse MultiplexFrame. Exception:" + << frame_exc.exception(); + } + auto frame = frame_exc.result(); + auto salted_service_id_hash = + ByteArray{std::move(frame.header().salted_service_id_hash())}; + auto service_id_hash_salt = frame.header().has_service_id_hash_salt() + ? frame.header().service_id_hash_salt() + : ""; + ASSERT_EQ(frame.frame_type(), MultiplexFrame::CONTROL_FRAME); + auto control_frame = frame.control_frame(); + ASSERT_EQ(control_frame.control_frame_type(), + MultiplexControlFrame::CONNECTION_REQUEST); + NEARBY_LOGS(INFO) << "Recieved MultiplexControlFrame::CONNECTION_REQUEST " + "frame, now send CONNECTION_RESPONSE frame."; + + ByteArray connection_response_frame = + ForConnectionResponse(salted_service_id_hash, service_id_hash_salt, + ConnectionResponseFrame::CONNECTION_ACCEPTED); + auto& writer = fake_socket.writer_1_; + NEARBY_LOGS(INFO) << "writer_1_ Write start"; + writer->Write(Base64Utils::IntToBytes(connection_response_frame.size())); + writer->Write(connection_response_frame); + writer->Flush(); + NEARBY_LOGS(INFO) << "writer_1_ Write end"; + absl::SleepFor(absl::Milliseconds(100)); + EXPECT_NE(multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_2)), + nullptr); + + fake_socket.reader_1_->Close(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 2); + multiplex_socket->Shutdown(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0); +} + +} // namespace multiplex +} // namespace mediums +} // namespace connections +} // namespace nearby diff --git a/connections/implementation/mediums/utils.cc b/connections/implementation/mediums/utils.cc index a6059145..9839fa45 100644 --- a/connections/implementation/mediums/utils.cc +++ b/connections/implementation/mediums/utils.cc @@ -14,14 +14,22 @@ #include "connections/implementation/mediums/utils.h" -#include +#include +#include #include +#include "internal/platform/base64_utils.h" +#include "internal/platform/byte_array.h" +#include "internal/platform/crypto.h" //NOLINT #include "internal/platform/prng.h" -#include "internal/platform/crypto.h" namespace nearby { namespace connections { + +namespace { +constexpr int kDefaultSaltLength = 16; +} // namespace + using ::location::nearby::connections::LocationHint; using ::location::nearby::connections::LocationStandard; @@ -70,5 +78,13 @@ LocationHint Utils::BuildLocationHint(const std::string& location) { return location_hint; } +// Generates salts. +std::string Utils::GenerateSalt() { return GenerateSalt(kDefaultSaltLength); } + +std::string Utils::GenerateSalt(size_t length) { + ByteArray salt = GenerateRandomBytes(length); + return Base64Utils::Encode(salt); +} + } // namespace connections } // namespace nearby diff --git a/connections/implementation/mediums/utils.h b/connections/implementation/mediums/utils.h index 578cf300..5097a123 100644 --- a/connections/implementation/mediums/utils.h +++ b/connections/implementation/mediums/utils.h @@ -15,7 +15,7 @@ #ifndef CORE_INTERNAL_MEDIUMS_UTILS_H_ #define CORE_INTERNAL_MEDIUMS_UTILS_H_ -#include +#include #include #include "connections/implementation/proto/offline_wire_formats.pb.h" @@ -31,6 +31,8 @@ class Utils { static ByteArray Sha256Hash(const std::string& source, size_t length); static location::nearby::connections::LocationHint BuildLocationHint( const std::string& location); + static std::string GenerateSalt(); + static std::string GenerateSalt(size_t length); }; } // namespace connections diff --git a/connections/implementation/offline_frames.cc b/connections/implementation/offline_frames.cc index 03ec322c..925025ab 100644 --- a/connections/implementation/offline_frames.cc +++ b/connections/implementation/offline_frames.cc @@ -171,8 +171,8 @@ ByteArray ForConnectionRequestPresence( return ToBytes(std::move(frame)); } -ByteArray ForConnectionResponse( - std::int32_t status, const OsInfo& os_info) { +ByteArray ForConnectionResponse(std::int32_t status, const OsInfo& os_info, + std::int32_t multiplex_socket_bitmask) { OfflineFrame frame; frame.set_version(OfflineFrame::V1); @@ -188,6 +188,7 @@ ByteArray ForConnectionResponse( ? ConnectionResponseFrame::ACCEPT : ConnectionResponseFrame::REJECT); *sub_frame->mutable_os_info() = os_info; + sub_frame->set_multiplex_socket_bitmask(multiplex_socket_bitmask); sub_frame->set_safe_to_disconnect_version( NearbyFlags::GetInstance().GetInt64Flag( config_package_nearby::nearby_connections_feature:: diff --git a/connections/implementation/offline_frames.h b/connections/implementation/offline_frames.h index 33b4a2a6..c9e9e71c 100644 --- a/connections/implementation/offline_frames.h +++ b/connections/implementation/offline_frames.h @@ -53,7 +53,8 @@ ByteArray ForConnectionRequestPresence( const location::nearby::connections::PresenceDevice& proto_presence_device, const ConnectionInfo& connection_info); ByteArray ForConnectionResponse( - std::int32_t status, const location::nearby::connections::OsInfo& os_info); + std::int32_t status, const location::nearby::connections::OsInfo& os_info, + std::int32_t multiplex_socket_bitmask); // Builds Payload transfer messages. ByteArray ForDataPayloadTransfer( diff --git a/connections/implementation/offline_frames_test.cc b/connections/implementation/offline_frames_test.cc index b393daf9..c200c6b5 100644 --- a/connections/implementation/offline_frames_test.cc +++ b/connections/implementation/offline_frames_test.cc @@ -270,6 +270,7 @@ TEST(OfflineFramesTest, CanGenerateConnectionResponse) { status: 1 response: REJECT os_info { type: LINUX } + multiplex_socket_bitmask: 0x01 safe_to_disconnect_version: 5 > >)pb"; @@ -277,9 +278,11 @@ TEST(OfflineFramesTest, CanGenerateConnectionResponse) { OsInfo os_info; os_info.set_type(OsInfo::LINUX); NearbyFlags::GetInstance().OverrideInt64FlagValue( - config_package_nearby::nearby_connections_feature:: - kSafeToDisconnectVersion, 5); - ByteArray bytes = ForConnectionResponse(1, os_info); + config_package_nearby::nearby_connections_feature:: + kSafeToDisconnectVersion, + 5); + ByteArray bytes = + ForConnectionResponse(1, os_info, /*multiplex_socket_bitmask=*/0x01); auto response = FromBytes(bytes); ASSERT_TRUE(response.ok()); OfflineFrame message = response.result(); diff --git a/connections/implementation/offline_frames_validator_test.cc b/connections/implementation/offline_frames_validator_test.cc index 02e89ca3..1940c9d0 100644 --- a/connections/implementation/offline_frames_validator_test.cc +++ b/connections/implementation/offline_frames_validator_test.cc @@ -174,7 +174,8 @@ TEST(OfflineFramesValidatorTest, OfflineFrame offline_frame; OsInfo os_info; - ByteArray bytes = ForConnectionResponse(kStatusAccepted, os_info); + ByteArray bytes = ForConnectionResponse(kStatusAccepted, os_info, + /*multiplex_socket_bitmask=*/0); offline_frame.ParseFromString(std::string(bytes)); auto ret_value = EnsureValidOfflineFrame(offline_frame); @@ -187,7 +188,8 @@ TEST(OfflineFramesValidatorTest, OfflineFrame offline_frame; OsInfo os_info; - ByteArray bytes = ForConnectionResponse(kStatusAccepted, os_info); + ByteArray bytes = ForConnectionResponse(kStatusAccepted, os_info, + /*multiplex_socket_bitmask=*/0); offline_frame.ParseFromString(std::string(bytes)); auto* v1_frame = offline_frame.mutable_v1(); @@ -203,7 +205,8 @@ TEST(OfflineFramesValidatorTest, OfflineFrame offline_frame; OsInfo os_info; - ByteArray bytes = ForConnectionResponse(-1, os_info); + ByteArray bytes = + ForConnectionResponse(-1, os_info, /*multiplex_socket_bitmask=*/0); offline_frame.ParseFromString(std::string(bytes)); auto ret_value = EnsureValidOfflineFrame(offline_frame); diff --git a/internal/platform/BUILD b/internal/platform/BUILD index af5a3641..a96b2863 100644 --- a/internal/platform/BUILD +++ b/internal/platform/BUILD @@ -59,6 +59,7 @@ cc_library( deps = [ "//proto:connections_enums_cc_proto", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/meta:type_traits", "@com_google_absl//absl/strings", diff --git a/internal/platform/feature_flags.h b/internal/platform/feature_flags.h index 5d18a178..c9e9c9b5 100644 --- a/internal/platform/feature_flags.h +++ b/internal/platform/feature_flags.h @@ -107,6 +107,9 @@ class FeatureFlags { // The new outgoing frame with the middle priority will wait for space to // become available if the queue is full.' std::uint32_t multiplex_socket_middle_priority_queue_capacity = 50; + // The maximum size of frame we'll attempt to read, to avoid a remote device + // from triggering an OutOfMemory error. + std::uint32_t connection_max_frame_length = 1048576; }; static const FeatureFlags& GetInstance() { diff --git a/internal/platform/socket.h b/internal/platform/socket.h index 43acd77c..16fc983e 100644 --- a/internal/platform/socket.h +++ b/internal/platform/socket.h @@ -15,8 +15,17 @@ #ifndef PLATFORM_BASE_SOCKET_H_ #define PLATFORM_BASE_SOCKET_H_ +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/functional/any_invocable.h" +#include "internal/platform/byte_array.h" #include "internal/platform/input_stream.h" #include "internal/platform/output_stream.h" +#include "proto/connections_enums.pb.h" namespace nearby { @@ -32,6 +41,72 @@ class Socket { virtual void Close() = 0; }; +class MediumSocket : public Socket { + public: + explicit MediumSocket(location::nearby::proto::connections::Medium medium) + : medium_(medium) {} + ~MediumSocket() override = default; + + /** Returns the medium of the socket. */ + virtual location::nearby::proto::connections::Medium GetMedium() const { + return medium_; + } + + /** Creates a virtual socket. */ + virtual MediumSocket* CreateVirtualSocket( + const std::string& salted_service_id_hash_key, OutputStream* outputstream, + location::nearby::proto::connections::Medium medium, + absl::flat_hash_map>* + virtual_sockets_ptr) { + return this; + } + + /** Feeds the received incoming data to the client. */ + virtual void FeedIncomingData(ByteArray data) {} + + /** Returns true if the socket is a virtual socket. */ + virtual bool IsVirtualSocket() { + return false; + } + + /** Adds a listener to be invoked when the socket is closed. */ + void AddOnSocketClosedListener( + std::unique_ptr> socket_closed_listener) { + socket_closed_listeners_.insert(std::move(socket_closed_listener)); + } + + /** Adds a listener to be invoked when the multiplex socket is enabled. */ + void RegisterMultiplexEnabledCallback( + std::shared_ptr> callback) { + multiplex_socket_enabled_cbs_.insert(std::move(callback)); + } + + /** Enables the multiplex socket. */ + void EnableMultiplexSocket() { + if (!IsVirtualSocket()) { + return; + } + + for (auto& callback : multiplex_socket_enabled_cbs_) { + callback.get(); + } + } + + /** Closes the local socket. */ + void CloseLocal() { + for (auto& listener : socket_closed_listeners_) { + listener.get(); + } + } + + private: + location::nearby::proto::connections::Medium medium_; + absl::flat_hash_set>> + socket_closed_listeners_; + absl::flat_hash_set>> + multiplex_socket_enabled_cbs_; +}; + } // namespace nearby #endif // PLATFORM_BASE_SOCKET_H_