Add timeout task to U2F HID state machine

device/u2f/u2f_hid_device.cc

 // Copyright 2017 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "u2f_hid_device.h"
 
 #include "base/bind.h"
 #include "base/bind_helpers.h"
 #include "base/command_line.h"
+#include "base/threading/thread_task_runner_handle.h"
 #include "crypto/random.h"
 #include "device/base/device_client.h"
 #include "device/hid/hid_connection.h"
 #include "u2f_apdu_command.h"
 #include "u2f_message.h"
 
 namespace device {
 
 namespace switches {
 static constexpr char kEnableU2fHidTest[] = "enable-u2f-hid-tests";
 }  // namespace switches
 
 U2fHidDevice::U2fHidDevice(scoped_refptr<HidDeviceInfo> device_info)
     : U2fDevice(),
       state_(State::INIT),
       device_info_(device_info),
       weak_factory_(this) {
   channel_id_ = kBroadcastChannel;
 }
 
 U2fHidDevice::~U2fHidDevice() {
   // Cleanup connection
-  if (connection_)
+  if (connection_ && !connection_->closed())
     connection_->Close();
 }
 
 void U2fHidDevice::DeviceTransact(std::unique_ptr<U2fApduCommand> command,
                                   const DeviceCallback& callback) {
   Transition(std::move(command), callback);
 }
 
 void U2fHidDevice::Transition(std::unique_ptr<U2fApduCommand> command,
                               const DeviceCallback& callback) {
   switch (state_) {
-    case State::INIT:
+    case State::INIT: {
       state_ = State::BUSY;
+      timeout_callback_.Reset(base::Bind(&U2fHidDevice::OnTimeout,
+                                         weak_factory_.GetWeakPtr(), callback));
+      auto self = weak_factory_.GetWeakPtr();
       Connect(base::Bind(&U2fHidDevice::OnConnect, weak_factory_.GetWeakPtr(),
                          base::Passed(&command), callback));
+      if (self && !timeout_callback_.IsCancelled()) {
+        // Setup timeout task for 3 seconds

[Reilly Grant (use Gerrit), 2017/04/17 22:25:17]

Set up the timeout task before calling Connect to avoid needing to protect
yourself against reentrancy here.

I would put this logic into a function like this:

void ArmTimeout(const DeviceCallback& callback) {
  timeout_callback_.Reset(base::Bind(&U2fHidDevice::OnTimeout,
                                     weak_factory_.GetWeakPtr(), callback));
  base::ThreadTaskRunnerHandle::Get()->PostDelayedTask(
      FROM_HERE, timeout_callback_.callback(),
      base::TimeDelta::FromMilliseconds(3000));
}

If I understand the state machine correctly you should also be able to add a
DCHECK(!timeout_callback_.IsCancelled()) to the top of this function.

[Casey Piper, 2017/04/17 23:02:38]

Done. Did you mean DCHECK(timeout_callback.IsCancelled()) since the callback
wouldn't yet be armed?

+        base::ThreadTaskRunnerHandle::Get()->PostDelayedTask(
+            FROM_HERE, timeout_callback_.callback(),
+            base::TimeDelta::FromMilliseconds(3000));
+      }
       break;
-    case State::CONNECTED:
+    }
+    case State::CONNECTED: {
       state_ = State::BUSY;
+      timeout_callback_.Reset(base::Bind(&U2fHidDevice::OnTimeout,
+                                         weak_factory_.GetWeakPtr(), callback));
+      auto self = weak_factory_.GetWeakPtr();
       AllocateChannel(std::move(command), callback);
+      if (self && !timeout_callback_.IsCancelled()) {
+        // Setup timeout task for 3 seconds
+        base::ThreadTaskRunnerHandle::Get()->PostDelayedTask(
+            FROM_HERE, timeout_callback_.callback(),
+            base::TimeDelta::FromMilliseconds(3000));
+      }
       break;
+    }
     case State::IDLE: {
       state_ = State::BUSY;
       std::unique_ptr<U2fMessage> msg = U2fMessage::Create(
           channel_id_, U2fMessage::Type::CMD_MSG, command->GetEncodedCommand());
+
+      timeout_callback_.Reset(base::Bind(&U2fHidDevice::OnTimeout,
+                                         weak_factory_.GetWeakPtr(), callback));
+      auto self = weak_factory_.GetWeakPtr();
+      // Write message to the device
       WriteMessage(std::move(msg), true,
                    base::Bind(&U2fHidDevice::MessageReceived,
                               weak_factory_.GetWeakPtr(), callback));
+
+      if (self && !timeout_callback_.IsCancelled()) {
+        // Setup timeout task for 3 seconds
+        base::ThreadTaskRunnerHandle::Get()->PostDelayedTask(
+            FROM_HERE, timeout_callback_.callback(),
+            base::TimeDelta::FromMilliseconds(3000));
+      }
       break;
     }
     case State::BUSY:
       pending_transactions_.push_back({std::move(command), callback});
       break;
     case State::DEVICE_ERROR:
     default:
       base::WeakPtr<U2fHidDevice> self = weak_factory_.GetWeakPtr();
       callback.Run(false, nullptr);
       // Executing callbacks may free |this|. Check |self| first.
       while (self && !pending_transactions_.empty()) {
         // Respond to any pending requests
         DeviceCallback pending_cb = pending_transactions_.front().second;
         pending_transactions_.pop_front();
         pending_cb.Run(false, nullptr);
       }
       break;
   }
 }
 
 void U2fHidDevice::Connect(const HidService::ConnectCallback& callback) {
   HidService* hid_service = DeviceClient::Get()->GetHidService();
 
   hid_service->Connect(device_info_->device_id(), callback);
 }
 
 void U2fHidDevice::OnConnect(std::unique_ptr<U2fApduCommand> command,
                              const DeviceCallback& callback,
                              scoped_refptr<HidConnection> connection) {
+  if (timeout_callback_.IsCancelled()) {
+    return;
+  }

[Reilly Grant (use Gerrit), 2017/04/17 22:25:17]

nit: no braces around single-line if.

The state machine may be clearer if you check state_ != State::DEVICE_ERROR here
instead of timeout_callback_ specifically.

[Casey Piper, 2017/04/17 23:02:38]

Done.

+  timeout_callback_.Cancel();
+
   if (connection) {
     connection_ = connection;
     state_ = State::CONNECTED;
   } else {
     state_ = State::DEVICE_ERROR;
   }
   Transition(std::move(command), callback);
 }
 
 void U2fHidDevice::AllocateChannel(std::unique_ptr<U2fApduCommand> command,
                                    const DeviceCallback& callback) {
   // Send random nonce to device to verify received message
   std::vector<uint8_t> nonce(8);
   crypto::RandBytes(nonce.data(), nonce.size());
   std::unique_ptr<U2fMessage> message =
       U2fMessage::Create(channel_id_, U2fMessage::Type::CMD_INIT, nonce);
 
   WriteMessage(
       std::move(message), true,
       base::Bind(&U2fHidDevice::OnAllocateChannel, weak_factory_.GetWeakPtr(),
                  nonce, base::Passed(&command), callback));
 }
 
 void U2fHidDevice::OnAllocateChannel(std::vector<uint8_t> nonce,
                                      std::unique_ptr<U2fApduCommand> command,
                                      const DeviceCallback& callback,
                                      bool success,
                                      std::unique_ptr<U2fMessage> message) {
+  if (timeout_callback_.IsCancelled()) {
+    return;
+  }

[Reilly Grant (use Gerrit), 2017/04/17 22:25:17]

nit: no braces around single-line if.

[Casey Piper, 2017/04/17 23:02:38]

Done.

+  timeout_callback_.Cancel();
+
   if (!success || !message) {
     state_ = State::DEVICE_ERROR;
     Transition(nullptr, callback);
     return;
   }
   // Channel allocation response is defined as:
   // 0: 8 byte nonce
   // 8: 4 byte channel id
   // 12: Protocol version id
   // 13: Major device version
@@ skipping 79 matching lines @@
   std::unique_ptr<U2fMessage> read_message =
       U2fMessage::CreateFromSerializedData(read_buffer);
 
   if (!read_message) {
     std::move(callback).Run(false, nullptr);
     return;
   }
 
   // Received a message from a different channel, so try again
   if (channel_id_ != read_message->channel_id()) {
-    connection_->Read(base::Bind(&U2fHidDevice::OnRead,
-                                 weak_factory_.GetWeakPtr(),
-                                 base::Passed(&callback)));
+    base::ThreadTaskRunnerHandle::Get()->PostTask(
+        FROM_HERE,
+        base::Bind(&U2fHidDevice::ReadMessage, weak_factory_.GetWeakPtr(),
+                   base::Passed(&callback)));
     return;
   }
 
   if (read_message->MessageComplete()) {
     std::move(callback).Run(success, std::move(read_message));
     return;
   }
 
   // Continue reading additional packets
   connection_->Read(
@@ skipping 15 matching lines @@
   message->AddContinuationPacket(read_buffer);
   if (message->MessageComplete()) {
     std::move(callback).Run(success, std::move(message));
     return;
   }
   connection_->Read(
       base::Bind(&U2fHidDevice::OnReadContinuation, weak_factory_.GetWeakPtr(),
                  base::Passed(&message), base::Passed(&callback)));
 }
 
-void U2fHidDevice::MessageReceived(const DeviceCallback& callback,
+void U2fHidDevice::MessageReceived(DeviceCallback callback,
                                    bool success,
                                    std::unique_ptr<U2fMessage> message) {
+  if (timeout_callback_.IsCancelled()) {
+    return;
+  }

[Reilly Grant (use Gerrit), 2017/04/17 22:25:17]

nit: no braces around single-line if.

[Casey Piper, 2017/04/17 23:02:38]

Done.

+  timeout_callback_.Cancel();
+
   if (!success) {
     state_ = State::DEVICE_ERROR;
     Transition(nullptr, callback);
     return;
   }
+
   std::unique_ptr<U2fApduResponse> response = nullptr;
   if (message)
     response = U2fApduResponse::CreateFromMessage(message->GetMessagePayload());
   state_ = State::IDLE;
   base::WeakPtr<U2fHidDevice> self = weak_factory_.GetWeakPtr();
   callback.Run(success, std::move(response));
 
   // Executing |callback| may have freed |this|. Check |self| first.
   if (self && !pending_transactions_.empty()) {
     // If any transactions were queued, process the first one
@@ skipping 18 matching lines @@
       std::move(wink_message), true,
       base::Bind(&U2fHidDevice::OnWink, weak_factory_.GetWeakPtr(), callback));
 }
 
 void U2fHidDevice::OnWink(const WinkCallback& callback,
                           bool success,
                           std::unique_ptr<U2fMessage> response) {
   callback.Run();
 }
 
+void U2fHidDevice::OnTimeout(DeviceCallback callback) {
+  timeout_callback_.Cancel();

[Reilly Grant (use Gerrit), 2017/04/17 22:25:17]

Since timeout_callback_ can only fire once it should not need to be cancelled
here.

[Casey Piper, 2017/04/17 23:02:38]

Done.

+  state_ = State::DEVICE_ERROR;
+  Transition(nullptr, callback);
+}
+
 std::string U2fHidDevice::GetId() {
-  std::ostringstream id("hid:");
+  std::ostringstream id("hid:", std::ios::ate);
   id << device_info_->device_id();
   return id.str();
 }
 
 // static
 bool U2fHidDevice::IsTestEnabled() {
   base::CommandLine* command_line = base::CommandLine::ForCurrentProcess();
   return command_line->HasSwitch(switches::kEnableU2fHidTest);
 }
 
 }  // namespace device