Fixed a few issues in WiFi LAN feature.

PiperOrigin-RevId: 400786291

cpp/platform/impl/windows/wifi_lan.h

@@ -16,8 +16,8 @@
 #define PLATFORM_IMPL_WINDOWS_WIFI_LAN_H_
 
 // Windows headers
-#include <windows.h> // NOLINT
-#include <win32/windns.h> // NOLINT
+#include <windows.h>       // NOLINT
+#include <win32/windns.h>  // NOLINT
 
 // Standard C/C++ headers
 #include <exception>
@@ -108,8 +108,13 @@ class WifiLanService : public api::WifiLanService {
 // remote WiFi LAN service, also will return a WifiLanSocket to caller.
 class WifiLanSocket : public api::WifiLanSocket {
  public:
-  WifiLanSocket(StreamSocket socket);
+  explicit WifiLanSocket(api::WifiLanService* wifi_lan_service,
+                         StreamSocket socket);
+  WifiLanSocket(WifiLanSocket&) = default;
+  WifiLanSocket(WifiLanSocket&&) = default;
   ~WifiLanSocket() override;
+  WifiLanSocket& operator=(const WifiLanSocket&) = default;
+  WifiLanSocket& operator=(WifiLanSocket&&) = default;
 
   // Returns the InputStream of the WifiLanSocket.
   // On error, returned stream will report Exception::kIo on any operation.
@@ -132,10 +137,6 @@ class WifiLanSocket : public api::WifiLanSocket {
   // nullptr otherwise.
   api::WifiLanService* GetRemoteWifiLanService() override;
 
-  // When connect to remove WiFi LAN servie, need to save remove WiFi LAN
-  // information, so that can return it based on ip address and port query
-  void SetRemoteWifiLanService(api::WifiLanService* wifi_lan_service);
-
   // Sets service id binding to the socket
   void SetServiceId(std::string service_id);
 
@@ -174,13 +175,13 @@ class WifiLanSocket : public api::WifiLanSocket {
     Exception Close() override;
 
    private:
-    IOutputStream output_stream_;
+    IOutputStream output_stream_{nullptr};
   };
 
   // Internal properties
   StreamSocket stream_soket_{nullptr};
-  std::unique_ptr<SocketInputStream> input_stream_{nullptr};
-  std::unique_ptr<SocketOutputStream> output_stream_{nullptr};
+  SocketInputStream input_stream_{nullptr};
+  SocketOutputStream output_stream_{nullptr};
 
   api::WifiLanService* remote_wifi_lan_service_ = nullptr;
   WifiLanMedium* medium_ = nullptr;
@@ -306,7 +307,7 @@ class WifiLanNsd {
   WifiLanMedium* medium_ = nullptr;
   WifiLanService wifi_lan_service_{};
   absl::flat_hash_map<std::string, std::unique_ptr<WifiLanService>>
-      remote_wifi_lan_services_{};
+      remote_wifi_lan_services_ ABSL_GUARDED_BY(mutex_);
 
   // NSD Status
   int nsd_status_ = NSD_STATUS_IDLE;

cpp/platform/impl/windows/wifi_lan_medium.cc

@@ -188,11 +188,11 @@ std::unique_ptr<api::WifiLanSocket> WifiLanMedium::Connect(
     try {
       socket.ConnectAsync(host_name, service_name).get();
       // connected need to keep connection
+
       std::unique_ptr<WifiLanSocket> wifi_lan_socket =
-          std::make_unique<WifiLanSocket>(std::move(socket));
+          std::make_unique<WifiLanSocket>(&wifi_lan_service, socket);
       wifi_lan_socket->SetServiceId(service_id);
       wifi_lan_socket->SetMedium(this);
-      wifi_lan_socket->SetRemoteWifiLanService(&wifi_lan_service);
       {
         MutexLock lock(&mutex_);
         wifi_lan_sockets_.insert(wifi_lan_socket.get());

cpp/platform/impl/windows/wifi_lan_nsd.cc

@@ -340,10 +340,12 @@ fire_and_forget WifiLanNsd::Listener_ConnectionReceived(
   StreamSocket stream_socket = args.Socket();
 
   // Send to callback
-  WifiLanSocket socket{stream_socket};
-  socket.SetMedium(medium_);
-  socket.SetServiceId(service_id_);
-  accepted_connection_callback_.accepted_cb(socket, service_id_);
+  std::unique_ptr<windows::WifiLanSocket> socket =
+      std::make_unique<windows::WifiLanSocket>(&wifi_lan_service_,
+                                               stream_socket);
+  socket->SetMedium(medium_);
+  socket->SetServiceId(service_id_);
+  accepted_connection_callback_.accepted_cb(*socket.get(), service_id_);
   return fire_and_forget{};
 }
 

cpp/platform/impl/windows/wifi_lan_socket.cc

@@ -19,10 +19,12 @@ namespace location {
 namespace nearby {
 namespace windows {
 
-WifiLanSocket::WifiLanSocket(StreamSocket socket) {
+WifiLanSocket::WifiLanSocket(api::WifiLanService* wifi_lan_service,
+                             StreamSocket socket) {
+  remote_wifi_lan_service_ = wifi_lan_service;
   stream_soket_ = socket;
-  input_stream_ = std::make_unique<SocketInputStream>(socket.InputStream());
-  output_stream_ = std::make_unique<SocketOutputStream>(socket.OutputStream());
+  input_stream_ = SocketInputStream(socket.InputStream());
+  output_stream_ = SocketOutputStream(socket.OutputStream());
 }
 
 WifiLanSocket::~WifiLanSocket() {
@@ -35,17 +37,14 @@ WifiLanSocket::~WifiLanSocket() {
   }
 }
 
-InputStream& WifiLanSocket::GetInputStream() { return *input_stream_.get(); }
+InputStream& WifiLanSocket::GetInputStream() { return input_stream_; }
 
-OutputStream& WifiLanSocket::GetOutputStream() { return *output_stream_.get(); }
+OutputStream& WifiLanSocket::GetOutputStream() { return output_stream_; }
 
 Exception WifiLanSocket::Close() {
   try {
     if (stream_soket_ != nullptr) {
       stream_soket_.Close();
-      stream_soket_ = nullptr;
-      input_stream_ = nullptr;
-      output_stream_ = nullptr;
       medium_->CloseConnection(*this);
     }
     return {Exception::kSuccess};
@@ -58,11 +57,6 @@ api::WifiLanService* WifiLanSocket::GetRemoteWifiLanService() {
   return remote_wifi_lan_service_;
 }
 
-void WifiLanSocket::SetRemoteWifiLanService(
-    api::WifiLanService* wifi_lan_service) {
-  remote_wifi_lan_service_ = wifi_lan_service;
-}
-
 void WifiLanSocket::SetServiceId(std::string service_id) {
   service_id_ = service_id;
 }
@@ -98,7 +92,13 @@ ExceptionOr<ByteArray> WifiLanSocket::SocketInputStream::Read(
 
     auto ibuffer =
         input_stream_.ReadAsync(buffer, size, InputStreamOptions::None).get();
+
+    if (ibuffer.Length() != size) {
+      NEARBY_LOGS(WARNING) << "Only got part of data of needed.";
+    }
+
     ByteArray data((char*)ibuffer.data(), ibuffer.Length());
+
     return ExceptionOr(data);
   } catch (...) {
     return Exception{Exception::kIo};
@@ -120,7 +120,7 @@ ExceptionOr<size_t> WifiLanSocket::SocketInputStream::Skip(size_t offset) {
 Exception WifiLanSocket::SocketInputStream::Close() {
   try {
     input_stream_.Close();
-  } catch (std::exception exception) {
+  } catch (...) {
     return {Exception::kIo};
   }
 
@@ -136,10 +136,11 @@ WifiLanSocket::SocketOutputStream::SocketOutputStream(
 Exception WifiLanSocket::SocketOutputStream::Write(const ByteArray& data) {
   Buffer buffer = Buffer(data.size());
   std::memcpy(buffer.data(), data.data(), data.size());
+  buffer.Length(data.size());
 
   try {
-    output_stream_.WriteAsync(buffer);
-  } catch (std::exception exception) {
+    output_stream_.WriteAsync(buffer).get();
+  } catch (...) {
     return {Exception::kIo};
   }
 
@@ -149,7 +150,7 @@ Exception WifiLanSocket::SocketOutputStream::Write(const ByteArray& data) {
 Exception WifiLanSocket::SocketOutputStream::Flush() {
   try {
     output_stream_.FlushAsync().get();
-  } catch (std::exception exception) {
+  } catch (...) {
     return {Exception::kIo};
   }
 
@@ -159,7 +160,7 @@ Exception WifiLanSocket::SocketOutputStream::Flush() {
 Exception WifiLanSocket::SocketOutputStream::Close() {
   try {
     output_stream_.Close();
-  } catch (std::exception exception) {
+  } catch (...) {
     return {Exception::kIo};
   }