[Nearby Connections][C++] Track bandwidth upgrade mediums for each endpoint (2/4)

Consolidate BWU service ID bookkeeping into the BaseBwuHandler class. This eliminates duplicated code in per-medium BWU handlers and allows us to easily augment the bookkeeping with endpoint IDs in a subsequent CL.

Also, localize all service ID wrapping to the BaseBwuHandler class. A postfix is appended to the service ID during bandwidth upgrade, e.g., "NearbyShare_INITIATOR_BWU"; this is done to distinguish a medium's use for advertising/discovery vs. its use for device-to-device connections.

This CL should theoretically be a no-op, but per-medium BWU handlers were inconsistent in their implementations.

PiperOrigin-RevId: 442019603

connections/implementation/BUILD

@@ -40,6 +40,7 @@ cmake(
 cc_library(
     name = "internal",
     srcs = [
+        "base_bwu_handler.cc",
         "base_endpoint_channel.cc",
         "base_pcp_handler.cc",
         "ble_advertisement.cc",

connections/implementation/base_bwu_handler.cc

@@ -0,0 +1,55 @@
+// Copyright 2022 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/base_bwu_handler.h"
+
+#include <string>
+#include <utility>
+
+#include "connections/implementation/service_id_constants.h"
+#include "internal/platform/logging.h"
+
+namespace location {
+namespace nearby {
+namespace connections {
+
+BaseBwuHandler::BaseBwuHandler(BwuNotifications bwu_notifications)
+    : bwu_notifications_(std::move(bwu_notifications)) {}
+
+ByteArray BaseBwuHandler::InitializeUpgradedMediumForEndpoint(
+    ClientProxy* client, const std::string& service_id,
+    const std::string& endpoint_id) {
+  std::string upgrade_service_id = WrapInitiatorUpgradeServiceId(service_id);
+
+  // Perform any medium-specific handling in the child class.
+  ByteArray upgrade_path_available_frame =
+      HandleInitializeUpgradedMediumForEndpoint(client, upgrade_service_id,
+                                                endpoint_id);
+  if (!upgrade_path_available_frame.Empty()) {
+    active_service_ids_.insert(upgrade_service_id);
+  }
+
+  return upgrade_path_available_frame;
+}
+
+void BaseBwuHandler::RevertInitiatorState() {
+  for (const auto& service_id : active_service_ids_) {
+    HandleRevertInitiatorStateForService(service_id);
+  }
+  active_service_ids_.clear();
+}
+
+}  // namespace connections
+}  // namespace nearby
+}  // namespace location

connections/implementation/base_bwu_handler.h

@@ -15,40 +15,51 @@
 #ifndef CORE_INTERNAL_BASE_BWU_HANDLER_H_
 #define CORE_INTERNAL_BASE_BWU_HANDLER_H_
 
-#include <cstdint>
 #include <memory>
 #include <string>
 
 #include "absl/container/flat_hash_map.h"
-#include "absl/time/clock.h"
 #include "connections/implementation/bwu_handler.h"
-#include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel_manager.h"
-#include "internal/platform/cancelable_alarm.h"
-#include "internal/platform/count_down_latch.h"
-#include "internal/platform/scheduled_executor.h"
-#include "internal/platform/single_thread_executor.h"
 
 namespace location {
 namespace nearby {
 namespace connections {
 
+// Manages the bookkeeping common to all medium handlers. Notably, it tracks all
+// of the service IDs (TODO: and endpoint IDs) that initiated a bandwidth
+// upgrade.
 class BaseBwuHandler : public BwuHandler {
  public:
-  using ClientIntroduction = BwuNegotiationFrame::ClientIntroduction;
+  explicit BaseBwuHandler(BwuNotifications bwu_notifications);
 
-  BaseBwuHandler(EndpointChannelManager& channel_manager,
-                 BwuNotifications bwu_notifications)
-      : channel_manager_(&channel_manager),
-        bwu_notifications_(std::move(bwu_notifications)) {}
-  ~BaseBwuHandler() override = default;
+  // BwuHandler implementation:
+  ByteArray InitializeUpgradedMediumForEndpoint(
+      ClientProxy* client, const std::string& service_id,
+      const std::string& endpoint_id) final;
+  void RevertInitiatorState() final;
 
  protected:
+  // Invoked by InitializeUpgradedMediumForEndpoint and RevertInitiatorState,
+  // respectively, to handle medium-specific logic.
+  virtual ByteArray HandleInitializeUpgradedMediumForEndpoint(
+      ClientProxy* client, const std::string& upgrade_service_id,
+      const std::string& endpoint_id) = 0;
+  virtual void HandleRevertInitiatorStateForService(
+      const std::string& upgrade_service_id) = 0;
+
   // Represents the incoming Socket the Initiator has gotten after initializing
   // its upgraded bandwidth medium.
   EndpointChannelManager* GetEndpointChannelManager();
-  EndpointChannelManager* channel_manager_;
+
   BwuNotifications bwu_notifications_;
+
+ private:
+  // The set of service IDs that are initiating a bandwidth upgrade. Not used
+  // for service IDs that respond to bandwidth upgrade requests from another
+  // device; only tracked by the initiator.
+  // TODO(nohle): track endpoint IDs as well.
+  absl::flat_hash_set<std::string> active_service_ids_;
 };
 
 }  // namespace connections

connections/implementation/bluetooth_bwu_handler.cc

@@ -14,6 +14,9 @@
 
 #include "connections/implementation/bluetooth_bwu_handler.h"
 
+#include <string>
+#include <utility>
+
 #include "absl/functional/bind_front.h"
 #include "connections/implementation/bluetooth_endpoint_channel.h"
 #include "connections/implementation/client_proxy.h"
@@ -26,48 +29,82 @@ namespace location {
 namespace nearby {
 namespace connections {
 
-BluetoothBwuHandler::BluetoothBwuHandler(
-    Mediums& mediums, EndpointChannelManager& channel_manager,
-    BwuNotifications notifications)
-    : BaseBwuHandler(channel_manager, std::move(notifications)),
-      mediums_(mediums) {}
+BluetoothBwuHandler::BluetoothBwuHandler(Mediums& mediums,
+                                         BwuNotifications notifications)
+    : BaseBwuHandler(std::move(notifications)), mediums_(mediums) {}
 
-void BluetoothBwuHandler::Revert() {
-  for (const std::string& service_id : active_service_ids_) {
-    bluetooth_medium_.StopAcceptingConnections(service_id);
+// Called by BWU target. Retrieves a new medium info from incoming message,
+// and establishes connection over BT using this info.
+// Returns a channel ready to exchange data or nullptr on error.
+std::unique_ptr<EndpointChannel>
+BluetoothBwuHandler::CreateUpgradedEndpointChannel(
+    ClientProxy* client, const std::string& service_id,
+    const std::string& endpoint_id, const UpgradePathInfo& upgrade_path_info) {
+  const UpgradePathInfo::BluetoothCredentials& bluetooth_credentials =
+      upgrade_path_info.bluetooth_credentials();
+  if (!bluetooth_credentials.has_service_name() ||
+      !bluetooth_credentials.has_mac_address()) {
+    NEARBY_LOG(ERROR, "BluetoothBwuHandler failed to parse UpgradePathInfo.");
+    return nullptr;
   }
-  active_service_ids_.clear();
-  NEARBY_LOG(INFO,
-             "BluetoothBwuHandler successfully reverted all Bluetooth state.");
-}
 
-// Accept Connection Callback.
-// Notifies that the remote party called BluetoothClassic::Connect()
-// for this socket.
-void BluetoothBwuHandler::OnIncomingBluetoothConnection(
-    ClientProxy* client, const std::string& service_id,
-    BluetoothSocket socket) {
-  auto channel = absl::make_unique<BluetoothEndpointChannel>(
+  const std::string& service_name = bluetooth_credentials.service_name();
+  const std::string& mac_address = bluetooth_credentials.mac_address();
+
+  NEARBY_LOGS(VERBOSE) << "BluetoothBwuHandler is attempting to connect to "
+                          "available Bluetooth device ("
+                       << service_name << ", " << mac_address
+                       << ") for endpoint " << endpoint_id << " and service ID "
+                       << service_id;
+
+  BluetoothDevice device = bluetooth_medium_.GetRemoteDevice(mac_address);
+  if (!device.IsValid()) {
+    NEARBY_LOGS(ERROR)
+        << "BluetoothBwuHandler failed to derive a valid Bluetooth device "
+           "from the MAC address ("
+        << mac_address << ") for endpoint " << endpoint_id;
+    return nullptr;
+  }
+
+  BluetoothSocket socket = bluetooth_medium_.Connect(
+      device, service_id, client->GetCancellationFlag(endpoint_id));
+  if (!socket.IsValid()) {
+    NEARBY_LOGS(ERROR)
+        << "BluetoothBwuHandler failed to connect to the Bluetooth device ("
+        << service_name << ", " << mac_address << ") for endpoint "
+        << endpoint_id << " and service ID " << service_id;
+    return nullptr;
+  }
+
+  NEARBY_LOGS(VERBOSE)
+      << "BluetoothBwuHandler successfully connected to Bluetooth device ("
+      << service_id << ", " << mac_address << ") while upgrading endpoint "
+      << endpoint_id;
+
+  auto channel = std::make_unique<BluetoothEndpointChannel>(
       service_id, /*channel_name=*/service_id, socket);
-  std::unique_ptr<IncomingSocketConnection> connection{
-      new IncomingSocketConnection{
-          .socket =
-              std::make_unique<BluetoothIncomingSocket>(service_id, socket),
-          .channel = std::move(channel),
-      }};
-  bwu_notifications_.incoming_connection_cb(client, std::move(connection));
+  if (channel == nullptr) {
+    NEARBY_LOGS(ERROR)
+        << "BluetoothBwuHandler failed to create Bluetooth endpoint "
+           "channel to the Bluetooth device ("
+        << service_name << ", " << mac_address << ") for endpoint "
+        << endpoint_id << " and service ID " << service_id;
+    socket.Close();
+    return nullptr;
+  }
+
+  return channel;
 }
 
-// Called by BWU initiator. BT Medium is set up, and BWU request is prepared,
-// with necessary info (service_id, MAC address) for remote party to perform
-// discovery.
-ByteArray BluetoothBwuHandler::InitializeUpgradedMediumForEndpoint(
-    ClientProxy* client, const std::string& service_id,
+ByteArray BluetoothBwuHandler::HandleInitializeUpgradedMediumForEndpoint(
+    ClientProxy* client, const std::string& upgrade_service_id,
     const std::string& endpoint_id) {
-  std::string upgrade_service_id = Utils::WrapUpgradeServiceId(service_id);
-
   std::string mac_address = bluetooth_medium_.GetMacAddress();
   if (mac_address.empty()) {
+    NEARBY_LOGS(ERROR) << "BluetoothBwuHandler couldn't initiate the "
+                          "BLUETOOTH upgrade for service ID "
+                       << upgrade_service_id << " and endpoint " << endpoint_id
+                       << " because MAC address is empty.";
     return {};
   }
 
@@ -89,75 +126,35 @@ ByteArray BluetoothBwuHandler::InitializeUpgradedMediumForEndpoint(
     }
     NEARBY_LOGS(VERBOSE)
         << "BluetoothBwuHandler successfully started listening for incoming "
-           "Bluetooth connections on serviceid="
+           "Bluetooth connections on service_id="
         << upgrade_service_id << " while upgrading endpoint " << endpoint_id;
   }
-  // cache service ID to revert
-  active_service_ids_.emplace(upgrade_service_id);
 
   return parser::ForBwuBluetoothPathAvailable(upgrade_service_id, mac_address);
 }
 
-// Called by BWU target. Retrieves a new medium info from incoming message,
-// and establishes connection over BT using this info.
-// Returns a channel ready to exchange data or nullptr on error.
-std::unique_ptr<EndpointChannel>
-BluetoothBwuHandler::CreateUpgradedEndpointChannel(
-    ClientProxy* client, const std::string& service_id,
-    const std::string& endpoint_id, const UpgradePathInfo& upgrade_path_info) {
-  const UpgradePathInfo::BluetoothCredentials& bluetooth_credentials =
-      upgrade_path_info.bluetooth_credentials();
-  if (!bluetooth_credentials.has_service_name() ||
-      !bluetooth_credentials.has_mac_address()) {
-    NEARBY_LOG(ERROR, "BluetoothBwuHandler failed to parse UpgradePathInfo.");
-    return nullptr;
-  }
+void BluetoothBwuHandler::HandleRevertInitiatorStateForService(
+    const std::string& upgrade_service_id) {
+  bluetooth_medium_.StopAcceptingConnections(upgrade_service_id);
+  NEARBY_LOG(INFO,
+             "BluetoothBwuHandler successfully reverted all Bluetooth state.");
+}
 
-  const std::string& service_name = bluetooth_credentials.service_name();
-  const std::string& mac_address = bluetooth_credentials.mac_address();
-
-  NEARBY_LOGS(VERBOSE) << "BluetoothBwuHandler is attempting to connect to "
-                          "available Bluetooth device "
-                       << service_name << ", " << mac_address
-                       << ") for endpoint " << endpoint_id;
-
-  BluetoothDevice device = bluetooth_medium_.GetRemoteDevice(mac_address);
-  if (!device.IsValid()) {
-    NEARBY_LOGS(ERROR)
-        << "BluetoothBwuHandler failed to derive a valid Bluetooth device "
-           "from the MAC address ("
-        << mac_address << ") for endpoint " << endpoint_id;
-    return nullptr;
-  }
-
-  BluetoothSocket socket = bluetooth_medium_.Connect(
-      device, service_name, client->GetCancellationFlag(endpoint_id));
-  if (!socket.IsValid()) {
-    NEARBY_LOGS(ERROR)
-        << "BluetoothBwuHandler failed to connect to the Bluetooth device ("
-        << service_name << ", " << mac_address << ") for endpoint "
-        << endpoint_id;
-    return nullptr;
-  }
-
-  NEARBY_LOGS(VERBOSE)
-      << "BluetoothBwuHandler successfully connected to Bluetooth device ("
-      << service_name << ", " << mac_address << ") while upgrading endpoint "
-      << endpoint_id;
-
-  auto channel = std::make_unique<BluetoothEndpointChannel>(
-      service_id, /*channel_name=*/service_name, socket);
-  if (channel == nullptr) {
-    NEARBY_LOGS(ERROR)
-        << "BluetoothBwuHandler failed to create Bluetooth endpoint "
-           "channel to the Bluetooth device ("
-        << service_name << ", " << mac_address << ") for endpoint "
-        << endpoint_id;
-    socket.Close();
-    return nullptr;
-  }
-
-  return channel;
+// Accept Connection Callback.
+// Notifies that the remote party called BluetoothClassic::Connect()
+// for this socket.
+void BluetoothBwuHandler::OnIncomingBluetoothConnection(
+    ClientProxy* client, const std::string& upgrade_service_id,
+    BluetoothSocket socket) {
+  auto channel = absl::make_unique<BluetoothEndpointChannel>(
+      upgrade_service_id, /*channel_name=*/upgrade_service_id, socket);
+  std::unique_ptr<IncomingSocketConnection> connection{
+      new IncomingSocketConnection{
+          .socket = std::make_unique<BluetoothIncomingSocket>(
+              upgrade_service_id, socket),
+          .channel = std::move(channel),
+      }};
+  bwu_notifications_.incoming_connection_cb(client, std::move(connection));
 }
 
 }  // namespace connections

connections/implementation/bluetooth_bwu_handler.h

@@ -32,32 +32,16 @@ namespace connections {
 // per-Medium-specific operations needed to upgrade an EndpointChannel.
 class BluetoothBwuHandler : public BaseBwuHandler {
  public:
-  BluetoothBwuHandler(Mediums& mediums, EndpointChannelManager& channel_manager,
-                      BwuNotifications notifications);
-  ~BluetoothBwuHandler() override = default;
+  explicit BluetoothBwuHandler(Mediums& mediums,
+                               BwuNotifications notifications);
 
  private:
-  constexpr static const int kServiceIdLength = 10;
-
-  // Implements BaseBwuHandler:
-  // Reverts any changes made to the device in the process of upgrading
-  // endpoints.
-  void Revert() override;
-
-  // Cleans up in-progress upgrades after endpoint disconnection.
-  void OnEndpointDisconnect(ClientProxy* client,
-                            const std::string& endpoint_id) override {}
-
-  void OnIncomingBluetoothConnection(ClientProxy* client,
-                                     const std::string& service_id,
-                                     BluetoothSocket socket);
-
   class BluetoothIncomingSocket : public IncomingSocket {
    public:
     explicit BluetoothIncomingSocket(const std::string& name,
                                      BluetoothSocket socket)
         : name_(name), socket_(socket) {}
-    ~BluetoothIncomingSocket() override = default;
+
     std::string ToString() override { return name_; }
     void Close() override { socket_.Close(); }
 
@@ -66,24 +50,27 @@ class BluetoothBwuHandler : public BaseBwuHandler {
     BluetoothSocket socket_;
   };
 
-  // First part of InitiateBwuForEndpoint implementation;
-  // returns a BWU request to remote party as byte array.
-  ByteArray InitializeUpgradedMediumForEndpoint(
-      ClientProxy* client, const std::string& service_id,
-      const std::string& endpoint_id) override;
-
-  // Invoked from OnBwuNegotiationFrame.
+  // BwuHandler implementation:
   std::unique_ptr<EndpointChannel> CreateUpgradedEndpointChannel(
       ClientProxy* client, const std::string& service_id,
       const std::string& endpoint_id,
-      const UpgradePathInfo& upgrade_path_info) override;
+      const UpgradePathInfo& upgrade_path_info) final;
+  Medium GetUpgradeMedium() const final { return Medium::BLUETOOTH; }
+  void OnEndpointDisconnect(ClientProxy* client,
+                            const std::string& endpoint_id) final {}
 
-  // Returns the upgrade medium of the BwuHandler.
-  // @BwuHandlerThread
-  Medium GetUpgradeMedium() const override { return Medium::BLUETOOTH; }
+  // BaseBwuHandler implementation:
+  ByteArray HandleInitializeUpgradedMediumForEndpoint(
+      ClientProxy* client, const std::string& upgrade_service_id,
+      const std::string& endpoint_id) final;
+  void HandleRevertInitiatorStateForService(
+      const std::string& upgrade_service_id) final;
+
+  void OnIncomingBluetoothConnection(ClientProxy* client,
+                                     const std::string& upgrade_service_id,
+                                     BluetoothSocket socket);
 
   Mediums& mediums_;
-  absl::flat_hash_set<std::string> active_service_ids_;
   BluetoothRadio& bluetooth_radio_{mediums_.GetBluetoothRadio()};
   BluetoothClassic& bluetooth_medium_{mediums_.GetBluetoothClassic()};
 };

connections/implementation/bwu_handler.h

@@ -45,7 +45,7 @@ class BwuHandler {
   // Called to revert any state changed by the Initiator to setup the upgraded
   // medium for an endpoint.
   // @BwuHandlerThread
-  virtual void Revert() = 0;
+  virtual void RevertInitiatorState() = 0;
 
   // Called by the Responder to setup the upgraded medium for this endpoint (if
   // that hasn't already been done) using the UpgradePathInfo sent by the

connections/implementation/bwu_manager.cc

@@ -94,19 +94,16 @@ void BwuManager::InitBwuHandlers() {
           absl::bind_front(&BwuManager::OnIncomingConnection, this),
   };
   if (config_.allow_upgrade_to.wifi_lan) {
-    handlers_.emplace(Medium::WIFI_LAN,
-                      std::make_unique<WifiLanBwuHandler>(
-                          *mediums_, *channel_manager_, notifications));
+    handlers_.emplace(Medium::WIFI_LAN, std::make_unique<WifiLanBwuHandler>(
+                                            *mediums_, notifications));
   }
   if (config_.allow_upgrade_to.web_rtc) {
-    handlers_.emplace(Medium::WEB_RTC,
-                      std::make_unique<WebrtcBwuHandler>(
-                          *mediums_, *channel_manager_, notifications));
+    handlers_.emplace(Medium::WEB_RTC, std::make_unique<WebrtcBwuHandler>(
+                                           *mediums_, notifications));
   }
   if (config_.allow_upgrade_to.bluetooth) {
-    handlers_.emplace(Medium::BLUETOOTH,
-                      std::make_unique<BluetoothBwuHandler>(
-                          *mediums_, *channel_manager_, notifications));
+    handlers_.emplace(Medium::BLUETOOTH, std::make_unique<BluetoothBwuHandler>(
+                                             *mediums_, notifications));
   }
 }
 
@@ -132,7 +129,7 @@ void BwuManager::Shutdown() {
   medium_ = Medium::UNKNOWN_MEDIUM;
   for (auto& item : handlers_) {
     BwuHandler& handler = *item.second;
-    handler.Revert();
+    handler.RevertInitiatorState();
   }
   handlers_.clear();
 
@@ -337,7 +334,7 @@ void BwuManager::Revert() {
   NEARBY_LOGS(INFO) << "Revert reseting medium "
                     << proto::connections::Medium_Name(medium_);
   if (handler_) {
-    handler_->Revert();
+    handler_->RevertInitiatorState();
     handler_ = nullptr;
   }
   medium_ = Medium::UNKNOWN_MEDIUM;

connections/implementation/mediums/utils.cc

@@ -24,10 +24,6 @@ namespace location {
 namespace nearby {
 namespace connections {
 
-namespace {
-constexpr absl::string_view kUpgradeServiceIdPostfix = "_UPGRADE";
-}
-
 ByteArray Utils::GenerateRandomBytes(size_t length) {
   Prng rng;
   std::string data;
@@ -58,22 +54,6 @@ ByteArray Utils::Sha256Hash(const std::string& source, size_t length) {
   return full_hash;
 }
 
-std::string Utils::WrapUpgradeServiceId(const std::string& service_id) {
-  if (service_id.empty()) {
-    return {};
-  }
-  return service_id + std::string(kUpgradeServiceIdPostfix);
-}
-
-std::string Utils::UnwrapUpgradeServiceId(
-    const std::string& upgrade_service_id) {
-  auto pos = upgrade_service_id.find(std::string(kUpgradeServiceIdPostfix));
-  if (pos != std::string::npos) {
-    return std::string(upgrade_service_id, 0, pos);
-  }
-  return upgrade_service_id;
-}
-
 LocationHint Utils::BuildLocationHint(const std::string& location) {
   LocationHint location_hint;
   location_hint.set_format(LocationStandard::UNKNOWN);

connections/implementation/mediums/utils.h

@@ -30,8 +30,6 @@ class Utils {
   static ByteArray GenerateRandomBytes(size_t length);
   static ByteArray Sha256Hash(const ByteArray& source, size_t length);
   static ByteArray Sha256Hash(const std::string& source, size_t length);
-  static std::string WrapUpgradeServiceId(const std::string& service_id);
-  static std::string UnwrapUpgradeServiceId(const std::string& service_id);
   static LocationHint BuildLocationHint(const std::string& location);
 };
 

connections/implementation/service_id_constants.h

@@ -15,6 +15,8 @@
 #ifndef CORE_INTERNAL_SERVICE_ID_CONSTANTS_H_
 #define CORE_INTERNAL_SERVICE_ID_CONSTANTS_H_
 
+#include <string>
+
 #include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
 
@@ -35,6 +37,17 @@ inline bool IsInitiatorUpgradeServiceId(absl::string_view service_id) {
          absl::EndsWith(service_id, kInitiatorUpgradeServiceIdPostfix);
 }
 
+// Appends the kInitiatorUpgradeServiceIdPostfix to |service_id| if necessary.
+inline std::string WrapInitiatorUpgradeServiceId(absl::string_view service_id) {
+  // If |service_id| is empty or already has the upgrade postfix, do nothing.
+  if (service_id.empty() || IsInitiatorUpgradeServiceId(service_id)) {
+    return std::string(service_id);
+  }
+
+  return std::string(service_id) +
+         std::string(kInitiatorUpgradeServiceIdPostfix);
+}
+
 }  // namespace connections
 }  // namespace nearby
 }  // namespace location

connections/implementation/webrtc_bwu_handler.cc

@@ -15,6 +15,7 @@
 #include "connections/implementation/webrtc_bwu_handler.h"
 
 #include <string>
+#include <utility>
 
 #include "absl/functional/bind_front.h"
 #include "connections/implementation/client_proxy.h"
@@ -27,82 +28,19 @@ namespace location {
 namespace nearby {
 namespace connections {
 
+WebrtcBwuHandler::WebrtcIncomingSocket::WebrtcIncomingSocket(
+    const std::string& name, mediums::WebRtcSocketWrapper socket)
+    : name_(name), socket_(socket) {}
+
+void WebrtcBwuHandler::WebrtcIncomingSocket::Close() { socket_.Close(); }
+
+std::string WebrtcBwuHandler::WebrtcIncomingSocket::ToString() { return name_; }
+
 WebrtcBwuHandler::WebrtcBwuHandler(Mediums& mediums,
-                                   EndpointChannelManager& channel_manager,
                                    BwuNotifications notifications)
-    : BaseBwuHandler(channel_manager, std::move(notifications)),
+    : BaseBwuHandler(std::move(notifications)),
       mediums_(mediums) {}
 
-void WebrtcBwuHandler::Revert() {
-  for (const auto& service_id : active_service_ids_) {
-    webrtc_.StopAcceptingConnections(service_id);
-  }
-  active_service_ids_.clear();
-
-  NEARBY_LOG(INFO, "WebrtcBwuHandler successfully reverted state.");
-}
-
-// Accept Connection Callback.
-// Notifies that the remote party called WebRtc::Connect()
-// for this socket.
-void WebrtcBwuHandler::OnIncomingWebrtcConnection(
-    ClientProxy* client, const std::string& upgrade_service_id,
-    mediums::WebRtcSocketWrapper socket) {
-  std::string service_id = Utils::UnwrapUpgradeServiceId(upgrade_service_id);
-  auto channel = std::make_unique<WebRtcEndpointChannel>(
-      upgrade_service_id, /*channel_name=*/service_id, socket);
-  auto webrtc_socket =
-      std::make_unique<WebrtcIncomingSocket>(service_id, socket);
-  std::unique_ptr<IncomingSocketConnection> connection(
-      new IncomingSocketConnection{std::move(webrtc_socket),
-                                   std::move(channel)});
-
-  bwu_notifications_.incoming_connection_cb(client, std::move(connection));
-}
-
-// Called by BWU initiator. Set up WebRTC upgraded medium for this endpoint,
-// and returns a upgrade path info (PeerId, LocationHint) for remote party to
-// perform discovery.
-ByteArray WebrtcBwuHandler::InitializeUpgradedMediumForEndpoint(
-    ClientProxy* client, const std::string& service_id,
-    const std::string& endpoint_id) {
-  // Use wrapped service ID to avoid have the same ID with the one for
-  // startAdvertising. Otherwise, the listening request would be ignored because
-  // the medium already start accepting the connection because the client not
-  // stop the advertising yet.
-  std::string upgrade_service_id = Utils::WrapUpgradeServiceId(service_id);
-
-  LocationHint location_hint =
-      Utils::BuildLocationHint(webrtc_.GetDefaultCountryCode());
-
-  mediums::WebrtcPeerId self_id{mediums::WebrtcPeerId::FromRandom()};
-  if (!webrtc_.IsAcceptingConnections(service_id)) {
-    if (!webrtc_.StartAcceptingConnections(
-            upgrade_service_id, self_id, location_hint,
-            {
-                .accepted_cb = absl::bind_front(
-                    &WebrtcBwuHandler::OnIncomingWebrtcConnection, this,
-                    client),
-            })) {
-      NEARBY_LOG(ERROR,
-                 "WebRtcBwuHandler couldn't initiate the WEB_RTC upgrade for "
-                 "endpoint %s because it failed to start listening for "
-                 "incoming WebRTC connections.",
-                 endpoint_id.c_str());
-      return {};
-    }
-    NEARBY_LOG(INFO,
-               "WebRtcBwuHandler successfully started listening for incoming "
-               "WebRTC connections while upgrading endpoint %s",
-               endpoint_id.c_str());
-  }
-
-  // cache service ID to revert
-  active_service_ids_.emplace(upgrade_service_id);
-
-  return parser::ForBwuWebrtcPathAvailable(self_id.GetId(), location_hint);
-}
-
 // Called by BWU target. Retrieves a new medium info from incoming message,
 // and establishes connection over WebRTC using this info.
 std::unique_ptr<EndpointChannel>
@@ -152,16 +90,64 @@ WebrtcBwuHandler::CreateUpgradedEndpointChannel(
   return channel;
 }
 
-void WebrtcBwuHandler::OnEndpointDisconnect(ClientProxy* client,
-                                            const std::string& endpoint_id) {}
+void WebrtcBwuHandler::HandleRevertInitiatorStateForService(
+    const std::string& upgrade_service_id) {
+  webrtc_.StopAcceptingConnections(upgrade_service_id);
+  NEARBY_LOGS(INFO)
+      << "WebrtcBwuHandler successfully reverted state for service "
+      << upgrade_service_id;
+}
 
-WebrtcBwuHandler::WebrtcIncomingSocket::WebrtcIncomingSocket(
-    const std::string& name, mediums::WebRtcSocketWrapper socket)
-    : name_(name), socket_(socket) {}
+// Called by BWU initiator. Set up WebRTC upgraded medium for this endpoint,
+// and returns a upgrade path info (PeerId, LocationHint) for remote party to
+// perform discovery.
+ByteArray WebrtcBwuHandler::HandleInitializeUpgradedMediumForEndpoint(
+    ClientProxy* client, const std::string& upgrade_service_id,
+    const std::string& endpoint_id) {
+  LocationHint location_hint =
+      Utils::BuildLocationHint(webrtc_.GetDefaultCountryCode());
 
-void WebrtcBwuHandler::WebrtcIncomingSocket::Close() { socket_.Close(); }
+  mediums::WebrtcPeerId self_id{mediums::WebrtcPeerId::FromRandom()};
+  if (!webrtc_.IsAcceptingConnections(upgrade_service_id)) {
+    if (!webrtc_.StartAcceptingConnections(
+            upgrade_service_id, self_id, location_hint,
+            {
+                .accepted_cb = absl::bind_front(
+                    &WebrtcBwuHandler::OnIncomingWebrtcConnection, this,
+                    client),
+            })) {
+      NEARBY_LOG(ERROR,
+                 "WebRtcBwuHandler couldn't initiate the WEB_RTC upgrade for "
+                 "endpoint %s because it failed to start listening for "
+                 "incoming WebRTC connections.",
+                 endpoint_id.c_str());
+      return {};
+    }
+    NEARBY_LOG(INFO,
+               "WebRtcBwuHandler successfully started listening for incoming "
+               "WebRTC connections while upgrading endpoint %s",
+               endpoint_id.c_str());
+  }
 
-std::string WebrtcBwuHandler::WebrtcIncomingSocket::ToString() { return name_; }
+  return parser::ForBwuWebrtcPathAvailable(self_id.GetId(), location_hint);
+}
+
+// Accept Connection Callback.
+// Notifies that the remote party called WebRtc::Connect()
+// for this socket.
+void WebrtcBwuHandler::OnIncomingWebrtcConnection(
+    ClientProxy* client, const std::string& upgrade_service_id,
+    mediums::WebRtcSocketWrapper socket) {
+  auto channel = std::make_unique<WebRtcEndpointChannel>(
+      upgrade_service_id, /*channel_name=*/upgrade_service_id, socket);
+  auto webrtc_socket =
+      std::make_unique<WebrtcIncomingSocket>(upgrade_service_id, socket);
+  std::unique_ptr<IncomingSocketConnection> connection(
+      new IncomingSocketConnection{std::move(webrtc_socket),
+                                   std::move(channel)});
+
+  bwu_notifications_.incoming_connection_cb(client, std::move(connection));
+}
 
 }  // namespace connections
 }  // namespace nearby

connections/implementation/webrtc_bwu_handler.h

@@ -15,6 +15,8 @@
 #ifndef CORE_INTERNAL_WEBRTC_BWU_HANDLER_H_
 #define CORE_INTERNAL_WEBRTC_BWU_HANDLER_H_
 
+#include <string>
+
 #include "connections/implementation/base_bwu_handler.h"
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel_manager.h"
@@ -29,55 +31,17 @@ namespace location {
 namespace nearby {
 namespace connections {
 
-using BwuNegotiationFrame = BandwidthUpgradeNegotiationFrame;
-
 // Defines the set of methods that need to be implemented to handle the
 // per-Medium-specific operations needed to upgrade an EndpointChannel.
 class WebrtcBwuHandler : public BaseBwuHandler {
  public:
-  WebrtcBwuHandler(Mediums& mediums, EndpointChannelManager& channel_manager,
-                   BwuNotifications notifications);
-  ~WebrtcBwuHandler() override = default;
+  explicit WebrtcBwuHandler(Mediums& mediums, BwuNotifications notifications);
 
  private:
-  // Called by the Initiator to setup the upgraded medium for this endpoint (if
-  // that hasn't already been done), and returns a serialized UpgradePathInfo
-  // that can be sent to the Responder.
-  // @BwuHandlerThread
-  ByteArray InitializeUpgradedMediumForEndpoint(
-      ClientProxy* client, const std::string& service_id,
-      const std::string& endpoint_id) override;
-  // Called to revert any state changed by the Initiator to setup the upgraded
-  // medium for an endpoint.
-  // @BwuHandlerThread
-  void Revert() override;
-
-  // Called by the Responder to setup the upgraded medium for this endpoint (if
-  // that hasn't already been done) using the UpgradePathInfo sent by the
-  // Initiator, and returns a new EndpointChannel for the upgraded medium.
-  // @BwuHandlerThread
-  std::unique_ptr<EndpointChannel> CreateUpgradedEndpointChannel(
-      ClientProxy* client, const std::string& service_id,
-      const std::string& endpoint_id,
-      const UpgradePathInfo& upgrade_path_info) override;
-  // Returns the upgrade medium of the BwuHandler.
-  // @BwuHandlerThread
-  Medium GetUpgradeMedium() const override { return Medium::WEB_RTC; }
-
-  void OnIncomingWebrtcConnection(ClientProxy* client,
-                                  const std::string& service_id,
-                                  mediums::WebRtcSocketWrapper socket);
-
-  void OnEndpointDisconnect(ClientProxy* client,
-                            const std::string& endpoint_id) override;
-
-  std::string GetCountryCode();
-
   class WebrtcIncomingSocket : public BwuHandler::IncomingSocket {
    public:
     explicit WebrtcIncomingSocket(const std::string& name,
                                   mediums::WebRtcSocketWrapper socket);
-    ~WebrtcIncomingSocket() override = default;
 
     std::string ToString() override;
     void Close() override;
@@ -87,9 +51,28 @@ class WebrtcBwuHandler : public BaseBwuHandler {
     mediums::WebRtcSocketWrapper socket_;
   };
 
+  // BwuHandler implementation:
+  std::unique_ptr<EndpointChannel> CreateUpgradedEndpointChannel(
+      ClientProxy* client, const std::string& service_id,
+      const std::string& endpoint_id,
+      const UpgradePathInfo& upgrade_path_info) final;
+  Medium GetUpgradeMedium() const final { return Medium::WEB_RTC; }
+  void OnEndpointDisconnect(ClientProxy* client,
+                            const std::string& endpoint_id) final {}
+
+  // BaseBwuHandler implementation:
+  ByteArray HandleInitializeUpgradedMediumForEndpoint(
+      ClientProxy* client, const std::string& upgrade_service_id,
+      const std::string& endpoint_id) final;
+  void HandleRevertInitiatorStateForService(
+      const std::string& upgrade_service_id) final;
+
+  void OnIncomingWebrtcConnection(ClientProxy* client,
+                                  const std::string& upgrade_service_id,
+                                  mediums::WebRtcSocketWrapper socket);
+
   Mediums& mediums_;
   mediums::WebRtc& webrtc_{mediums_.GetWebRtc()};
-  absl::flat_hash_set<std::string> active_service_ids_;
 };
 
 }  // namespace connections

connections/implementation/wifi_lan_bwu_handler.cc

@@ -16,10 +16,10 @@
 
 #include <locale>
 #include <string>
+#include <utility>
 
 #include "absl/functional/bind_front.h"
 #include "connections/implementation/client_proxy.h"
-#include "connections/implementation/mediums/utils.h"
 #include "connections/implementation/offline_frames.h"
 #include "connections/implementation/wifi_lan_endpoint_channel.h"
 #include "internal/platform/wifi_lan.h"
@@ -29,76 +29,8 @@ namespace nearby {
 namespace connections {
 
 WifiLanBwuHandler::WifiLanBwuHandler(Mediums& mediums,
-                                     EndpointChannelManager& channel_manager,
                                      BwuNotifications notifications)
-    : BaseBwuHandler(channel_manager, std::move(notifications)),
-      mediums_(mediums) {}
-
-// Called by BWU initiator. Set up WifiLan upgraded medium for this endpoint,
-// and returns a upgrade path info (ip address, port) for remote party to
-// perform discovery.
-ByteArray WifiLanBwuHandler::InitializeUpgradedMediumForEndpoint(
-    ClientProxy* client, const std::string& service_id,
-    const std::string& endpoint_id) {
-  // Use wrapped service ID to avoid have the same ID with the one for
-  // startAdvertising. Otherwise, the listening request would be ignored because
-  // the medium already start accepting the connection because the client not
-  // stop the advertising yet.
-  std::string upgrade_service_id = Utils::WrapUpgradeServiceId(service_id);
-
-  if (!wifi_lan_medium_.IsAcceptingConnections(upgrade_service_id)) {
-    if (!wifi_lan_medium_.StartAcceptingConnections(
-            upgrade_service_id,
-            {
-                .accepted_cb = absl::bind_front(
-                    &WifiLanBwuHandler::OnIncomingWifiLanConnection, this,
-                    client),
-            })) {
-      NEARBY_LOGS(ERROR)
-          << "WifiLanBwuHandler couldn't initiate the WifiLan upgrade for "
-             "endpoint "
-          << endpoint_id
-          << " because it failed to start listening for "
-             "incoming WifiLan connections.";
-      return {};
-    }
-    NEARBY_LOGS(INFO)
-        << "WifiLanBwuHandler successfully started listening for incoming "
-           "WifiLan connections while upgrading endpoint "
-        << endpoint_id;
-  }
-
-  // cache service ID to revert
-  active_service_ids_.insert(upgrade_service_id);
-
-  auto credential = wifi_lan_medium_.GetCredentials(upgrade_service_id);
-  auto ip_address = credential.first;
-  auto port = credential.second;
-
-  if (ip_address.empty()) {
-    NEARBY_LOGS(INFO)
-        << "WifiLanBwuHandler couldn't initiate the wifi_lan upgrade for "
-           "endpoint "
-        << endpoint_id
-        << " because the wifi_lan ip address were unable to be obtained.";
-    return {};
-  }
-  NEARBY_LOGS(INFO)
-      << "WifiLanBwuHandler retrieved WIFI_LAN credential. IP addr: "
-      << ip_address[0] << "." << ip_address[1] << "." << ip_address[2] << "."
-      << ip_address[3] << ",  Port: " << port;
-
-  return parser::ForBwuWifiLanPathAvailable(ip_address, port);
-}
-
-void WifiLanBwuHandler::Revert() {
-  for (const std::string& service_id : active_service_ids_) {
-    wifi_lan_medium_.StopAcceptingConnections(service_id);
-  }
-  active_service_ids_.clear();
-
-  NEARBY_LOG(INFO, "WifiLanBwuHandler successfully reverted all states.");
-}
+    : BaseBwuHandler(std::move(notifications)), mediums_(mediums) {}
 
 // Called by BWU target. Retrieves a new medium info from incoming message,
 // and establishes connection over WifiLan using this info.
@@ -121,9 +53,8 @@ WifiLanBwuHandler::CreateUpgradedEndpointChannel(
   std::int32_t port = upgrade_path_info_socket.wifi_port();
 
   NEARBY_LOGS(VERBOSE) << "WifiLanBwuHandler is attempting to connect to "
-                          "available WifiLan service ("
-                       << ip_address << ":" << port << ") for endpoint "
-                       << endpoint_id;
+                       << "available WifiLan service (" << ip_address << ":"
+                       << port << ") for endpoint " << endpoint_id;
 
   WifiLanSocket socket = wifi_lan_medium_.Connect(
       service_id, ip_address, port, client->GetCancellationFlag(endpoint_id));
@@ -144,9 +75,8 @@ WifiLanBwuHandler::CreateUpgradedEndpointChannel(
       service_id, /*channel_name=*/service_id, socket);
   if (channel == nullptr) {
     NEARBY_LOGS(ERROR) << "WifiLanBwuHandler failed to create WifiLan endpoint "
-                          "channel to the WifiLan service ("
-                       << ip_address << ":" << port << ") for endpoint "
-                       << endpoint_id;
+                       << "channel to the WifiLan service (" << ip_address
+                       << ":" << port << ") for endpoint " << endpoint_id;
     socket.Close();
     return nullptr;
   }
@@ -154,15 +84,72 @@ WifiLanBwuHandler::CreateUpgradedEndpointChannel(
   return channel;
 }
 
+// Called by BWU initiator. Set up WifiLan upgraded medium for this endpoint,
+// and returns a upgrade path info (ip address, port) for remote party to
+// perform discovery.
+ByteArray WifiLanBwuHandler::HandleInitializeUpgradedMediumForEndpoint(
+    ClientProxy* client, const std::string& upgrade_service_id,
+    const std::string& endpoint_id) {
+  if (!wifi_lan_medium_.IsAcceptingConnections(upgrade_service_id)) {
+    if (!wifi_lan_medium_.StartAcceptingConnections(
+            upgrade_service_id,
+            {
+                .accepted_cb = absl::bind_front(
+                    &WifiLanBwuHandler::OnIncomingWifiLanConnection, this,
+                    client),
+            })) {
+      NEARBY_LOGS(ERROR)
+          << "WifiLanBwuHandler couldn't initiate the WifiLan upgrade for "
+          << "service " << upgrade_service_id << " and endpoint " << endpoint_id
+          << " because it failed to start listening for incoming WifiLan "
+             "connections.";
+      return {};
+    }
+    NEARBY_LOGS(INFO)
+        << "WifiLanBwuHandler successfully started listening for incoming "
+           "WifiLan connections while upgrading endpoint "
+        << endpoint_id;
+  }
+
+  // Note: Credentials are not populated until StartAcceptingConnections() is
+  // called and the server socket is created. Be careful moving this codeblock
+  // around.
+  auto credential = wifi_lan_medium_.GetCredentials(upgrade_service_id);
+  auto ip_address = credential.first;
+  auto port = credential.second;
+  if (ip_address.empty()) {
+    NEARBY_LOGS(INFO)
+        << "WifiLanBwuHandler couldn't initiate the wifi_lan upgrade for "
+        << "service " << upgrade_service_id << " and endpoint " << endpoint_id
+        << " because the wifi_lan ip address were unable to be obtained.";
+    return {};
+  }
+
+  NEARBY_LOGS(INFO)
+      << "WifiLanBwuHandler retrieved WIFI_LAN credentials. IP addr: "
+      << ip_address[0] << "." << ip_address[1] << "." << ip_address[2] << "."
+      << ip_address[3] << ",  Port: " << port;
+
+  return parser::ForBwuWifiLanPathAvailable(ip_address, port);
+}
+
+void WifiLanBwuHandler::HandleRevertInitiatorStateForService(
+    const std::string& upgrade_service_id) {
+  wifi_lan_medium_.StopAcceptingConnections(upgrade_service_id);
+  NEARBY_LOGS(INFO) << "WifiLanBwuHandler successfully reverted all states for "
+                    << "upgrade service ID " << upgrade_service_id;
+}
+
 // Accept Connection Callback.
 void WifiLanBwuHandler::OnIncomingWifiLanConnection(
-    ClientProxy* client, const std::string& service_id, WifiLanSocket socket) {
+    ClientProxy* client, const std::string& upgrade_service_id,
+    WifiLanSocket socket) {
   auto channel = absl::make_unique<WifiLanEndpointChannel>(
-      service_id, /*channel_name=*/service_id, socket);
+      upgrade_service_id, /*channel_name=*/upgrade_service_id, socket);
   std::unique_ptr<IncomingSocketConnection> connection(
       new IncomingSocketConnection{
-          .socket =
-              absl::make_unique<WifiLanIncomingSocket>(service_id, socket),
+          .socket = absl::make_unique<WifiLanIncomingSocket>(upgrade_service_id,
+                                                             socket),
           .channel = std::move(channel),
       });
   bwu_notifications_.incoming_connection_cb(client, std::move(connection));

connections/implementation/wifi_lan_bwu_handler.h

@@ -15,6 +15,8 @@
 #ifndef CORE_INTERNAL_WIFI_LAN_BWU_HANDLER_H_
 #define CORE_INTERNAL_WIFI_LAN_BWU_HANDLER_H_
 
+#include <string>
+
 #include "connections/implementation/base_bwu_handler.h"
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel_manager.h"
@@ -28,37 +30,14 @@ namespace connections {
 // per-Medium-specific operations needed to upgrade an EndpointChannel.
 class WifiLanBwuHandler : public BaseBwuHandler {
  public:
-  WifiLanBwuHandler(Mediums& mediums, EndpointChannelManager& channel_manager,
-                    BwuNotifications notifications);
-  ~WifiLanBwuHandler() override = default;
+  explicit WifiLanBwuHandler(Mediums& mediums, BwuNotifications notifications);
 
  private:
-  ByteArray InitializeUpgradedMediumForEndpoint(
-      ClientProxy* client, const std::string& service_id,
-      const std::string& endpoint_id) override;
-
-  void Revert() override;
-
-  std::unique_ptr<EndpointChannel> CreateUpgradedEndpointChannel(
-      ClientProxy* client, const std::string& service_id,
-      const std::string& endpoint_id,
-      const UpgradePathInfo& upgrade_path_info) override;
-
-  Medium GetUpgradeMedium() const override { return Medium::WIFI_LAN; }
-
-  void OnEndpointDisconnect(ClientProxy* client,
-                            const std::string& endpoint_id) override {}
-
-  void OnIncomingWifiLanConnection(ClientProxy* client,
-                                   const std::string& service_id,
-                                   WifiLanSocket socket);
-
   class WifiLanIncomingSocket : public BwuHandler::IncomingSocket {
    public:
     explicit WifiLanIncomingSocket(const std::string& name,
                                    WifiLanSocket socket)
         : name_(name), socket_(socket) {}
-    ~WifiLanIncomingSocket() override = default;
 
     std::string ToString() override { return name_; }
     void Close() override { socket_.Close(); }
@@ -68,9 +47,28 @@ class WifiLanBwuHandler : public BaseBwuHandler {
     WifiLanSocket socket_;
   };
 
+  // BwuHandler implementation:
+  std::unique_ptr<EndpointChannel> CreateUpgradedEndpointChannel(
+      ClientProxy* client, const std::string& service_id,
+      const std::string& endpoint_id,
+      const UpgradePathInfo& upgrade_path_info) final;
+  Medium GetUpgradeMedium() const final { return Medium::WIFI_LAN; }
+  void OnEndpointDisconnect(ClientProxy* client,
+                            const std::string& endpoint_id) final {}
+
+  // BaseBwuHandler implementation:
+  ByteArray HandleInitializeUpgradedMediumForEndpoint(
+      ClientProxy* client, const std::string& upgrade_service_id,
+      const std::string& endpoint_id) final;
+  void HandleRevertInitiatorStateForService(
+      const std::string& upgrade_service_id) final;
+
+  void OnIncomingWifiLanConnection(ClientProxy* client,
+                                   const std::string& upgrade_service_id,
+                                   WifiLanSocket socket);
+
   Mediums& mediums_;
   WifiLan& wifi_lan_medium_{mediums_.GetWifiLan()};
-  absl::flat_hash_set<std::string> active_service_ids_;
 };
 
 }  // namespace connections