Pass task_runner_ to server and send 404 when incorrect url

components/ui_devtools/devtools_server.cc

 // Copyright 2016 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 "components/ui_devtools/devtools_server.h"
 
 #include "base/bind.h"
 #include "base/bind_helpers.h"
 #include "base/command_line.h"
 #include "base/format_macros.h"
@@ skipping 10 matching lines @@
 
 namespace ui {
 namespace devtools {
 
 namespace {
 const char kChromeDeveloperToolsPrefix[] =
     "chrome-devtools://devtools/bundled/inspector.html?ws=";
 }  // namespace
 
 UiDevToolsServer::UiDevToolsServer()
-    : thread_(new base::Thread("UiDevToolsServerThread")) {}
+    : thread_(new base::Thread("UiDevToolsServerThread")) {
+  base::Thread::Options options;
+  options.message_loop_type = base::MessageLoop::TYPE_IO;
+  CHECK(thread_->StartWithOptions(options));
+  task_runner_ = thread_->task_runner();
+}
+
+UiDevToolsServer::UiDevToolsServer(
+    scoped_refptr<base::SingleThreadTaskRunner> task_runner)
+    : task_runner_(task_runner) {
+  DCHECK(task_runner_);
+}
 
 UiDevToolsServer::~UiDevToolsServer() {}
 
 // static
 std::unique_ptr<UiDevToolsServer> UiDevToolsServer::Create() {
+  return Create(nullptr);
+}
+
+// static
+std::unique_ptr<UiDevToolsServer> UiDevToolsServer::Create(
+    scoped_refptr<base::SingleThreadTaskRunner> task_runner) {
   std::unique_ptr<UiDevToolsServer> server;
   if (base::CommandLine::ForCurrentProcess()->HasSwitch(kEnableUiDevTools)) {
     // TODO(mhashmi): Change port if more than one inspectable clients
     int port = 9223;  // Default port is 9223
     base::StringToInt(
         base::CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
             kEnableUiDevTools),
         &port);
-    server.reset(new UiDevToolsServer());
+    server.reset(task_runner ? new UiDevToolsServer(task_runner)
+                             : new UiDevToolsServer());
     server->Start("127.0.0.1", port);
   }
   return server;
 }
 
 void UiDevToolsServer::AttachClient(std::unique_ptr<UiDevToolsClient> client) {
   clients_.push_back(std::move(client));
 }
 
 void UiDevToolsServer::SendOverWebSocket(int connection_id,
                                          const String& message) {
-  thread_->task_runner()->PostTask(
+  task_runner_->PostTask(
       FROM_HERE,
       base::Bind(&net::HttpServer::SendOverWebSocket,
                  base::Unretained(server_.get()), connection_id, message));
 }
 
 void UiDevToolsServer::Start(const std::string& address_string, uint16_t port) {
-  if (thread_ && thread_->IsRunning())
-    return;
-
-  // Start IO thread upon which all the methods will run
-  base::Thread::Options options;
-  options.message_loop_type = base::MessageLoop::TYPE_IO;
-  if (thread_->StartWithOptions(options)) {
-    thread_->task_runner()->PostTask(
-        FROM_HERE, base::Bind(&UiDevToolsServer::StartServer,
-                              base::Unretained(this), address_string, port));
-  }
+  task_runner_->PostTask(
+      FROM_HERE, base::Bind(&UiDevToolsServer::StartServer,
+                            base::Unretained(this), address_string, port));
 }
 
 void UiDevToolsServer::StartServer(const std::string& address_string,
                                    uint16_t port) {
   std::unique_ptr<net::ServerSocket> socket(
       new net::TCPServerSocket(nullptr, net::NetLogSource()));
   constexpr int kBacklog = 1;
   if (socket->ListenWithAddressAndPort(address_string, port, kBacklog) !=
       net::OK)
     return;
@@ skipping 17 matching lines @@
         "them:</h3>";
     net::IPEndPoint ip;
     server_->GetLocalAddress(&ip);
     for (ClientsList::size_type i = 0; i != clients_.size(); i++) {
       clientHTML += base::StringPrintf(
           "<p><strong>%s</strong> (%s%s/%" PRIuS ")</p>",
           clients_[i]->name().c_str(), kChromeDeveloperToolsPrefix,
           ip.ToString().c_str(), i);
     }
     clientHTML += "</html>";
-    thread_->task_runner()->PostTask(
+    task_runner_->PostTask(
         FROM_HERE,
         base::Bind(&net::HttpServer::Send200, base::Unretained(server_.get()),
                    connection_id, clientHTML, "text/html"));
+    return;
   }
+  task_runner_->PostTask(
+      FROM_HERE, base::Bind(&net::HttpServer::Send404,
+                            base::Unretained(server_.get()), connection_id));
 }
 
 void UiDevToolsServer::OnWebSocketRequest(
     int connection_id,
     const net::HttpServerRequestInfo& info) {
   size_t target_id = 0;
   if (info.path.empty() ||
       !base::StringToSizeT(info.path.substr(1), &target_id) ||
       target_id > clients_.size())
     return;
 
   UiDevToolsClient* client = clients_[target_id].get();
   // Only one user can inspect the client at a time
   if (client->connected())
     return;
   client->set_connection_id(connection_id);
   connections_[connection_id] = client;
-  thread_->task_runner()->PostTask(
+  task_runner_->PostTask(
       FROM_HERE,
       base::Bind(&net::HttpServer::AcceptWebSocket,
                  base::Unretained(server_.get()), connection_id, info));
 }
 
 void UiDevToolsServer::OnWebSocketMessage(int connection_id,
                                           const std::string& data) {
   ConnectionsMap::iterator it = connections_.find(connection_id);
   DCHECK(it != connections_.end());
   UiDevToolsClient* client = it->second;
   DCHECK(client);
-  thread_->task_runner()->PostTask(
-      FROM_HERE,
-      base::Bind(&UiDevToolsClient::Dispatch, base::Unretained(client), data));
+  task_runner_->PostTask(FROM_HERE, base::Bind(&UiDevToolsClient::Dispatch,
+                                               base::Unretained(client), data));
 }
 
 void UiDevToolsServer::OnClose(int connection_id) {
   ConnectionsMap::iterator it = connections_.find(connection_id);
-  DCHECK(it != connections_.end());
+  if (it == connections_.end())
+    return;
   UiDevToolsClient* client = it->second;
   DCHECK(client);
   client->set_connection_id(UiDevToolsClient::kNotConnected);
   connections_.erase(it);
 }
 
 }  // namespace devtools
 }  // namespace ui