Tether MessageTransferOperation: Only wait for a response from a host for a certain amount of time …

chromeos/components/tether/message_transfer_operation_unittest.cc

Index: chromeos/components/tether/message_transfer_operation_unittest.cc
diff --git a/chromeos/components/tether/message_transfer_operation_unittest.cc b/chromeos/components/tether/message_transfer_operation_unittest.cc
index 5db05a38893b61d6e1f8589d3786a34fb20c0bd3..ab39735677fdccd7b47ff9bae0a71ca78025d902 100644
--- a/chromeos/components/tether/message_transfer_operation_unittest.cc
+++ b/chromeos/components/tether/message_transfer_operation_unittest.cc
@@ -4,8 +4,10 @@
 
 #include "chromeos/components/tether/message_transfer_operation.h"
 
+#include "base/timer/mock_timer.h"
 #include "chromeos/components/tether/fake_ble_connection_manager.h"
 #include "chromeos/components/tether/message_wrapper.h"
+#include "chromeos/components/tether/timer_factory.h"
 #include "components/cryptauth/remote_device_test_util.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -20,6 +22,8 @@ namespace {
 // except that it must be consistent throughout the test.
 const MessageType kTestMessageType = MessageType::TETHER_AVAILABILITY_REQUEST;
 
+const uint32_t kTestResponseTimeoutSeconds = 5;
+
 // A test double for MessageTransferOperation is needed because
 // MessageTransferOperation has pure virtual methods which must be overridden in
 // order to create a concrete instantiation of the class.
@@ -27,9 +31,7 @@ class TestOperation : public MessageTransferOperation {
  public:
   TestOperation(const std::vector<cryptauth::RemoteDevice>& devices_to_connect,
                 BleConnectionManager* connection_manager)
-      : MessageTransferOperation(devices_to_connect, connection_manager),
-        has_operation_started_(false),
-        has_operation_finished_(false) {}
+      : MessageTransferOperation(devices_to_connect, connection_manager) {}
   ~TestOperation() override {}
 
   bool HasDeviceAuthenticated(const cryptauth::RemoteDevice& remote_device) {
@@ -62,6 +64,9 @@ class TestOperation : public MessageTransferOperation {
       const cryptauth::RemoteDevice& remote_device) override {
     device_map_[remote_device].received_messages.push_back(
         std::move(message_wrapper));
+
+    if (should_unregister_device_on_message_received_)
+      UnregisterDevice(remote_device);
   }
 
   void OnOperationStarted() override { has_operation_started_ = true; }
@@ -72,6 +77,26 @@ class TestOperation : public MessageTransferOperation {
     return kTestMessageType;
   }
 
+  bool ShouldWaitForResponse() override { return should_wait_for_response_; }
+
+  uint32_t GetResponseTimeoutSeconds() override {
+    return response_timeout_seconds_;
+  }
+
+  void set_should_wait_for_response(bool should_wait_for_response) {
+    should_wait_for_response_ = should_wait_for_response;
+  }
+
+  void set_response_timeout_seconds(uint32_t response_timeout_seconds) {
+    response_timeout_seconds_ = response_timeout_seconds;
+  }
+
+  void set_should_unregister_device_on_message_received(
+      bool should_unregister_device_on_message_received) {
+    should_unregister_device_on_message_received_ =
+        should_unregister_device_on_message_received;
+  }
+
   bool has_operation_started() { return has_operation_started_; }
 
   bool has_operation_finished() { return has_operation_finished_; }
@@ -87,8 +112,38 @@ class TestOperation : public MessageTransferOperation {
 
   std::map<cryptauth::RemoteDevice, DeviceMapValue> device_map_;
 
-  bool has_operation_started_;
-  bool has_operation_finished_;
+  bool should_wait_for_response_ = true;
+  uint32_t response_timeout_seconds_ = kTestResponseTimeoutSeconds;
+  bool should_unregister_device_on_message_received_ = false;
+  bool has_operation_started_ = false;
+  bool has_operation_finished_ = false;
+};
+
+class TestTimerFactory : public TimerFactory {
+ public:
+  ~TestTimerFactory() override {}
+
+  // TimerFactory:
+  std::unique_ptr<base::Timer> CreateOneShotTimer() override {
+    EXPECT_FALSE(device_id_for_next_timer_.empty());
+    base::MockTimer* mock_timer = new base::MockTimer(
+        false /* retain_user_task */, false /* is_repeating */);
+    device_id_to_timer_map_[device_id_for_next_timer_] = mock_timer;
+    return base::WrapUnique(mock_timer);
+  }
+
+  base::MockTimer* GetResponseTimerForDeviceId(const std::string& device_id) {
+    return device_id_to_timer_map_[device_id_for_next_timer_];
+  }
+
+  void set_device_id_for_next_timer(
+      const std::string& device_id_for_next_timer) {
+    device_id_for_next_timer_ = device_id_for_next_timer;
+  }
+
+ private:
+  std::string device_id_for_next_timer_;
+  std::unordered_map<std::string, base::MockTimer*> device_id_to_timer_map_;
 };
 
 DeviceStatus CreateFakeDeviceStatus() {
@@ -131,8 +186,10 @@ class MessageTransferOperationTest : public testing::Test {
   }
 
   void ConstructOperation(std::vector<cryptauth::RemoteDevice> remote_devices) {
+    test_timer_factory_ = new TestTimerFactory();
     operation_ = base::WrapUnique(
         new TestOperation(remote_devices, fake_ble_connection_manager_.get()));
+    operation_->SetTimerFactoryForTest(base::WrapUnique(test_timer_factory_));
     VerifyOperationStartedAndFinished(false /* has_started */,
                                       false /* has_finished */);
   }
@@ -157,9 +214,43 @@ class MessageTransferOperationTest : public testing::Test {
     EXPECT_EQ(has_finished, operation_->has_operation_finished());
   }
 
+  void TransitionDeviceStatusFromDisconnectedToAuthenticated(
+      const cryptauth::RemoteDevice& remote_device) {
+    test_timer_factory_->set_device_id_for_next_timer(
+        remote_device.GetDeviceId());
+
+    fake_ble_connection_manager_->SetDeviceStatus(
+        remote_device, cryptauth::SecureChannel::Status::CONNECTING);
+    fake_ble_connection_manager_->SetDeviceStatus(
+        remote_device, cryptauth::SecureChannel::Status::CONNECTED);
+    fake_ble_connection_manager_->SetDeviceStatus(
+        remote_device, cryptauth::SecureChannel::Status::AUTHENTICATING);
+    fake_ble_connection_manager_->SetDeviceStatus(
+        remote_device, cryptauth::SecureChannel::Status::AUTHENTICATED);
+  }
+
+  base::MockTimer* GetResponseTimerForDevice(
+      const cryptauth::RemoteDevice& remote_device) {
+    return test_timer_factory_->GetResponseTimerForDeviceId(
+        remote_device.GetDeviceId());
+  }
+
+  void VerifyDefaultTimerCreatedForDevice(
+      const cryptauth::RemoteDevice& remote_device) {
+    VerifyTimerCreatedForDevice(remote_device, kTestResponseTimeoutSeconds);
+  }
+
+  void VerifyTimerCreatedForDevice(const cryptauth::RemoteDevice& remote_device,
+                                   uint32_t response_timeout_seconds) {
+    EXPECT_TRUE(GetResponseTimerForDevice(remote_device));
+    EXPECT_EQ(base::TimeDelta::FromSeconds(response_timeout_seconds),
+              GetResponseTimerForDevice(remote_device)->GetCurrentDelay());
+  }
+
   const std::vector<cryptauth::RemoteDevice> test_devices_;
 
   std::unique_ptr<FakeBleConnectionManager> fake_ble_connection_manager_;
+  TestTimerFactory* test_timer_factory_;
   std::unique_ptr<TestOperation> operation_;
 
  private:
@@ -236,16 +327,10 @@ TEST_F(MessageTransferOperationTest, TestFailsThenConnects) {
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
 
   // Try again and succeed.
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTED);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATED);
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
   EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[0]);
 
   EXPECT_TRUE(operation_->GetReceivedMessages(test_devices_[0]).empty());
 }
@@ -256,17 +341,75 @@ TEST_F(MessageTransferOperationTest,
   InitializeOperation();
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
 
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTED);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATED);
+  // Simulate how subclasses behave after a successful response: unregister the
+  // device.
+  operation_->set_should_unregister_device_on_message_received(true);
+
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
+  EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+  EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[0]);
+
+  fake_ble_connection_manager_->ReceiveMessage(
+      test_devices_[0],
+      MessageWrapper(CreateTetherAvailabilityResponse()).ToRawMessage());
+
+  EXPECT_EQ(1u, operation_->GetReceivedMessages(test_devices_[0]).size());
+  std::shared_ptr<MessageWrapper> message =
+      operation_->GetReceivedMessages(test_devices_[0])[0];
+  EXPECT_EQ(MessageType::TETHER_AVAILABILITY_RESPONSE,
+            message->GetMessageType());
+  EXPECT_EQ(CreateTetherAvailabilityResponse().SerializeAsString(),
+            message->GetProto()->SerializeAsString());
+}
+
+TEST_F(MessageTransferOperationTest,
+       TestSuccessfulConnectionAndReceiveMessage_ResponseTimeoutSeconds) {
+  const uint32_t response_timeout_seconds = 90;
+
+  ConstructOperation(std::vector<cryptauth::RemoteDevice>{test_devices_[0]});
+  InitializeOperation();
+  EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+
+  operation_->set_response_timeout_seconds(response_timeout_seconds);
+
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
+  EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+  EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
+  VerifyTimerCreatedForDevice(test_devices_[0], response_timeout_seconds);
+
+  EXPECT_EQ(base::TimeDelta::FromSeconds(response_timeout_seconds),
+            GetResponseTimerForDevice(test_devices_[0])->GetCurrentDelay());
+
+  fake_ble_connection_manager_->ReceiveMessage(
+      test_devices_[0],
+      MessageWrapper(CreateTetherAvailabilityResponse()).ToRawMessage());
+
+  EXPECT_EQ(1u, operation_->GetReceivedMessages(test_devices_[0]).size());
+  std::shared_ptr<MessageWrapper> message =
+      operation_->GetReceivedMessages(test_devices_[0])[0];
+  EXPECT_EQ(MessageType::TETHER_AVAILABILITY_RESPONSE,
+            message->GetMessageType());
+  EXPECT_EQ(CreateTetherAvailabilityResponse().SerializeAsString(),
+            message->GetProto()->SerializeAsString());
+}
+
+TEST_F(MessageTransferOperationTest,
+       TestSuccessfulConnectionAndReceiveMessage_ShouldNotWaitForResponse) {
+  ConstructOperation(std::vector<cryptauth::RemoteDevice>{test_devices_[0]});
+  InitializeOperation();
+  EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+
+  operation_->set_should_wait_for_response(false);
+
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
   EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
 
+  // A Timer should not have been created because the operation should not wait
+  // for a response.
+  EXPECT_FALSE(GetResponseTimerForDevice(test_devices_[0]));
+
   fake_ble_connection_manager_->ReceiveMessage(
       test_devices_[0],
       MessageWrapper(CreateTetherAvailabilityResponse()).ToRawMessage());
@@ -280,6 +423,22 @@ TEST_F(MessageTransferOperationTest,
             message->GetProto()->SerializeAsString());
 }
 
+TEST_F(MessageTransferOperationTest, TestAuthenticatesButTimesOut) {
+  ConstructOperation(std::vector<cryptauth::RemoteDevice>{test_devices_[0]});
+  InitializeOperation();
+  EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
+  EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+  EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[0]);
+
+  GetResponseTimerForDevice(test_devices_[0])->Fire();
+
+  EXPECT_FALSE(IsDeviceRegistered(test_devices_[0]));
+  EXPECT_TRUE(operation_->has_operation_finished());
+}
+
 TEST_F(MessageTransferOperationTest, TestRepeatedInputDevice) {
   // Construct with two copies of the same device.
   ConstructOperation(
@@ -287,16 +446,10 @@ TEST_F(MessageTransferOperationTest, TestRepeatedInputDevice) {
   InitializeOperation();
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
 
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTED);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATED);
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
   EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[0]);
 
   fake_ble_connection_manager_->ReceiveMessage(
       test_devices_[0],
@@ -323,14 +476,7 @@ TEST_F(MessageTransferOperationTest, TestReceiveEventForOtherDevice) {
   // should not be affected.
   fake_ble_connection_manager_->RegisterRemoteDevice(
       test_devices_[1], MessageType::CONNECT_TETHERING_REQUEST);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[1], cryptauth::SecureChannel::Status::CONNECTING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[1], cryptauth::SecureChannel::Status::CONNECTED);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[1], cryptauth::SecureChannel::Status::AUTHENTICATING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[1], cryptauth::SecureChannel::Status::AUTHENTICATED);
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[1]);
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
   EXPECT_FALSE(IsDeviceRegistered(test_devices_[1]));
   EXPECT_FALSE(operation_->HasDeviceAuthenticated(test_devices_[0]));
@@ -353,19 +499,13 @@ TEST_F(MessageTransferOperationTest,
   // initialization.
   fake_ble_connection_manager_->RegisterRemoteDevice(
       test_devices_[0], MessageType::CONNECT_TETHERING_REQUEST);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTED);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATED);
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
 
   // Now initialize; the authentication handler should have been invoked.
   InitializeOperation();
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
   EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[0]);
 
   // Receiving a message should work at this point.
   fake_ble_connection_manager_->ReceiveMessage(
@@ -381,6 +521,28 @@ TEST_F(MessageTransferOperationTest,
             message->GetProto()->SerializeAsString());
 }
 
+TEST_F(MessageTransferOperationTest,
+       AlreadyAuthenticatedBeforeInitialization_TimesOutWaitingForResponse) {
+  ConstructOperation(std::vector<cryptauth::RemoteDevice>{test_devices_[0]});
+
+  // Simulate the authentication of |test_devices_[0]|'s channel before
+  // initialization.
+  fake_ble_connection_manager_->RegisterRemoteDevice(
+      test_devices_[0], MessageType::CONNECT_TETHERING_REQUEST);
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
+
+  // Now initialize; the authentication handler should have been invoked.
+  InitializeOperation();
+  EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+  EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[0]);
+
+  GetResponseTimerForDevice(test_devices_[0])->Fire();
+
+  EXPECT_FALSE(IsDeviceRegistered(test_devices_[0]));
+  EXPECT_TRUE(operation_->has_operation_finished());
+}
+
 TEST_F(MessageTransferOperationTest, MultipleDevices) {
   ConstructOperation(test_devices_);
   InitializeOperation();
@@ -392,18 +554,14 @@ TEST_F(MessageTransferOperationTest, MultipleDevices) {
   // Authenticate |test_devices_[0]|'s channel.
   fake_ble_connection_manager_->RegisterRemoteDevice(
       test_devices_[0], MessageType::CONNECT_TETHERING_REQUEST);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::CONNECTED);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[0], cryptauth::SecureChannel::Status::AUTHENTICATED);
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[0]);
   EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[0]));
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[0]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[0]);
 
   // Fail 3 times to connect to |test_devices_[1]|.
+  test_timer_factory_->set_device_id_for_next_timer(
+      test_devices_[1].GetDeviceId());
   fake_ble_connection_manager_->SetDeviceStatus(
       test_devices_[1], cryptauth::SecureChannel::Status::CONNECTING);
   fake_ble_connection_manager_->SetDeviceStatus(
@@ -418,22 +576,19 @@ TEST_F(MessageTransferOperationTest, MultipleDevices) {
       test_devices_[1], cryptauth::SecureChannel::Status::DISCONNECTED);
   EXPECT_FALSE(operation_->HasDeviceAuthenticated(test_devices_[1]));
   EXPECT_FALSE(IsDeviceRegistered(test_devices_[1]));
+  EXPECT_FALSE(GetResponseTimerForDevice(test_devices_[1]));
 
   // Authenticate |test_devices_[2]|'s channel.
   fake_ble_connection_manager_->RegisterRemoteDevice(
       test_devices_[2], MessageType::CONNECT_TETHERING_REQUEST);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[2], cryptauth::SecureChannel::Status::CONNECTING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[2], cryptauth::SecureChannel::Status::CONNECTED);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[2], cryptauth::SecureChannel::Status::AUTHENTICATING);
-  fake_ble_connection_manager_->SetDeviceStatus(
-      test_devices_[2], cryptauth::SecureChannel::Status::AUTHENTICATED);
+  TransitionDeviceStatusFromDisconnectedToAuthenticated(test_devices_[2]);
   EXPECT_TRUE(operation_->HasDeviceAuthenticated(test_devices_[2]));
   EXPECT_TRUE(IsDeviceRegistered(test_devices_[2]));
+  VerifyDefaultTimerCreatedForDevice(test_devices_[2]);
 
   // Fail to authenticate |test_devices_[3]|'s channel.
+  test_timer_factory_->set_device_id_for_next_timer(
+      test_devices_[3].GetDeviceId());
   fake_ble_connection_manager_->RegisterRemoteDevice(
       test_devices_[3], MessageType::CONNECT_TETHERING_REQUEST);
   fake_ble_connection_manager_->SetDeviceStatus(
@@ -446,6 +601,7 @@ TEST_F(MessageTransferOperationTest, MultipleDevices) {
       test_devices_[3], cryptauth::SecureChannel::Status::DISCONNECTED);
   EXPECT_FALSE(operation_->HasDeviceAuthenticated(test_devices_[3]));
   EXPECT_FALSE(IsDeviceRegistered(test_devices_[3]));
+  EXPECT_FALSE(GetResponseTimerForDevice(test_devices_[3]));
 }
 
 }  // namespace tether