Add additional state management while in the PCP flow

PiperOrigin-RevId: 945918080

connections/implementation/endpoint_manager.cc

@@ -28,11 +28,13 @@
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel.h"
 #include "connections/implementation/endpoint_channel_manager.h"
+#include "connections/implementation/flags/nearby_connections_feature_flags.h"
 #include "connections/implementation/offline_frames.h"
 #include "connections/implementation/proto/offline_wire_formats.pb.h"
 #include "connections/implementation/service_id_constants.h"
 #include "connections/listeners.h"
 #include "connections/medium_selector.h"
+#include "internal/flags/nearby_flags.h"
 #include "internal/platform/byte_array.h"
 #include "internal/platform/count_down_latch.h"
 #include "internal/platform/exception.h"
@@ -64,6 +66,15 @@ constexpr absl::Time kInvalidTimestamp = absl::InfinitePast();
 // The maximum time we will wait for the encryption setup during negotiating a
 // connection.
 constexpr absl::Duration kDecryptRetryTimeout = absl::Seconds(3);
+
+// Returns true if the given `frame_type` is allowed before the connection to
+// the endpoint is confirmed (i.e., KEEP_ALIVE, CONNECTION_RESPONSE, and
+// DISCONNECTION frames).
+bool IsAllowedPreConfirmationFrameType(V1Frame::FrameType frame_type) {
+  return frame_type == V1Frame::KEEP_ALIVE ||
+         frame_type == V1Frame::CONNECTION_RESPONSE ||
+         frame_type == V1Frame::DISCONNECTION;
+}
 }  // namespace
 
 class EndpointManager::LockedFrameProcessor {
@@ -275,6 +286,17 @@ ExceptionOr<bool> EndpointManager::HandleData(
 
     // Route the incoming offlineFrame to its registered processor.
     V1Frame::FrameType frame_type = parser::GetFrameType(frame);
+    if (NearbyFlags::GetInstance().GetBoolFlag(
+            config_package_nearby::nearby_connections_feature::
+                kFilterUnconfirmedEndpointFrames) &&
+        client->HasPendingConnectionToEndpoint(endpoint_id) &&
+        !IsAllowedPreConfirmationFrameType(frame_type)) {
+      LOG(WARNING) << "EndpointManager discarded unauthorized frame ("
+                   << V1Frame::FrameType_Name(frame_type)
+                   << ") from unconfirmed endpoint " << endpoint_id << ".";
+      continue;
+    }
+
     LockedFrameProcessor frame_processor = GetFrameProcessor(frame_type);
     if (!frame_processor) {
       // report messages without handlers, except KEEP_ALIVE, which has

connections/implementation/endpoint_manager_test.cc

@@ -94,6 +94,10 @@ class SetSafeToDisconnect {
         config_package_nearby::nearby_connections_feature::
             kSafeToDisconnectVersion,
         safe_to_disconnect_version);
+    NearbyFlags::GetInstance().OverrideBoolFlagValue(
+        config_package_nearby::nearby_connections_feature::
+            kFilterUnconfirmedEndpointFrames,
+        false);
   }
 };
 
@@ -108,11 +112,11 @@ class EndpointManagerTest : public ::testing::Test {
  protected:
   void RegisterEndpoint(std::unique_ptr<MockEndpointChannel> channel,
                         bool should_close = true) {
-    CountDownLatch done(1);
+    auto done = std::make_shared<CountDownLatch>(1);
     if (should_close) {
       ON_CALL(*channel, Close(_))
           .WillByDefault(
-              [&done](DisconnectionReason reason) { done.CountDown(); });
+              [done](DisconnectionReason reason) { done->CountDown(); });
     }
     EXPECT_CALL(*channel, GetMedium()).WillRepeatedly(Return(Medium::BLE));
     EXPECT_CALL(*channel, GetLastReadTimestamp())
@@ -124,7 +128,7 @@ class EndpointManagerTest : public ::testing::Test {
                          connection_options_, std::move(channel), listener_,
                          connection_token_);
     if (should_close) {
-      EXPECT_TRUE(done.Await(absl::Milliseconds(1000)).result());
+      EXPECT_TRUE(done->Await(absl::Milliseconds(1000)).result());
     }
   }
   SetSafeToDisconnect set_safe_to_disconnect_{true, true, 5};
@@ -429,6 +433,70 @@ TEST_F(EndpointManagerTest, TryDecrypt) {
   RegisterEndpoint(std::move(endpoint_channel));
 }
 
+TEST_F(EndpointManagerTest,
+       FilterUnconfirmedFrames_DiscardsPayloadBeforeConfirmation) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kFilterUnconfirmedEndpointFrames,
+      true);
+  PayloadTransferFrame::PayloadHeader header;
+  header.set_id(12345);
+  header.set_type(PayloadTransferFrame::PayloadHeader::BYTES);
+  header.set_total_size(1024);
+  PayloadTransferFrame::PayloadChunk chunk;
+  chunk.set_body("payload data");
+  chunk.set_offset(150);
+  chunk.set_flags(1);
+  std::string payload_bytes = parser::ForDataPayloadTransfer(header, chunk);
+  auto endpoint_channel = std::make_unique<MockEndpointChannel>();
+  auto payload_processor = std::make_unique<MockFrameProcessor>();
+  EXPECT_CALL(*payload_processor, OnIncomingFrame).Times(0);
+  EXPECT_CALL(*payload_processor, OnEndpointDisconnect);
+  EXPECT_CALL(*endpoint_channel, Read())
+      .WillOnce(Return(ExceptionOr<ByteArray>(ByteArray(payload_bytes))))
+      .WillRepeatedly(Return(ExceptionOr<ByteArray>(Exception::kIo)));
+  EXPECT_CALL(*endpoint_channel, Write(_))
+      .WillRepeatedly(Return(Exception{Exception::kSuccess}));
+  em_.RegisterFrameProcessor(V1Frame::PAYLOAD_TRANSFER,
+                             payload_processor.get());
+  processors_.emplace_back(std::move(payload_processor));
+  RegisterEndpoint(std::move(endpoint_channel));
+}
+
+TEST_F(EndpointManagerTest,
+       FilterUnconfirmedFrames_AllowsPayloadAfterConfirmation) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kFilterUnconfirmedEndpointFrames,
+      true);
+  PayloadTransferFrame::PayloadHeader header;
+  header.set_id(12345);
+  header.set_type(PayloadTransferFrame::PayloadHeader::BYTES);
+  header.set_total_size(1024);
+  PayloadTransferFrame::PayloadChunk chunk;
+  chunk.set_body("payload data");
+  chunk.set_offset(150);
+  chunk.set_flags(1);
+  std::string payload_bytes = parser::ForDataPayloadTransfer(header, chunk);
+  auto endpoint_channel = std::make_unique<MockEndpointChannel>();
+  auto payload_processor = std::make_unique<MockFrameProcessor>();
+  EXPECT_CALL(mock_listener_.accepted_cb, Call).Times(1);
+  EXPECT_CALL(*payload_processor, OnIncomingFrame).Times(1);
+  EXPECT_CALL(*payload_processor, OnEndpointDisconnect);
+  EXPECT_CALL(*endpoint_channel, Read())
+      .WillOnce([this, payload_bytes]() {
+        client_->OnConnectionAccepted(endpoint_id_);
+        return ExceptionOr<ByteArray>(ByteArray(payload_bytes));
+      })
+      .WillRepeatedly(Return(ExceptionOr<ByteArray>(Exception::kIo)));
+  EXPECT_CALL(*endpoint_channel, Write(_))
+      .WillRepeatedly(Return(Exception{Exception::kSuccess}));
+  em_.RegisterFrameProcessor(V1Frame::PAYLOAD_TRANSFER,
+                             payload_processor.get());
+  processors_.emplace_back(std::move(payload_processor));
+  RegisterEndpoint(std::move(endpoint_channel));
+}
+
 // Regression test for b/278729669.
 //
 // During the destruction of NearbyConnections, Core (which owns ClientProxy)

connections/implementation/flags/nearby_connections_feature_flags.h

@@ -80,6 +80,11 @@ constexpr auto kEnableWifiDirect =
 // by default, enable Wi-Fi Hotspot client.
 constexpr auto kEnableWifiHotspotClient =
     flags::Flag<bool>(kConfigPackage, "45648734", true);
+// Enforces frame filtering on unconfirmed endpoints in EndpointManager and
+// BaseEndpointChannel so application payloads and upgrade requests are blocked
+// before connection acceptance.
+constexpr auto kFilterUnconfirmedEndpointFrames =
+    flags::Flag<bool>(kConfigPackage, "45813128", true);
 // When true, fix the BleServerSocket deadlock/use-after-free (b/494335036).
 constexpr auto kFixBleServerSocketDeadlock =
     flags::Flag<bool>(kConfigPackage, "45782647", true);