[Multiplex] Fix the issue that the physical socket is not handled properly.

PiperOrigin-RevId: 660594414

connections/implementation/mediums/bluetooth_classic.cc

@@ -84,6 +84,7 @@ BluetoothClassic::~BluetoothClassic() {
         multiplex_socket->Shutdown();
       }
     }
+    multiplex_sockets_.clear();
   }
 
   // All the AcceptLoopRunnable objects in here should already have gotten an
@@ -403,9 +404,13 @@ bool BluetoothClassic::StartAcceptingConnections(
           {
             MutexLock lock(&mutex_);
             if (is_multiplex_enabled_) {
+              BluetoothSocket client_socket_bak = client_socket;
+              auto physical_socket_ptr =
+                  std::make_shared<BluetoothSocket>(client_socket_bak);
               MultiplexSocket* multiplex_socket =
-                  MultiplexSocket::CreateIncomingSocket(&client_socket,
-                                                        service_id);
+                  MultiplexSocket::CreateIncomingSocket(
+                      physical_socket_ptr, service_id);
+
               if (multiplex_socket != nullptr &&
                   multiplex_socket->GetVirtualSocket(service_id)) {
                 multiplex_sockets_.emplace(
@@ -588,8 +593,11 @@ BluetoothSocket BluetoothClassic::AttemptToConnect(
   if (is_multiplex_enabled_) {
     // New MultiplexSocket but default disabled, should be enabled after
     // negotiated
-    MultiplexSocket* multiplex_socket =
-        MultiplexSocket::CreateOutgoingSocket(&socket, service_id);
+    auto physical_socket_ptr =
+        std::make_shared<BluetoothSocket>(socket);
+    MultiplexSocket* multiplex_socket = MultiplexSocket::CreateOutgoingSocket(
+        std::move(physical_socket_ptr), service_id);
+
     auto* virtual_socket = multiplex_socket->GetVirtualSocket(service_id);
     // Should not happen.
     auto* bluetooth_socket = down_cast<BluetoothSocket*>(virtual_socket);

connections/implementation/mediums/multiplex/multiplex_output_stream.cc

@@ -178,7 +178,6 @@ MultiplexOutputStream::MultiplexWriter::MultiplexWriter(
 
 MultiplexOutputStream::MultiplexWriter::~MultiplexWriter() {
   Close();
-  // writer_thread_.Shutdown();
   physical_writer_ = nullptr;
 }
 
@@ -338,6 +337,7 @@ Exception MultiplexOutputStream::VirtualOutputStream::Flush() {
 }
 
 Exception MultiplexOutputStream::VirtualOutputStream::Close() {
+  NEARBY_LOGS(INFO) << "MultiplexOutputStream::VirtualOutputStream::Close";
   is_closed_.Set(true);
   return {Exception::kSuccess};
 }

connections/implementation/mediums/multiplex/multiplex_socket.cc

@@ -79,26 +79,12 @@ void MultiplexSocket::StopListeningForIncomingConnection(
       std::pair<std::string, Medium>(service_id, type));
 }
 
-MultiplexSocket::MultiplexSocket(MediumSocket* physical_socket)
-    : physical_socket_(physical_socket),
-      multiplex_output_stream_{&physical_socket->GetOutputStream(), enabled_},
-      physical_reader_(&physical_socket->GetInputStream()) {
-  if (physical_socket->IsFakeSocket()) {
-    return;
-  }
-  NEARBY_LOGS(INFO) << "physical_socket_: " << physical_socket_;
-  switch (physical_socket_->GetMedium()) {
-    case Medium::BLUETOOTH:
-      medium_ = Medium::BLUETOOTH;
-      bluetooth_socket_ =
-          std::move(*static_cast<BluetoothSocket*>(physical_socket));
-      break;
-    default:
-      medium_ = Medium::UNKNOWN_MEDIUM;
-      NEARBY_LOGS(ERROR) << __func__ << "Unsupported medium: "
-                         << physical_socket_->GetMedium();
-  }
-}
+MultiplexSocket::MultiplexSocket(std::shared_ptr<MediumSocket> physical_socket)
+    : physical_socket_ptr_(physical_socket),
+      multiplex_output_stream_{&physical_socket_ptr_->GetOutputStream(),
+                               enabled_},
+      physical_reader_(&physical_socket_ptr_->GetInputStream()),
+      medium_(physical_socket_ptr_->GetMedium()){}
 
 absl::flat_hash_map<std::pair<std::string, Medium>,
                     MultiplexIncomingConnectionCb>&
@@ -119,7 +105,8 @@ MultiplexSocket::GetIncomingConnectionCallbacks() {
 }
 
 MultiplexSocket* MultiplexSocket::CreateIncomingSocket(
-    MediumSocket* physical_socket, const std::string& service_id) {
+    std::shared_ptr<MediumSocket> physical_socket,
+    const std::string& service_id) {
   static MultiplexSocket* multiplex_incoming_socket = nullptr;
   switch (physical_socket->GetMedium()) {
     case Medium::BLUETOOTH:
@@ -161,8 +148,8 @@ MultiplexSocket* MultiplexSocket::CreateIncomingSocket(
 }
 
 MultiplexSocket* MultiplexSocket::CreateOutgoingSocket(
-    MediumSocket* physical_socket, const std::string& service_id,
-    const std::string& service_id_hash_salt) {
+    std::shared_ptr<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:
@@ -201,7 +188,8 @@ MultiplexSocket* MultiplexSocket::CreateOutgoingSocket(
 }
 
 MultiplexSocket* MultiplexSocket::CreateOutgoingSocket(
-    MediumSocket* physical_socket, const std::string& service_id) {
+    std::shared_ptr<MediumSocket> physical_socket,
+    const std::string& service_id) {
   return CreateOutgoingSocket(physical_socket, service_id,
                               Utils::GenerateSalt());
 }
@@ -219,9 +207,9 @@ MediumSocket* MultiplexSocket::CreateFirstVirtualSocket(
                     << ", salt=" << service_id_hash_salt
                     << ", salted_service_id_hash_key="
                     << salted_service_id_hash_key;
-  MediumSocket* virtual_socket = physical_socket_->CreateVirtualSocket(
-      salted_service_id_hash_key, output_stream, physical_socket_->GetMedium(),
-      &virtual_sockets_);
+  MediumSocket* virtual_socket = physical_socket_ptr_->CreateVirtualSocket(
+      salted_service_id_hash_key, output_stream,
+      medium_, &virtual_sockets_);
 
   virtual_socket->AddOnSocketClosedListener(
       std::make_unique<absl::AnyInvocable<void()>>(
@@ -248,9 +236,9 @@ MediumSocket* MultiplexSocket::CreateVirtualSocket(
                     << ", salted_service_id_hash_key="
                     << salted_service_id_hash_key;
 
-  MediumSocket* virtual_socket = physical_socket_->CreateVirtualSocket(
-      salted_service_id_hash_key, output_stream, physical_socket_->GetMedium(),
-      &virtual_sockets_);
+  MediumSocket* virtual_socket = physical_socket_ptr_->CreateVirtualSocket(
+      salted_service_id_hash_key, output_stream,
+      medium_, &virtual_sockets_);
 
   virtual_socket->AddOnSocketClosedListener(
       std::make_unique<absl::AnyInvocable<void()>>(
@@ -498,7 +486,7 @@ void MultiplexSocket::HandleConnectionRequest(
     const std::string& service_id_hash_salt) {
   if (!IsEnabled()) {
     NEARBY_LOGS(WARNING) << "Received a CONNECTION_REQUEST frame on medium "
-                         << Medium_Name(physical_socket_->GetMedium())
+                         << Medium_Name(medium_)
                          << " but status is disabled, ignore it.";
     return;
   }
@@ -522,7 +510,7 @@ void MultiplexSocket::HandleConnectionRequest(
                       << service_id_hash_salt
                       << ", hash key : " << salted_service_id_hash_key
                       << " on medium "
-                      << Medium_Name(physical_socket_->GetMedium());
+                      << Medium_Name(medium_);
 
     NEARBY_LOGS(INFO) << "The size of incomingConnectionCallbacks : "
                       << GetIncomingConnectionCallbacks().size();
@@ -538,7 +526,7 @@ void MultiplexSocket::HandleConnectionRequest(
                     << ", hash salt : " << service_id_hash_salt
                     << ", hash key : " << salted_service_id_hash_key
                     << " on medium "
-                    << Medium_Name(physical_socket_->GetMedium());
+                    << Medium_Name(medium_);
 
   if (!multiplex_output_stream_.WriteConnectionResponseFrame(
           salted_service_id_hash, service_id_hash_salt,
@@ -744,19 +732,8 @@ void MultiplexSocket::Shutdown() {
   }
 
   multiplex_output_stream_.Shutdown();
+  physical_socket_ptr_->Close();
 
-  if (!physical_socket_->IsFakeSocket()) {
-    switch (medium_) {
-      case Medium::BLUETOOTH:
-        bluetooth_socket_.Close();
-        break;
-      case Medium::UNKNOWN_MEDIUM:
-        NEARBY_LOGS(INFO) << __func__ << " Unknown medium";
-        break;
-      default:
-        break;
-    }
-  }
   if (reader_thread_shutdown_barrier_) {
     reader_thread_shutdown_barrier_->Await(kTimeoutForReaderThreadStop);
   }

connections/implementation/mediums/multiplex/multiplex_socket.h

@@ -19,14 +19,11 @@
 #include <string>
 #include <utility>
 
-#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 "connections/medium_selector.h"
 #include "internal/platform/atomic_boolean.h"
-#include "internal/platform/ble.h"
-#include "internal/platform/bluetooth_classic.h"
 #include "internal/platform/byte_array.h"
 #include "internal/platform/count_down_latch.h"
 #include "internal/platform/future.h"
@@ -35,7 +32,6 @@
 #include "internal/platform/mutex.h"
 #include "internal/platform/single_thread_executor.h"
 #include "internal/platform/socket.h"
-#include "internal/platform/wifi_lan.h"
 #include "proto/connections_enums.pb.h"
 #include "proto/mediums/multiplex_frames.pb.h"
 
@@ -54,16 +50,18 @@ class MultiplexSocket {
   MultiplexSocket& operator=(const MultiplexSocket&) = delete;
 
   // Creates a new incoming MultiplexSocket.
-  static MultiplexSocket* CreateIncomingSocket(MediumSocket* physical_socket,
-                                               const std::string& service_id);
+  static MultiplexSocket* CreateIncomingSocket(
+      std::shared_ptr<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);
+      std::shared_ptr<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);
+  static MultiplexSocket* CreateOutgoingSocket(
+      std::shared_ptr<MediumSocket> physical_socket,
+      const std::string& service_id);
 
   // A Table of service Id as row key, medium type as column key, and
   // MultiplexIncomingConnectionCb as value. Non-empty while the client starts
@@ -96,8 +94,6 @@ class MultiplexSocket {
     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.
@@ -113,7 +109,7 @@ class MultiplexSocket {
   void SetShutdown(bool is_shutdown) { is_shutdown_ = is_shutdown; }
 
  private:
-  explicit MultiplexSocket(MediumSocket* physical_socket);
+  explicit MultiplexSocket(std::shared_ptr<MediumSocket> physical_socket);
   ~MultiplexSocket() = default;
 
   // Creates the first virtual socket for the service id. The first virtual
@@ -167,7 +163,7 @@ class MultiplexSocket {
                         absl::AnyInvocable<void()> runnable);
 
   // The physical socket connect to the remote device.
-  MediumSocket* physical_socket_;
+  std::shared_ptr<MediumSocket> physical_socket_ptr_;
 
   // The output stream to manage all outgoing frames from all clients.
   MultiplexOutputStream multiplex_output_stream_;
@@ -176,10 +172,6 @@ class MultiplexSocket {
   InputStream* physical_reader_;
   // The medium type of the physical socket.
   Medium medium_;
-  // Save the phyical socket here, so it can be closed when all the virtual
-  // socket is gone.
-  BluetoothSocket bluetooth_socket_;
-  WifiLanSocket wifi_lan_socket_;
 
   // The callback to enable the MultiplexSocket.
   std::shared_ptr<absl::AnyInvocable<void()>> enable_cb_ =

connections/implementation/mediums/multiplex/multiplex_socket_test.cc

@@ -88,9 +88,6 @@ class FakeSocket : public MediumSocket {
     writer_2_ = std::move(pipe_2_.second);
   }
 
-  bool IsFakeSocket() override{
-    return true;
-  }
   InputStream& GetInputStream() override { return *reader_1_; }
   OutputStream& GetOutputStream() override { return *writer_2_; }
   Exception Close() override {
@@ -158,11 +155,12 @@ class FakeSocket : public MediumSocket {
 };
 
 TEST(MultiplexSocketTest, CreateSuccessAndReaderThreadStarted) {
-  testing::NiceMock<FakeSocket> fake_socket{Medium::BLUETOOTH};
+  auto fake_socket_ptr =
+      std::make_shared<FakeSocket>(Medium::BLUETOOTH);
   MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_1),
                                                       Medium::BLUETOOTH);
   MultiplexSocket* multiplex_socket_incoming =
-      MultiplexSocket::CreateIncomingSocket(&fake_socket,
+      MultiplexSocket::CreateIncomingSocket(fake_socket_ptr,
                                             std::string(SERVICE_ID_1));
   ASSERT_NE(multiplex_socket_incoming, nullptr);
   FakeSocket* virtual_socket =
@@ -174,14 +172,15 @@ TEST(MultiplexSocketTest, CreateSuccessAndReaderThreadStarted) {
   }
 
   SingleThreadExecutor executor;
+  FakeSocket* socket = fake_socket_ptr.get();
   CountDownLatch latch(1);
-  executor.Execute([&fake_socket, &latch]() {
+  executor.Execute([socket, &latch]() {
     ByteArray connection_req_frame = parser::ForConnectionRequestConnections(
         {}, {
                 .local_endpoint_id = "endpoint1",
                 .local_endpoint_info = ByteArray("endpoint1 info"),
             });
-    auto& writer = fake_socket.writer_1_;
+    auto& writer = socket->writer_1_;
     NEARBY_LOGS(INFO) << "writer_1_ Write start";
     writer->Write(Base64Utils::IntToBytes(connection_req_frame.size()));
     writer->Write(connection_req_frame);
@@ -199,18 +198,18 @@ TEST(MultiplexSocketTest, CreateSuccessAndReaderThreadStarted) {
   NEARBY_LOGS(INFO) << "Received " << data.size() << " bytes of data.";
   EXPECT_NE(data.size(), 0);
   absl::SleepFor(absl::Milliseconds(100));
-  fake_socket.reader_1_->Close();
+  socket->reader_1_->Close();
   EXPECT_EQ(multiplex_socket_incoming->GetVirtualSocketCount(), 1);
   virtual_socket->Close();
   EXPECT_EQ(multiplex_socket_incoming->GetVirtualSocketCount(), 0);
 }
-
 TEST(MultiplexSocketTest, CreateFail_MediumNotSupport) {
-  testing::NiceMock<FakeSocket> fake_socket{Medium::WEB_RTC};
+  auto fake_socket_ptr =
+      std::make_shared<FakeSocket>(Medium::WEB_RTC);
   MultiplexSocket::StopListeningForIncomingConnection(std::string(SERVICE_ID_1),
                                                       Medium::WEB_RTC);
   MultiplexSocket* multiplex_socket_incoming =
-      MultiplexSocket::CreateIncomingSocket(&fake_socket,
+      MultiplexSocket::CreateIncomingSocket(fake_socket_ptr,
                                             std::string(SERVICE_ID_1));
 
   ASSERT_EQ(multiplex_socket_incoming, nullptr);
@@ -218,13 +217,14 @@ TEST(MultiplexSocketTest, CreateFail_MediumNotSupport) {
 
 TEST(MultiplexSocketTest,
      EstablishVirtualSocket_ReturnNullWhenMultiplexSocketDisabled) {
-  testing::NiceMock<FakeSocket> fake_socket{Medium::BLE};
+  auto fake_socket_ptr = std::make_shared<FakeSocket>(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));
+      fake_socket_ptr, std::string(SERVICE_ID_1));
   ASSERT_NE(multiplex_socket, nullptr);
 
   MediumSocket* socket =
@@ -238,7 +238,7 @@ TEST(MultiplexSocketTest,
     NEARBY_LOGS(INFO) << "Virtual socket not found for " << SERVICE_ID_1;
     return;
   }
-  fake_socket.reader_1_->Close();
+  fake_socket_ptr->reader_1_->Close();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 1);
   virtual_socket->Close();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0);
@@ -246,13 +246,13 @@ TEST(MultiplexSocketTest,
 
 TEST(MultiplexSocketTest,
      EstablishVirtualSocket_TimeoutBecauseNoConnectionResponse) {
-  testing::NiceMock<FakeSocket> fake_socket{Medium::WIFI_LAN};
+  auto fake_socket_ptr = std::make_shared<FakeSocket>(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));
+      fake_socket_ptr, std::string(SERVICE_ID_1));
   ASSERT_NE(multiplex_socket, nullptr);
   multiplex_socket->Enable();
   FakeSocket* virtual_socket = (FakeSocket*)multiplex_socket->GetVirtualSocket(
@@ -274,7 +274,7 @@ TEST(MultiplexSocketTest,
   });
   latch.Await(absl::Milliseconds(3000));
 
-  auto reader = fake_socket.reader_2_.get();
+  auto reader = fake_socket_ptr->reader_2_.get();
   NEARBY_LOGS(INFO) << "reader_2_ Read start";
   ExceptionOr<std::int32_t> read_int = Base64Utils::ReadInt(reader);
   if (!read_int.ok()) {
@@ -288,7 +288,7 @@ TEST(MultiplexSocketTest,
             nullptr);
 
   absl::SleepFor(absl::Milliseconds(100));
-  fake_socket.reader_1_->Close();
+  fake_socket_ptr->reader_1_->Close();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 1);
   virtual_socket->Close();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0);
@@ -296,14 +296,14 @@ TEST(MultiplexSocketTest,
 
 TEST(MultiplexSocketTest,
      EstablishVirtualSocket_RemoteAccepted) {
-  testing::NiceMock<FakeSocket> fake_socket{Medium::BLUETOOTH};
+  auto fake_socket_ptr = std::make_shared<FakeSocket>(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));
+      fake_socket_ptr, std::string(SERVICE_ID_1));
   ASSERT_NE(multiplex_socket, nullptr);
   multiplex_socket->Enable();
 
@@ -315,7 +315,7 @@ TEST(MultiplexSocketTest,
     EXPECT_NE(socket, nullptr);
   });
 
-  auto reader = fake_socket.reader_2_.get();
+  auto reader = fake_socket_ptr->reader_2_.get();
   NEARBY_LOGS(INFO) << "reader_2_ Waiting for CONNECTION_REQUEST frame.";
   ExceptionOr<std::int32_t> read_int = Base64Utils::ReadInt(reader);
   if (!read_int.ok()) {
@@ -361,7 +361,7 @@ TEST(MultiplexSocketTest,
   ByteArray connection_response_frame =
       ForConnectionResponse(salted_service_id_hash, service_id_hash_salt,
                             ConnectionResponseFrame::CONNECTION_ACCEPTED);
-  auto& writer = fake_socket.writer_1_;
+  auto& writer = fake_socket_ptr->writer_1_;
   NEARBY_LOGS(INFO) << "writer_1_ Write start";
   writer->Write(Base64Utils::IntToBytes(connection_response_frame.size()));
   writer->Write(connection_response_frame);
@@ -371,7 +371,7 @@ TEST(MultiplexSocketTest,
   EXPECT_NE(multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_2)),
             nullptr);
 
-  fake_socket.reader_1_->Close();
+  fake_socket_ptr->reader_1_->Close();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 2);
   multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_2))->Close();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 1);

internal/platform/blocking_queue_stream.cc

@@ -36,7 +36,8 @@ ExceptionOr<ByteArray> BlockingQueueStream::Read(std::int64_t size) {
         << "Failed to read BlockingQueueStream because it was closed.";
     return ExceptionOr<ByteArray>(Exception::kInterrupted);
   }
-  NEARBY_LOGS(INFO) << "BlockingQueueStream read " << size << " bytes";
+  NEARBY_LOGS(INFO) << "BlockingQueueStream expect to read " << size
+                    << " bytes";
   return ExceptionOr<ByteArray>(blocking_queue_.Take());
 }
 

internal/platform/socket.h

@@ -77,11 +77,6 @@ class MediumSocket : public Socket {
     return false;
   }
 
-  /** Returns true if the socket is a Fake socket for unit test. */
-  virtual bool IsFakeSocket() {
-    return false;
-  }
-
   /** Adds a listener to be invoked when the socket is closed. */
   void AddOnSocketClosedListener(
       std::unique_ptr<absl::AnyInvocable<void()>> socket_closed_listener) {