[chromedriver] Bidi WebSocket connection(I)

Re-organize (ChromeDriver) http server into its own class so that
http handler can be a friend of http server. Then we can handle
the http server methods by posting on cmd thread to call
http handler methods. The reason to post on cmd thread is
that we can ensure thread safety when manipulating session map.

Also, pass cmd task runner and pointer to session map to
InitSessionParams so that session_commands can manipulate
the session map upon ExecuteInitSession.

Will make follow up CL to add condition check before manipulating the
session map. Follow up CL will also need to parse user WebSocketUrl
parameter.

There is also pending discussion about the behavior WebSocket connection
if the session terminates in the CL code review comments.

Bug: chromedriver:3588
Change-Id: Ibbd8b12842cf2e5b93d67bf68e1b869bb93f20c6
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2382299
Commit-Queue: Shengfa Lin <shengfa@google.com>
Reviewed-by: John Chen <johnchen@chromium.org>
Cr-Commit-Position: refs/heads/master@{#803644}

chrome/test/chromedriver/BUILD.gn

@@ -225,10 +225,13 @@ source_set("lib") {
     "performance_logger.h",
     "server/http_handler.cc",
     "server/http_handler.h",
+    "server/http_server.cc",
+    "server/http_server.h",
     "session.cc",
     "session.h",
     "session_commands.cc",
     "session_commands.h",
+    "session_connection_map.h",
     "session_thread_map.h",
     "util.cc",
     "util.h",

chrome/test/chromedriver/server/http_handler.cc

@@ -30,6 +30,7 @@
 #include "chrome/test/chromedriver/chrome/status.h"
 #include "chrome/test/chromedriver/constants/version.h"
 #include "chrome/test/chromedriver/net/url_request_context_getter.h"
+#include "chrome/test/chromedriver/server/http_server.h"
 #include "chrome/test/chromedriver/session.h"
 #include "chrome/test/chromedriver/session_thread_map.h"
 #include "chrome/test/chromedriver/util.h"
@@ -139,6 +140,7 @@ HttpHandler::HttpHandler(const std::string& url_base)
 HttpHandler::HttpHandler(
     const base::RepeatingClosure& quit_func,
     const scoped_refptr<base::SingleThreadTaskRunner> io_task_runner,
+    const scoped_refptr<base::SingleThreadTaskRunner> cmd_task_runner,
     const std::string& url_base,
     int adb_port)
     : quit_func_(quit_func), url_base_(url_base), received_shutdown_(false) {
@@ -163,12 +165,14 @@ HttpHandler::HttpHandler(
           kPost, internal::kNewSessionPathPattern,
           base::BindRepeating(
               &ExecuteCreateSession, &session_thread_map_,
-              WrapToCommand("InitSession",
-                            base::BindRepeating(
-                                &ExecuteInitSession,
-                                InitSessionParams(
-                                    wrapper_url_loader_factory_.get(),
-                                    socket_factory_, device_manager_.get()))))),
+              WrapToCommand(
+                  "InitSession",
+                  base::BindRepeating(
+                      &ExecuteInitSession,
+                      InitSessionParams(wrapper_url_loader_factory_.get(),
+                                        socket_factory_, device_manager_.get(),
+                                        cmd_task_runner,
+                                        &session_connection_map_))))),
       CommandMapping(kDelete, "session/:sessionId",
                      base::BindRepeating(
                          &ExecuteSessionCommand, &session_thread_map_, "Quit",
@@ -975,6 +979,10 @@ void HttpHandler::Handle(const net::HttpServerRequestInfo& request,
     received_shutdown_ = true;
 }
 
+base::WeakPtr<HttpHandler> HttpHandler::WeakPtr() {
+  return weak_ptr_factory_.GetWeakPtr();
+}
+
 Command HttpHandler::WrapToCommand(const char* name,
                                    const SessionCommand& session_command,
                                    bool w3c_standard_command) {
@@ -1279,6 +1287,10 @@ HttpHandler::PrepareStandardResponse(
   return response;
 }
 
+void HttpHandler::OnWebSocketRequest(int connection_id,
+                                     const net::HttpServerRequestInfo& info) {}
+
+void HttpHandler::OnClose(int connection_id) {}
 
 namespace internal {
 

chrome/test/chromedriver/server/http_handler.h

@@ -21,6 +21,7 @@
 #include "chrome/test/chromedriver/element_commands.h"
 #include "chrome/test/chromedriver/net/sync_websocket_factory.h"
 #include "chrome/test/chromedriver/session_commands.h"
+#include "chrome/test/chromedriver/session_connection_map.h"
 #include "chrome/test/chromedriver/session_thread_map.h"
 #include "chrome/test/chromedriver/window_commands.h"
 
@@ -43,6 +44,8 @@ class DeviceManager;
 class URLRequestContextGetter;
 class WrapperURLLoaderFactory;
 
+class HttpServer;
+
 enum HttpMethod {
   kGet,
   kPost,
@@ -73,6 +76,7 @@ class HttpHandler {
   explicit HttpHandler(const std::string& url_base);
   HttpHandler(const base::RepeatingClosure& quit_func,
               const scoped_refptr<base::SingleThreadTaskRunner> io_task_runner,
+              const scoped_refptr<base::SingleThreadTaskRunner> cmd_task_runner,
               const std::string& url_base,
               int adb_port);
   ~HttpHandler();
@@ -80,6 +84,8 @@ class HttpHandler {
   void Handle(const net::HttpServerRequestInfo& request,
               const HttpResponseSenderFunc& send_response_func);
 
+  base::WeakPtr<HttpHandler> WeakPtr();
+
  private:
   FRIEND_TEST_ALL_PREFIXES(HttpHandlerTest, HandleUnknownCommand);
   FRIEND_TEST_ALL_PREFIXES(HttpHandlerTest, HandleNewSession);
@@ -89,6 +95,8 @@ class HttpHandler {
   FRIEND_TEST_ALL_PREFIXES(HttpHandlerTest, StandardResponse_ErrorNoMessage);
   typedef std::vector<CommandMapping> CommandMap;
 
+  friend class HttpServer;
+
   Command WrapToCommand(const char* name,
                         const SessionCommand& session_command,
                         bool w3c_standard_command = true);
@@ -119,6 +127,11 @@ class HttpHandler {
       std::unique_ptr<base::Value> value,
       const std::string& session_id);
 
+  void OnWebSocketRequest(int connection_id,
+                          const net::HttpServerRequestInfo& info);
+
+  void OnClose(int connection_id);
+
   base::ThreadChecker thread_checker_;
   base::RepeatingClosure quit_func_;
   std::string url_base_;
@@ -129,6 +142,7 @@ class HttpHandler {
   std::unique_ptr<WrapperURLLoaderFactory> wrapper_url_loader_factory_;
   SyncWebSocketFactory socket_factory_;
   SessionThreadMap session_thread_map_;
+  SessionConnectionMap session_connection_map_;
   std::unique_ptr<CommandMap> command_map_;
   std::unique_ptr<Adb> adb_;
   std::unique_ptr<DeviceManager> device_manager_;

chrome/test/chromedriver/server/http_server.cc

+// Copyright (c) 2020 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 "chrome/test/chromedriver/server/http_server.h"
+
+namespace {
+
+// Maximum message size between app and ChromeDriver. Data larger than 150 MB
+// or so can cause crashes in Chrome (https://crbug.com/890854), so there is no
+// need to support messages that are too large.
+const int kBufferSize = 256 * 1024 * 1024;  // 256 MB
+
+int ListenOnIPv4(net::ServerSocket* socket, uint16_t port, bool allow_remote) {
+  std::string binding_ip = net::IPAddress::IPv4Localhost().ToString();
+  if (allow_remote)
+    binding_ip = net::IPAddress::IPv4AllZeros().ToString();
+  return socket->ListenWithAddressAndPort(binding_ip, port, 5);
+}
+
+int ListenOnIPv6(net::ServerSocket* socket, uint16_t port, bool allow_remote) {
+  std::string binding_ip = net::IPAddress::IPv6Localhost().ToString();
+  if (allow_remote)
+    binding_ip = net::IPAddress::IPv6AllZeros().ToString();
+  return socket->ListenWithAddressAndPort(binding_ip, port, 5);
+}
+
+bool RequestIsSafeToServe(const net::HttpServerRequestInfo& info,
+                          bool allow_remote,
+                          const std::vector<net::IPAddress>& whitelisted_ips) {
+  // To guard against browser-originating cross-site requests, when host header
+  // and/or origin header are present, serve only those coming from localhost
+  // or from an explicitly whitelisted ip.
+  std::string origin_header = info.GetHeaderValue("origin");
+  bool local_origin = false;
+  if (!origin_header.empty()) {
+    GURL url = GURL(origin_header);
+    local_origin = net::IsLocalhost(url);
+    if (!local_origin) {
+      if (!allow_remote) {
+        LOG(ERROR)
+            << "Remote connections not allowed; rejecting request with origin: "
+            << origin_header;
+        return false;
+      }
+      if (!whitelisted_ips.empty()) {
+        net::IPAddress address = net::IPAddress();
+        if (!ParseURLHostnameToAddress(origin_header, &address)) {
+          LOG(ERROR) << "Unable to parse origin to IPAddress: "
+                     << origin_header;
+          return false;
+        }
+        if (!base::Contains(whitelisted_ips, address)) {
+          LOG(ERROR) << "Rejecting request with origin: " << origin_header;
+          return false;
+        }
+      }
+    }
+  }
+  // TODO https://crbug.com/chromedriver/3389
+  //  When remote access is allowed and origin is not specified,
+  // we should confirm that host is current machines ip or hostname
+
+  if (local_origin || !allow_remote) {
+    // when origin is localhost host must be localhost
+    // when origin is not set, and no remote access, host must be localhost
+    std::string host_header = info.GetHeaderValue("host");
+    if (!host_header.empty()) {
+      GURL url = GURL("http://" + host_header);
+      if (!net::IsLocalhost(url)) {
+        LOG(ERROR) << "Rejecting request with host: " << host_header
+                   << ". origin is " << origin_header;
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+}  // namespace
+
+HttpServer::HttpServer(const std::string& url_base,
+                       const std::vector<net::IPAddress>& whitelisted_ips,
+                       const HttpRequestHandlerFunc& handle_request_func,
+                       base::WeakPtr<HttpHandler> handler,
+                       scoped_refptr<base::SingleThreadTaskRunner> cmd_runner)
+    : url_base_(url_base),
+      handle_request_func_(handle_request_func),
+      allow_remote_(false),
+      whitelisted_ips_(whitelisted_ips),
+      handler_(handler),
+      cmd_runner_(cmd_runner) {}
+
+int HttpServer::Start(uint16_t port, bool allow_remote, bool use_ipv4) {
+  allow_remote_ = allow_remote;
+  std::unique_ptr<net::ServerSocket> server_socket(
+      new net::TCPServerSocket(nullptr, net::NetLogSource()));
+  int status = use_ipv4 ? ListenOnIPv4(server_socket.get(), port, allow_remote)
+                        : ListenOnIPv6(server_socket.get(), port, allow_remote);
+  if (status != net::OK) {
+    VLOG(0) << "listen on " << (use_ipv4 ? "IPv4" : "IPv6")
+            << " failed with error " << net::ErrorToShortString(status);
+    return status;
+  }
+  server_ = std::make_unique<net::HttpServer>(std::move(server_socket), this);
+  net::IPEndPoint address;
+  return server_->GetLocalAddress(&address);
+}
+
+void HttpServer::OnConnect(int connection_id) {
+  server_->SetSendBufferSize(connection_id, kBufferSize);
+  server_->SetReceiveBufferSize(connection_id, kBufferSize);
+}
+
+void HttpServer::OnHttpRequest(int connection_id,
+                               const net::HttpServerRequestInfo& info) {
+  if (!RequestIsSafeToServe(info, allow_remote_, whitelisted_ips_)) {
+    server_->Send500(connection_id,
+                     "Host header or origin header is specified and is not "
+                     "whitelisted or localhost.",
+                     TRAFFIC_ANNOTATION_FOR_TESTS);
+    return;
+  }
+  handle_request_func_.Run(
+      info, base::BindRepeating(&HttpServer::OnResponse,
+                                weak_factory_.GetWeakPtr(), connection_id,
+                                !info.HasHeaderValue("connection", "close")));
+}
+
+HttpServer::~HttpServer() = default;
+
+void HttpServer::OnWebSocketRequest(int connection_id,
+                                    const net::HttpServerRequestInfo& info) {
+  cmd_runner_->PostTask(
+      FROM_HERE, base::BindOnce(&HttpHandler::OnWebSocketRequest, handler_,
+                                connection_id, info));
+
+  std::string path = info.path;
+  std::string session_id;
+
+  if (!base::StartsWith(path, url_base_, base::CompareCase::SENSITIVE)) {
+    net::HttpServerResponseInfo response(net::HTTP_BAD_REQUEST);
+    response.SetBody("invalid websocket request url path", "text/plain");
+    server_->SendResponse(connection_id, response,
+                          TRAFFIC_ANNOTATION_FOR_TESTS);
+    return;
+  }
+  path.erase(0, url_base_.length());
+
+  std::vector<std::string> path_parts =
+      base::SplitString(path, "/", base::TRIM_WHITESPACE, base::SPLIT_WANT_ALL);
+  std::vector<std::string> command_path_parts = base::SplitString(
+      kCreateWebSocketPath, "/", base::TRIM_WHITESPACE, base::SPLIT_WANT_ALL);
+
+  if (path_parts.size() != command_path_parts.size()) {
+    net::HttpServerResponseInfo response(net::HTTP_BAD_REQUEST);
+    response.SetBody("invalid websocket request url path", "text/plain");
+    server_->SendResponse(connection_id, response,
+                          TRAFFIC_ANNOTATION_FOR_TESTS);
+    return;
+  }
+
+  for (size_t i = 0; i < path_parts.size(); ++i) {
+    if (command_path_parts[i][0] == ':') {
+      std::string name = command_path_parts[i];
+      name.erase(0, 1);
+      CHECK(name.length());
+      if (name == "sessionId")
+        session_id = path_parts[i];
+    } else if (command_path_parts[i] != path_parts[i]) {
+      net::HttpServerResponseInfo response(net::HTTP_BAD_REQUEST);
+      response.SetBody("invalid websocket request url path", "text/plain");
+      server_->SendResponse(connection_id, response,
+                            TRAFFIC_ANNOTATION_FOR_TESTS);
+      return;
+    }
+  }
+
+  server_->AcceptWebSocket(connection_id, info, TRAFFIC_ANNOTATION_FOR_TESTS);
+  connection_to_session_map[connection_id] = session_id;
+}
+
+void HttpServer::OnWebSocketMessage(int connection_id, std::string data) {
+  base::Optional<base::Value> parsed_data = base::JSONReader::Read(data);
+  std::string path = url_base_ + kSendCommandFromWebSocket;
+  base::ReplaceFirstSubstringAfterOffset(
+      &path, 0, ":sessionId", connection_to_session_map[connection_id]);
+
+  net::HttpServerRequestInfo request;
+  request.method = "post";
+  request.path = path;
+  request.data = data;
+  OnHttpRequest(connection_id, request);
+}
+
+void HttpServer::OnClose(int connection_id) {
+  cmd_runner_->PostTask(FROM_HERE, base::BindOnce(&HttpHandler::OnClose,
+                                                  handler_, connection_id));
+}
+
+void HttpServer::OnResponse(
+    int connection_id,
+    bool keep_alive,
+    std::unique_ptr<net::HttpServerResponseInfo> response) {
+  if (!keep_alive)
+    response->AddHeader("Connection", "close");
+  server_->SendResponse(connection_id, *response, TRAFFIC_ANNOTATION_FOR_TESTS);
+  // Don't need to call server_->Close(), since SendResponse() will handle
+  // this for us.
+}

chrome/test/chromedriver/server/http_server.h

+// Copyright (c) 2020 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.
+
+#ifndef CHROME_TEST_CHROMEDRIVER_SERVER_HTTP_SERVER_H_
+#define CHROME_TEST_CHROMEDRIVER_SERVER_HTTP_SERVER_H_
+
+#include "base/json/json_reader.h"
+#include "chrome/test/chromedriver/server/http_handler.h"
+#include "net/base/url_util.h"
+#include "net/server/http_server.h"
+#include "net/server/http_server_request_info.h"
+#include "net/server/http_server_response_info.h"
+#include "net/socket/tcp_server_socket.h"
+#include "net/traffic_annotation/network_traffic_annotation_test_helper.h"
+
+typedef base::RepeatingCallback<void(const net::HttpServerRequestInfo&,
+                                     const HttpResponseSenderFunc&)>
+    HttpRequestHandlerFunc;
+
+class HttpServer : public net::HttpServer::Delegate {
+ public:
+  explicit HttpServer(const std::string& url_base,
+                      const std::vector<net::IPAddress>& whitelisted_ips,
+                      const HttpRequestHandlerFunc& handle_request_func,
+                      base::WeakPtr<HttpHandler> handler,
+                      scoped_refptr<base::SingleThreadTaskRunner> cmd_runner);
+
+  ~HttpServer() override;
+
+  int Start(uint16_t port, bool allow_remote, bool use_ipv4);
+
+  // Overridden from net::HttpServer::Delegate:
+  void OnConnect(int connection_id) override;
+
+  void OnHttpRequest(int connection_id,
+                     const net::HttpServerRequestInfo& info) override;
+
+  void OnWebSocketRequest(int connection_id,
+                          const net::HttpServerRequestInfo& info) override;
+
+  void OnWebSocketMessage(int connection_id, std::string data) override;
+
+  void OnClose(int connection_id) override;
+
+ private:
+  void OnResponse(int connection_id,
+                  bool keep_alive,
+                  std::unique_ptr<net::HttpServerResponseInfo> response);
+  const std::string url_base_;
+  HttpRequestHandlerFunc handle_request_func_;
+  std::unique_ptr<net::HttpServer> server_;
+  std::map<int, std::string> connection_to_session_map;
+  bool allow_remote_;
+  const std::vector<net::IPAddress> whitelisted_ips_;
+  base::WeakPtr<HttpHandler> handler_;
+  scoped_refptr<base::SingleThreadTaskRunner> cmd_runner_;
+  base::WeakPtrFactory<HttpServer> weak_factory_{this};  // Should be last.
+};
+
+#endif  // CHROME_TEST_CHROMEDRIVER_SERVER_HTTP_SERVER_H_

chrome/test/chromedriver/session_commands.cc

@@ -77,14 +77,24 @@ Status EvaluateScriptAndIgnoreResult(Session* session,
   return web_view->EvaluateScript(frame_id, expression, awaitPromise, &result);
 }
 
+void InitSessionForWebSocketConnection(SessionConnectionMap* session_map,
+                                       std::string session_id) {
+  session_map->insert({session_id, -1});
+}
+
 }  // namespace
 
-InitSessionParams::InitSessionParams(network::mojom::URLLoaderFactory* factory,
-                                     const SyncWebSocketFactory& socket_factory,
-                                     DeviceManager* device_manager)
+InitSessionParams::InitSessionParams(
+    network::mojom::URLLoaderFactory* factory,
+    const SyncWebSocketFactory& socket_factory,
+    DeviceManager* device_manager,
+    const scoped_refptr<base::SingleThreadTaskRunner> cmd_task_runner,
+    SessionConnectionMap* session_map)
     : url_loader_factory(factory),
       socket_factory(socket_factory),
-      device_manager(device_manager) {}
+      device_manager(device_manager),
+      cmd_task_runner(cmd_task_runner),
+      session_map(session_map) {}
 
 InitSessionParams::InitSessionParams(const InitSessionParams& other) = default;
 
@@ -595,6 +605,12 @@ Status ExecuteInitSession(const InitSessionParams& bound_params,
     session->quit = true;
     if (session->chrome != NULL)
       session->chrome->Quit();
+  } else {
+    // TODO only do this when WebSocketUrl capability is specified
+    // https://bugs.chromium.org/p/chromedriver/issues/detail?id=3588
+    bound_params.cmd_task_runner->PostTask(
+        FROM_HERE, base::BindOnce(&InitSessionForWebSocketConnection,
+                                  bound_params.session_map, session->id));
   }
   return status;
 }

chrome/test/chromedriver/session_commands.h

@@ -11,6 +11,7 @@
 #include "base/callback_forward.h"
 #include "chrome/test/chromedriver/command.h"
 #include "chrome/test/chromedriver/net/sync_websocket_factory.h"
+#include "chrome/test/chromedriver/session_connection_map.h"
 #include "services/network/public/mojom/url_loader_factory.mojom.h"
 
 namespace base {
@@ -24,15 +25,20 @@ struct Session;
 class Status;
 
 struct InitSessionParams {
-  InitSessionParams(network::mojom::URLLoaderFactory* factory,
-                    const SyncWebSocketFactory& socket_factory,
-                    DeviceManager* device_manager);
+  InitSessionParams(
+      network::mojom::URLLoaderFactory* factory,
+      const SyncWebSocketFactory& socket_factory,
+      DeviceManager* device_manager,
+      const scoped_refptr<base::SingleThreadTaskRunner> cmd_task_runner,
+      SessionConnectionMap* session_map);
   InitSessionParams(const InitSessionParams& other);
   ~InitSessionParams();
 
   network::mojom::URLLoaderFactory* url_loader_factory;
   SyncWebSocketFactory socket_factory;
   DeviceManager* device_manager;
+  scoped_refptr<base::SingleThreadTaskRunner> cmd_task_runner;
+  SessionConnectionMap* session_map;
 };
 
 bool GetW3CSetting(const base::DictionaryValue& params);

chrome/test/chromedriver/session_connection_map.h

+// Copyright 2020 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.
+
+#ifndef CHROME_TEST_CHROMEDRIVER_SESSION_CONNECTION_MAP_H_
+#define CHROME_TEST_CHROMEDRIVER_SESSION_CONNECTION_MAP_H_
+
+#include <string>
+#include <unordered_map>
+
+using SessionConnectionMap = std::unordered_map<std::string, int>;
+
+#endif  // CHROME_TEST_CHROMEDRIVER_SESSION_CONNECTION_MAP_H_