Store IncomingFramesReader in IncomingShareTargetInfo

- NearbyConnection is owned by NearbyConnectionsManager, and is
  destroyed when it's disconnected by either end from various methods
- NearbyConnection offers a disconnection listener to listen for
  destruction
- NearbyServiceManagerImpl expects to overwrite NearbyConnection's
  disconnection listener time to time, but NearbyConnectionImpl
  currently allows multiple listeners, causing
  NearbyServiceManagerImpl's disconnection handler to be called multiple
  times
- Changed NearbyConnectionImpl's disconnection listener back to a single
  member that can be overwritten
- Made IncomingFramesReader owned by IncomingShareTargetInfo so it'll be
  cleaned up in UnregisterShareTarget

Bug: 1076008
Change-Id: I16effec3526e0e55a958a96bc3893f5d619e5a1f
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2339982Reviewed-by: Richard Knoll <knollr@chromium.org>
Reviewed-by: Himanshu Jaju <himanshujaju@chromium.org>
Commit-Queue: Alex Chau <alexchau@chromium.org>
Cr-Commit-Position: refs/heads/master@{#795970}

chrome/browser/nearby_sharing/fake_nearby_connection.cc

@@ -24,16 +24,17 @@ void FakeNearbyConnection::Write(std::vector<uint8_t> bytes) {
 void FakeNearbyConnection::Close() {
   DCHECK(!closed_);
   closed_ = true;
+  if (disconnect_listener_)
+    std::move(disconnect_listener_).Run();
+
   if (callback_)
     std::move(callback_).Run(base::nullopt);
-  for (auto& disconnect_listener : disconnect_listeners_)
-    std::move(disconnect_listener).Run();
 }
 
-void FakeNearbyConnection::RegisterForDisconnection(
+void FakeNearbyConnection::SetDisconnectionListener(
     base::OnceClosure listener) {
   DCHECK(!closed_);
-  disconnect_listeners_.push_back(std::move(listener));
+  disconnect_listener_ = std::move(listener);
 }
 
 void FakeNearbyConnection::AppendReadableData(std::vector<uint8_t> bytes) {

chrome/browser/nearby_sharing/fake_nearby_connection.h

@@ -19,7 +19,7 @@ class FakeNearbyConnection : public NearbyConnection {
   void Read(ReadCallback callback) override;
   void Write(std::vector<uint8_t> bytes) override;
   void Close() override;
-  void RegisterForDisconnection(base::OnceClosure listener) override;
+  void SetDisconnectionListener(base::OnceClosure listener) override;
 
   void AppendReadableData(std::vector<uint8_t> bytes);
   std::vector<uint8_t> GetWrittenData();
@@ -31,7 +31,7 @@ class FakeNearbyConnection : public NearbyConnection {
   ReadCallback callback_;
   std::queue<std::vector<uint8_t>> read_data_;
   std::queue<std::vector<uint8_t>> write_data_;
-  std::vector<base::OnceClosure> disconnect_listeners_;
+  base::OnceClosure disconnect_listener_;
 };
 
 #endif  // CHROME_BROWSER_NEARBY_SHARING_FAKE_NEARBY_CONNECTION_H_

chrome/browser/nearby_sharing/incoming_frames_reader.cc

@@ -36,9 +36,6 @@ IncomingFramesReader::IncomingFramesReader(
   DCHECK(connection_);
 
   nearby_process_observer_.Add(process_manager);
-  connection->RegisterForDisconnection(
-      base::BindOnce(&IncomingFramesReader::OnConnectionClosed,
-                     weak_ptr_factory_.GetWeakPtr()));
 }
 
 IncomingFramesReader::~IncomingFramesReader() = default;
@@ -180,10 +177,6 @@ void IncomingFramesReader::Done(
   }
 }
 
-void IncomingFramesReader::OnConnectionClosed() {
-  connection_ = nullptr;
-}
-
 base::Optional<sharing::mojom::V1FramePtr> IncomingFramesReader::GetCachedFrame(
     base::Optional<sharing::mojom::V1Frame::Tag> frame_type) {
   NS_LOG(VERBOSE) << __func__ << ": Fetching cached frame";

chrome/browser/nearby_sharing/incoming_frames_reader.h

@@ -61,7 +61,6 @@ class IncomingFramesReader : public NearbyProcessManager::Observer {
   void OnFrameDecoded(sharing::mojom::FramePtr mojo_frame);
   void OnTimeout();
   void Done(base::Optional<sharing::mojom::V1FramePtr> frame);
-  void OnConnectionClosed();
   base::Optional<sharing::mojom::V1FramePtr> GetCachedFrame(
       base::Optional<sharing::mojom::V1Frame::Tag> frame_type);
 

chrome/browser/nearby_sharing/incoming_frames_reader_unittest.cc

@@ -283,15 +283,4 @@ TEST_F(IncomingFramesReaderTest, ReadAfterConnectionClosed) {
 
   connection().Close();
   run_loop_before_close.Run();
-
-  base::RunLoop run_loop_after_close;
-  frames_reader().ReadFrame(
-      sharing::mojom::V1Frame::Tag::INTRODUCTION,
-      base::BindLambdaForTesting(
-          [&](base::Optional<sharing::mojom::V1FramePtr> frame) {
-            EXPECT_FALSE(frame);
-            run_loop_after_close.Quit();
-          }),
-      kTimeout);
-  run_loop_after_close.Run();
 }

chrome/browser/nearby_sharing/incoming_share_target_info.h

@@ -5,10 +5,12 @@
 #ifndef CHROME_BROWSER_NEARBY_SHARING_INCOMING_SHARE_TARGET_INFO_H_
 #define CHROME_BROWSER_NEARBY_SHARING_INCOMING_SHARE_TARGET_INFO_H_
 
+#include <memory>
 #include <string>
 
 #include "base/optional.h"
 #include "chrome/browser/nearby_sharing/certificates/nearby_share_decrypted_public_certificate.h"
+#include "chrome/browser/nearby_sharing/incoming_frames_reader.h"
 
 class NearbyConnection;
 
@@ -50,11 +52,18 @@ class IncomingShareTargetInfo {
 
   const base::Optional<std::string>& token() const { return token_; }
 
+  IncomingFramesReader* frames_reader() { return frames_reader_.get(); }
+
+  void set_frames_reader(std::unique_ptr<IncomingFramesReader> frames_reader) {
+    frames_reader_ = std::move(frames_reader);
+  }
+
  private:
   base::Optional<std::string> endpoint_id_;
   base::Optional<NearbyShareDecryptedPublicCertificate> certificate_;
   NearbyConnection* connection_ = nullptr;
   base::Optional<std::string> token_;
+  std::unique_ptr<IncomingFramesReader> frames_reader_;
 };
 
 #endif  // CHROME_BROWSER_NEARBY_SHARING_INCOMING_SHARE_TARGET_INFO_H_

chrome/browser/nearby_sharing/nearby_connection.h

@@ -32,12 +32,13 @@ class NearbyConnection {
   virtual void Write(std::vector<uint8_t> bytes) = 0;
 
   // Closes the socket and disconnects from the remote device. This object will
-  // be invalidated after |callback| in RegisterForDisconnection is invoked.
+  // be invalidated after |callback| in SetDisconnectionListener is invoked.
   virtual void Close() = 0;
 
   // Listens to the socket being closed. Invoke |callback| when the socket is
   // closed. This object will be invalidated after |listener| is invoked.
-  virtual void RegisterForDisconnection(base::OnceClosure listener) = 0;
+  // Previously set listener will be replaced by |listener|.
+  virtual void SetDisconnectionListener(base::OnceClosure listener) = 0;
 };
 
 #endif  // CHROME_BROWSER_NEARBY_SHARING_NEARBY_CONNECTION_H_

chrome/browser/nearby_sharing/nearby_connection_impl.cc

@@ -15,11 +15,11 @@ NearbyConnectionImpl::NearbyConnectionImpl(
 
 NearbyConnectionImpl::~NearbyConnectionImpl() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  if (disconnect_listener_)
+    std::move(disconnect_listener_).Run();
+
   if (read_callback_)
     std::move(read_callback_).Run(base::nullopt);
-
-  for (auto& disconnect_listener : disconnect_listeners_)
-    std::move(disconnect_listener).Run();
 }
 
 void NearbyConnectionImpl::Read(ReadCallback callback) {
@@ -53,10 +53,10 @@ void NearbyConnectionImpl::Close() {
   nearby_connections_manager_->Disconnect(std::string(endpoint_id_));
 }
 
-void NearbyConnectionImpl::RegisterForDisconnection(
+void NearbyConnectionImpl::SetDisconnectionListener(
     base::OnceClosure listener) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  disconnect_listeners_.push_back(std::move(listener));
+  disconnect_listener_ = std::move(listener);
 }
 
 void NearbyConnectionImpl::WriteMessage(std::vector<uint8_t> bytes) {

chrome/browser/nearby_sharing/nearby_connection_impl.h

@@ -24,7 +24,7 @@ class NearbyConnectionImpl : public NearbyConnection {
   void Read(ReadCallback callback) override;
   void Write(std::vector<uint8_t> bytes) override;
   void Close() override;
-  void RegisterForDisconnection(base::OnceClosure listener) override;
+  void SetDisconnectionListener(base::OnceClosure listener) override;
 
   // Add |bytes| to the read queue, notifying ReadCallback.
   void WriteMessage(std::vector<uint8_t> bytes);
@@ -38,7 +38,7 @@ class NearbyConnectionImpl : public NearbyConnection {
   NearbyConnectionsManager* nearby_connections_manager_;
   std::string endpoint_id_;
   ReadCallback read_callback_;
-  std::vector<base::OnceClosure> disconnect_listeners_;
+  base::OnceClosure disconnect_listener_;
 
   // A read queue. The data that we've read from the remote device ends up here
   // until Read() is called to dequeue it.

chrome/browser/nearby_sharing/nearby_connections_manager_impl_unittest.cc

@@ -389,7 +389,7 @@ TEST_F(NearbyConnectionsManagerImplTest, ConnectClosed) {
 
   // Close should invoke disconnection callback and read callback.
   base::RunLoop close_run_loop;
-  nearby_connection->RegisterForDisconnection(
+  nearby_connection->SetDisconnectionListener(
       base::BindLambdaForTesting([&]() { close_run_loop.Quit(); }));
   base::RunLoop read_run_loop_3;
   nearby_connection->Read(base::BindLambdaForTesting(
@@ -427,7 +427,7 @@ TEST_F(NearbyConnectionsManagerImplTest, ConnectClosedByRemote) {
 
   // Remote closing should invoke disconnection callback and read callback.
   base::RunLoop close_run_loop;
-  nearby_connection->RegisterForDisconnection(
+  nearby_connection->SetDisconnectionListener(
       base::BindLambdaForTesting([&]() { close_run_loop.Quit(); }));
   base::RunLoop read_run_loop;
   nearby_connection->Read(base::BindLambdaForTesting(
@@ -458,7 +458,7 @@ TEST_F(NearbyConnectionsManagerImplTest, ConnectClosedByClient) {
 
   // Remote closing should invoke disconnection callback and read callback.
   base::RunLoop close_run_loop;
-  nearby_connection->RegisterForDisconnection(
+  nearby_connection->SetDisconnectionListener(
       base::BindLambdaForTesting([&]() { close_run_loop.Quit(); }));
   base::RunLoop read_run_loop;
   nearby_connection->Read(base::BindLambdaForTesting(

chrome/browser/nearby_sharing/nearby_sharing_service_impl.cc

@@ -325,7 +325,7 @@ void NearbySharingServiceImpl::Reject(
                      weak_ptr_factory_.GetWeakPtr(), share_target),
       kIncomingRejectionDelay);
 
-  connection->RegisterForDisconnection(
+  connection->SetDisconnectionListener(
       base::BindOnce(&NearbySharingServiceImpl::UnregisterShareTarget,
                      weak_ptr_factory_.GetWeakPtr(), share_target));
 
@@ -391,7 +391,7 @@ void NearbySharingServiceImpl::OnIncomingConnection(
 
   incoming_share_target_info_map_[share_target.id].set_connection(connection);
   incoming_share_target_info_map_[share_target.id].set_endpoint_id(endpoint_id);
-  connection->RegisterForDisconnection(
+  connection->SetDisconnectionListener(
       base::BindOnce(&NearbySharingServiceImpl::UnregisterShareTarget,
                      weak_ptr_factory_.GetWeakPtr(), share_target));
 
@@ -832,7 +832,7 @@ void NearbySharingServiceImpl::Fail(const ShareTarget& share_target,
     return;
   }
 
-  // TODO(himanshujaju) - Create alarm and cancel in RegisterForDisconnection().
+  // TODO(himanshujaju) - Create alarm and cancel in SetDisconnectionListener().
 
   // Send response to remote device.
   sharing::nearby::ConnectionResponseFrame::Status response_status;
@@ -889,21 +889,20 @@ void NearbySharingServiceImpl::ReceiveIntroduction(
     return;
   }
 
-  auto frames_reader = std::make_unique<IncomingFramesReader>(
-      process_manager_, profile_, connection);
-
-  frames_reader->ReadFrame(
+  auto& share_target_info = GetIncomingShareTargetInfo(share_target);
+  share_target_info.set_frames_reader(std::make_unique<IncomingFramesReader>(
+      process_manager_, profile_, connection));
+  share_target_info.frames_reader()->ReadFrame(
       sharing::mojom::V1Frame::Tag::INTRODUCTION,
       base::BindOnce(&NearbySharingServiceImpl::OnReceivedIntroduction,
                      weak_ptr_factory_.GetWeakPtr(), std::move(share_target),
-                     std::move(token), std::move(frames_reader)),
+                     std::move(token)),
       kReadFramesTimeout);
 }
 
 void NearbySharingServiceImpl::OnReceivedIntroduction(
     ShareTarget share_target,
     base::Optional<std::string> token,
-    std::unique_ptr<IncomingFramesReader> frames_reader,
     base::Optional<sharing::mojom::V1FramePtr> frame) {
   NearbyConnection* connection = GetIncomingConnection(share_target);
   if (!connection) {
@@ -1000,17 +999,25 @@ void NearbySharingServiceImpl::OnReceivedIntroduction(
     return;
   }
 
-  connection->RegisterForDisconnection(base::BindOnce(
+  connection->SetDisconnectionListener(base::BindOnce(
       &NearbySharingServiceImpl::OnIncomingConnectionDisconnected,
       weak_ptr_factory_.GetWeakPtr(), share_target));
 
-  frames_reader->ReadFrame(base::BindOnce(
-      &NearbySharingServiceImpl::OnFrameRead, weak_ptr_factory_.GetWeakPtr(),
-      std::move(frames_reader), std::move(share_target)));
+  auto* frames_reader =
+      GetIncomingShareTargetInfo(share_target).frames_reader();
+  if (!frames_reader) {
+    NS_LOG(WARNING) << __func__
+                    << ": Stopped reading further frames, due to no connection "
+                       "established.";
+    return;
+  }
+
+  frames_reader->ReadFrame(
+      base::BindOnce(&NearbySharingServiceImpl::OnFrameRead,
+                     weak_ptr_factory_.GetWeakPtr(), std::move(share_target)));
 }
 
 void NearbySharingServiceImpl::OnFrameRead(
-    std::unique_ptr<IncomingFramesReader> frames_reader,
     ShareTarget share_target,
     base::Optional<sharing::mojom::V1FramePtr> frame) {
   if (!frame) {
@@ -1036,9 +1043,18 @@ void NearbySharingServiceImpl::OnFrameRead(
       break;
   }
 
-  frames_reader->ReadFrame(base::BindOnce(
-      &NearbySharingServiceImpl::OnFrameRead, weak_ptr_factory_.GetWeakPtr(),
-      std::move(frames_reader), std::move(share_target)));
+  auto* frames_reader =
+      GetIncomingShareTargetInfo(share_target).frames_reader();
+  if (!frames_reader) {
+    NS_LOG(WARNING) << __func__
+                    << ": Stopped reading further frames, due to no connection "
+                       "established.";
+    return;
+  }
+
+  frames_reader->ReadFrame(
+      base::BindOnce(&NearbySharingServiceImpl::OnFrameRead,
+                     weak_ptr_factory_.GetWeakPtr(), std::move(share_target)));
 }
 
 void NearbySharingServiceImpl::HandleCertificateInfoFrame(

chrome/browser/nearby_sharing/nearby_sharing_service_impl.h

@@ -151,10 +151,8 @@ class NearbySharingServiceImpl
   void OnReceivedIntroduction(
       ShareTarget share_target,
       base::Optional<std::string> token,
-      std::unique_ptr<IncomingFramesReader> frames_reader,
       base::Optional<sharing::mojom::V1FramePtr> frame);
-  void OnFrameRead(std::unique_ptr<IncomingFramesReader> frames_reader,
-                   ShareTarget share_target,
+  void OnFrameRead(ShareTarget share_target,
                    base::Optional<sharing::mojom::V1FramePtr> frame);
   void HandleCertificateInfoFrame(
       const sharing::mojom::CertificateInfoFramePtr& certificate_frame);

chrome/browser/nearby_sharing/nearby_sharing_service_impl_unittest.cc

@@ -747,6 +747,28 @@ TEST_F(NearbySharingServiceImplTest, UnregisterReceiveSurfaceNeverRegistered) {
   EXPECT_FALSE(fake_nearby_connections_manager_->IsAdvertising());
 }
 
+TEST_F(NearbySharingServiceImplTest,
+       IncomingConnection_ClosedReadingIntroduction) {
+  FakeNearbyConnection connection;
+  ui::ScopedSetIdleState unlocked(ui::IDLE_STATE_IDLE);
+  SetConnectionType(net::NetworkChangeNotifier::CONNECTION_WIFI);
+  NiceMock<MockTransferUpdateCallback> callback;
+  EXPECT_CALL(callback, OnTransferUpdate(testing::_, testing::_)).Times(0);
+
+  NearbySharingService::StatusCodes result = service_->RegisterReceiveSurface(
+      &callback, NearbySharingService::ReceiveSurfaceState::kForeground);
+  EXPECT_EQ(result, NearbySharingService::StatusCodes::kOk);
+  EXPECT_TRUE(fake_nearby_connections_manager_->IsAdvertising());
+
+  service_->OnIncomingConnection(kEndpointId, {}, &connection);
+  connection.Close();
+
+  // Introduction is ignored without any side effect.
+
+  // To avoid UAF in OnIncomingTransferUpdate().
+  service_->UnregisterReceiveSurface(&callback);
+}
+
 TEST_F(NearbySharingServiceImplTest,
        IncomingConnection_EmptyIntroductionFrame) {
   std::string intro = "introduction_frame";
@@ -849,6 +871,59 @@ TEST_F(NearbySharingServiceImplTest,
   service_->UnregisterReceiveSurface(&callback);
 }
 
+TEST_F(NearbySharingServiceImplTest,
+       IncomingConnection_ClosedWaitingLocalConfirmation) {
+  std::string intro = "introduction_frame";
+  std::vector<uint8_t> bytes(intro.begin(), intro.end());
+  NiceMock<MockNearbySharingDecoder> mock_decoder;
+  EXPECT_CALL(mock_decoder, DecodeFrame(testing::Eq(bytes), testing::_))
+      .WillOnce(testing::Invoke(
+          [&](const std::vector<uint8_t>& data,
+              MockNearbySharingDecoder::DecodeFrameCallback callback) {
+            std::move(callback).Run(GetValidIntroductionFrame());
+          }));
+
+  EXPECT_CALL(mock_nearby_process_manager(),
+              GetOrStartNearbySharingDecoder(testing::_))
+      .WillRepeatedly(testing::Return(&mock_decoder));
+
+  FakeNearbyConnection connection;
+  connection.AppendReadableData(bytes);
+  ui::ScopedSetIdleState unlocked(ui::IDLE_STATE_IDLE);
+  SetConnectionType(net::NetworkChangeNotifier::CONNECTION_WIFI);
+  NiceMock<MockTransferUpdateCallback> callback;
+  base::RunLoop run_loop;
+  EXPECT_CALL(callback, OnTransferUpdate(testing::_, testing::_))
+      .WillOnce(testing::Invoke([&run_loop](const ShareTarget& share_target,
+                                            TransferMetadata metadata) {
+        EXPECT_EQ(TransferMetadata::Status::kAwaitingLocalConfirmation,
+                  metadata.status());
+        run_loop.Quit();
+      }));
+
+  NearbySharingService::StatusCodes result = service_->RegisterReceiveSurface(
+      &callback, NearbySharingService::ReceiveSurfaceState::kForeground);
+  EXPECT_EQ(result, NearbySharingService::StatusCodes::kOk);
+  EXPECT_TRUE(fake_nearby_connections_manager_->IsAdvertising());
+
+  service_->OnIncomingConnection(kEndpointId, {}, &connection);
+  run_loop.Run();
+
+  base::RunLoop run_loop_2;
+  EXPECT_CALL(callback, OnTransferUpdate(testing::_, testing::_))
+      .WillOnce(testing::Invoke([&run_loop_2](const ShareTarget& share_target,
+                                              TransferMetadata metadata) {
+        EXPECT_EQ(TransferMetadata::Status::kFailed, metadata.status());
+        run_loop_2.Quit();
+      }));
+
+  connection.Close();
+  run_loop_2.Run();
+
+  // To avoid UAF in OnIncomingTransferUpdate().
+  service_->UnregisterReceiveSurface(&callback);
+}
+
 TEST_F(NearbySharingServiceImplTest, AcceptInvalidShareTarget) {
   ShareTarget share_target;
   base::RunLoop run_loop;