fix: make webRequest work with WebSocket (#22040)

* fix: web request support proxying websocket

* fix: make tests work

* chore: do not use api:: code outside api/ folder

* chore: do not create proxy when no listener

* test: use separate session to avoid conflicts

* chore: address review

shell/browser/api/electron_api_web_request.h

@@ -13,7 +13,7 @@
 #include "gin/arguments.h"
 #include "gin/handle.h"
 #include "gin/wrappable.h"
-#include "shell/browser/net/proxying_url_loader_factory.h"
+#include "shell/browser/net/web_request_api_interface.h"
 
 namespace content {
 class BrowserContext;
@@ -49,10 +49,6 @@ class WebRequest : public gin::Wrappable<WebRequest>, public WebRequestAPI {
       v8::Isolate* isolate) override;
   const char* GetTypeName() override;
 
- private:
-  WebRequest(v8::Isolate* isolate, content::BrowserContext* browser_context);
-  ~WebRequest() override;
-
   // WebRequestAPI:
   bool HasListener() const override;
   int OnBeforeRequest(extensions::WebRequestInfo* info,
@@ -86,6 +82,10 @@ class WebRequest : public gin::Wrappable<WebRequest>, public WebRequestAPI {
                    int net_error) override;
   void OnRequestWillBeDestroyed(extensions::WebRequestInfo* info) override;
 
+ private:
+  WebRequest(v8::Isolate* isolate, content::BrowserContext* browser_context);
+  ~WebRequest() override;
+
   enum SimpleEvent {
     kOnSendHeaders,
     kOnBeforeRedirect,

shell/browser/electron_browser_client.cc

@@ -75,6 +75,7 @@
 #include "shell/browser/net/network_context_service.h"
 #include "shell/browser/net/network_context_service_factory.h"
 #include "shell/browser/net/proxying_url_loader_factory.h"
+#include "shell/browser/net/proxying_websocket.h"
 #include "shell/browser/net/system_network_context_manager.h"
 #include "shell/browser/network_hints_handler_impl.h"
 #include "shell/browser/notifications/notification_presenter.h"
@@ -1167,6 +1168,43 @@ void ElectronBrowserClient::RegisterNonNetworkSubresourceURLLoaderFactories(
   }
 }
 
+bool ElectronBrowserClient::WillInterceptWebSocket(
+    content::RenderFrameHost* frame) {
+  if (!frame)
+    return false;
+
+  v8::Isolate* isolate = v8::Isolate::GetCurrent();
+  auto* browser_context = frame->GetProcess()->GetBrowserContext();
+  auto web_request = api::WebRequest::FromOrCreate(isolate, browser_context);
+
+  // NOTE: Some unit test environments do not initialize
+  // BrowserContextKeyedAPI factories for e.g. WebRequest.
+  if (!web_request.get())
+    return false;
+
+  return web_request->HasListener();
+}
+
+void ElectronBrowserClient::CreateWebSocket(
+    content::RenderFrameHost* frame,
+    WebSocketFactory factory,
+    const GURL& url,
+    const net::SiteForCookies& site_for_cookies,
+    const base::Optional<std::string>& user_agent,
+    mojo::PendingRemote<network::mojom::WebSocketHandshakeClient>
+        handshake_client) {
+  v8::Isolate* isolate = v8::Isolate::GetCurrent();
+  auto* browser_context = frame->GetProcess()->GetBrowserContext();
+  auto web_request = api::WebRequest::FromOrCreate(isolate, browser_context);
+  DCHECK(web_request.get());
+  ProxyingWebSocket::StartProxying(
+      web_request.get(), std::move(factory), url,
+      site_for_cookies.RepresentativeUrl(), user_agent,
+      std::move(handshake_client), true, frame->GetProcess()->GetID(),
+      frame->GetRoutingID(), frame->GetLastCommittedOrigin(), browser_context,
+      &next_id_);
+}
+
 bool ElectronBrowserClient::WillCreateURLLoaderFactory(
     content::BrowserContext* browser_context,
     content::RenderFrameHost* frame_host,
@@ -1209,7 +1247,7 @@ bool ElectronBrowserClient::WillCreateURLLoaderFactory(
 
   new ProxyingURLLoaderFactory(
       web_request.get(), protocol->intercept_handlers(), browser_context,
-      render_process_id, std::move(navigation_ui_data),
+      render_process_id, &next_id_, std::move(navigation_ui_data),
       std::move(navigation_id), std::move(proxied_receiver),
       std::move(target_factory_remote), std::move(header_client_receiver),
       type);

shell/browser/electron_browser_client.h

@@ -173,6 +173,15 @@ class ElectronBrowserClient : public content::ContentBrowserClient,
       int render_process_id,
       int render_frame_id,
       NonNetworkURLLoaderFactoryMap* factories) override;
+  void CreateWebSocket(
+      content::RenderFrameHost* frame,
+      WebSocketFactory factory,
+      const GURL& url,
+      const net::SiteForCookies& site_for_cookies,
+      const base::Optional<std::string>& user_agent,
+      mojo::PendingRemote<network::mojom::WebSocketHandshakeClient>
+          handshake_client) override;
+  bool WillInterceptWebSocket(content::RenderFrameHost*) override;
   bool WillCreateURLLoaderFactory(
       content::BrowserContext* browser_context,
       content::RenderFrameHost* frame,
@@ -290,6 +299,10 @@ class ElectronBrowserClient : public content::ContentBrowserClient,
 
   bool disable_process_restart_tricks_ = false;
 
+  // Simple shared ID generator, used by ProxyingURLLoaderFactory and
+  // ProxyingWebSocket classes.
+  uint64_t next_id_ = 0;
+
   DISALLOW_COPY_AND_ASSIGN(ElectronBrowserClient);
 };
 

shell/browser/net/proxying_url_loader_factory.cc

@@ -20,13 +20,6 @@
 #include "shell/common/options_switches.h"
 
 namespace electron {
-
-namespace {
-
-int64_t g_request_id = 0;
-
-}  // namespace
-
 ProxyingURLLoaderFactory::InProgressRequest::FollowRedirectParams::
     FollowRedirectParams() = default;
 ProxyingURLLoaderFactory::InProgressRequest::FollowRedirectParams::
@@ -677,6 +670,7 @@ ProxyingURLLoaderFactory::ProxyingURLLoaderFactory(
     const HandlersMap& intercepted_handlers,
     content::BrowserContext* browser_context,
     int render_process_id,
+    uint64_t* request_id_generator,
     std::unique_ptr<extensions::ExtensionNavigationUIData> navigation_ui_data,
     base::Optional<int64_t> navigation_id,
     network::mojom::URLLoaderFactoryRequest loader_request,
@@ -688,6 +682,7 @@ ProxyingURLLoaderFactory::ProxyingURLLoaderFactory(
       intercepted_handlers_(intercepted_handlers),
       browser_context_(browser_context),
       render_process_id_(render_process_id),
+      request_id_generator_(request_id_generator),
       navigation_ui_data_(std::move(navigation_ui_data)),
       navigation_id_(std::move(navigation_id)),
       loader_factory_type_(loader_factory_type) {
@@ -765,7 +760,7 @@ void ProxyingURLLoaderFactory::CreateLoaderAndStart(
   // per-BrowserContext so extensions can make sense of it.  Note that
   // |network_service_request_id_| by contrast is not necessarily unique, so we
   // don't use it for identity here.
-  const uint64_t web_request_id = ++g_request_id;
+  const uint64_t web_request_id = ++(*request_id_generator_);
 
   // Notes: Chromium assumes that requests with zero-ID would never use the
   // "extraHeaders" code path, however in Electron requests started from

shell/browser/net/proxying_url_loader_factory.h

@@ -23,53 +23,12 @@
 #include "services/network/public/mojom/network_context.mojom.h"
 #include "services/network/public/mojom/url_loader.mojom.h"
 #include "services/network/public/mojom/url_response_head.mojom.h"
+#include "shell/browser/api/electron_api_web_request.h"
 #include "shell/browser/net/electron_url_loader_factory.h"
+#include "shell/browser/net/web_request_api_interface.h"
 
 namespace electron {
 
-// Defines the interface for WebRequest API, implemented by api::WebRequest.
-class WebRequestAPI {
- public:
-  virtual ~WebRequestAPI() {}
-
-  using BeforeSendHeadersCallback =
-      base::OnceCallback<void(const std::set<std::string>& removed_headers,
-                              const std::set<std::string>& set_headers,
-                              int error_code)>;
-
-  virtual bool HasListener() const = 0;
-  virtual int OnBeforeRequest(extensions::WebRequestInfo* info,
-                              const network::ResourceRequest& request,
-                              net::CompletionOnceCallback callback,
-                              GURL* new_url) = 0;
-  virtual int OnBeforeSendHeaders(extensions::WebRequestInfo* info,
-                                  const network::ResourceRequest& request,
-                                  BeforeSendHeadersCallback callback,
-                                  net::HttpRequestHeaders* headers) = 0;
-  virtual int OnHeadersReceived(
-      extensions::WebRequestInfo* info,
-      const network::ResourceRequest& request,
-      net::CompletionOnceCallback callback,
-      const net::HttpResponseHeaders* original_response_headers,
-      scoped_refptr<net::HttpResponseHeaders>* override_response_headers,
-      GURL* allowed_unsafe_redirect_url) = 0;
-  virtual void OnSendHeaders(extensions::WebRequestInfo* info,
-                             const network::ResourceRequest& request,
-                             const net::HttpRequestHeaders& headers) = 0;
-  virtual void OnBeforeRedirect(extensions::WebRequestInfo* info,
-                                const network::ResourceRequest& request,
-                                const GURL& new_location) = 0;
-  virtual void OnResponseStarted(extensions::WebRequestInfo* info,
-                                 const network::ResourceRequest& request) = 0;
-  virtual void OnErrorOccurred(extensions::WebRequestInfo* info,
-                               const network::ResourceRequest& request,
-                               int net_error) = 0;
-  virtual void OnCompleted(extensions::WebRequestInfo* info,
-                           const network::ResourceRequest& request,
-                           int net_error) = 0;
-  virtual void OnRequestWillBeDestroyed(extensions::WebRequestInfo* info) = 0;
-};
-
 // This class is responsible for following tasks when NetworkService is enabled:
 // 1. handling intercepted protocols;
 // 2. implementing webRequest module;
@@ -210,6 +169,7 @@ class ProxyingURLLoaderFactory
       const HandlersMap& intercepted_handlers,
       content::BrowserContext* browser_context,
       int render_process_id,
+      uint64_t* request_id_generator,
       std::unique_ptr<extensions::ExtensionNavigationUIData> navigation_ui_data,
       base::Optional<int64_t> navigation_id,
       network::mojom::URLLoaderFactoryRequest loader_request,
@@ -269,6 +229,7 @@ class ProxyingURLLoaderFactory
 
   content::BrowserContext* const browser_context_;
   const int render_process_id_;
+  uint64_t* request_id_generator_;  // managed by AtomBrowserClient
   std::unique_ptr<extensions::ExtensionNavigationUIData> navigation_ui_data_;
   base::Optional<int64_t> navigation_id_;
   mojo::ReceiverSet<network::mojom::URLLoaderFactory> proxy_receivers_;

shell/browser/net/proxying_websocket.cc

@@ -0,0 +1,456 @@
+// 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.
+
+#include "electron/shell/browser/net/proxying_websocket.h"
+
+#include <utility>
+
+#include "base/bind.h"
+#include "base/strings/string_util.h"
+#include "base/strings/stringprintf.h"
+#include "components/keyed_service/content/browser_context_keyed_service_shutdown_notifier_factory.h"
+#include "content/public/browser/browser_thread.h"
+#include "extensions/browser/extension_navigation_ui_data.h"
+#include "net/base/ip_endpoint.h"
+#include "net/http/http_util.h"
+
+namespace electron {
+
+ProxyingWebSocket::ProxyingWebSocket(
+    WebRequestAPI* web_request_api,
+    WebSocketFactory factory,
+    const network::ResourceRequest& request,
+    mojo::PendingRemote<network::mojom::WebSocketHandshakeClient>
+        handshake_client,
+    bool has_extra_headers,
+    int process_id,
+    int render_frame_id,
+    content::BrowserContext* browser_context,
+    uint64_t* request_id_generator)
+    : web_request_api_(web_request_api),
+      request_(request),
+      factory_(std::move(factory)),
+      forwarding_handshake_client_(std::move(handshake_client)),
+      request_headers_(request.headers),
+      response_(network::mojom::URLResponseHead::New()),
+      has_extra_headers_(has_extra_headers),
+      info_(extensions::WebRequestInfoInitParams(
+          ++(*request_id_generator),
+          process_id,
+          render_frame_id,
+          nullptr,
+          MSG_ROUTING_NONE,
+          request,
+          /*is_download=*/false,
+          /*is_async=*/true,
+          /*is_service_worker_script=*/false,
+          /*navigation_id=*/base::nullopt)) {}
+
+ProxyingWebSocket::~ProxyingWebSocket() {
+  if (on_before_send_headers_callback_) {
+    std::move(on_before_send_headers_callback_)
+        .Run(net::ERR_ABORTED, base::nullopt);
+  }
+  if (on_headers_received_callback_) {
+    std::move(on_headers_received_callback_)
+        .Run(net::ERR_ABORTED, base::nullopt, GURL());
+  }
+}
+
+void ProxyingWebSocket::Start() {
+  // If the header client will be used, we start the request immediately, and
+  // OnBeforeSendHeaders and OnSendHeaders will be handled there. Otherwise,
+  // send these events before the request starts.
+  base::RepeatingCallback<void(int)> continuation;
+  if (has_extra_headers_) {
+    continuation = base::BindRepeating(
+        &ProxyingWebSocket::ContinueToStartRequest, weak_factory_.GetWeakPtr());
+  } else {
+    continuation =
+        base::BindRepeating(&ProxyingWebSocket::OnBeforeRequestComplete,
+                            weak_factory_.GetWeakPtr());
+  }
+
+  int result = web_request_api_->OnBeforeRequest(&info_, request_, continuation,
+                                                 &redirect_url_);
+
+  if (result == net::ERR_BLOCKED_BY_CLIENT) {
+    OnError(result);
+    return;
+  }
+
+  if (result == net::ERR_IO_PENDING) {
+    return;
+  }
+
+  DCHECK_EQ(net::OK, result);
+  continuation.Run(net::OK);
+}
+
+void ProxyingWebSocket::OnOpeningHandshakeStarted(
+    network::mojom::WebSocketHandshakeRequestPtr request) {
+  DCHECK(forwarding_handshake_client_);
+  forwarding_handshake_client_->OnOpeningHandshakeStarted(std::move(request));
+}
+
+void ProxyingWebSocket::ContinueToHeadersReceived() {
+  auto continuation =
+      base::BindRepeating(&ProxyingWebSocket::OnHeadersReceivedComplete,
+                          weak_factory_.GetWeakPtr());
+  info_.AddResponseInfoFromResourceResponse(*response_);
+  int result = web_request_api_->OnHeadersReceived(
+      &info_, request_, continuation, response_->headers.get(),
+      &override_headers_, &redirect_url_);
+
+  if (result == net::ERR_BLOCKED_BY_CLIENT) {
+    OnError(result);
+    return;
+  }
+
+  PauseIncomingMethodCallProcessing();
+  if (result == net::ERR_IO_PENDING)
+    return;
+
+  DCHECK_EQ(net::OK, result);
+  OnHeadersReceivedComplete(net::OK);
+}
+
+void ProxyingWebSocket::OnConnectionEstablished(
+    mojo::PendingRemote<network::mojom::WebSocket> websocket,
+    mojo::PendingReceiver<network::mojom::WebSocketClient> client_receiver,
+    network::mojom::WebSocketHandshakeResponsePtr response,
+    mojo::ScopedDataPipeConsumerHandle readable) {
+  DCHECK(forwarding_handshake_client_);
+  DCHECK(!is_done_);
+  is_done_ = true;
+  websocket_ = std::move(websocket);
+  client_receiver_ = std::move(client_receiver);
+  handshake_response_ = std::move(response);
+  readable_ = std::move(readable);
+
+  response_->remote_endpoint = handshake_response_->remote_endpoint;
+
+  // response_->headers will be set in OnBeforeSendHeaders if
+  // |receiver_as_header_client_| is set.
+  if (receiver_as_header_client_.is_bound()) {
+    ContinueToCompleted();
+    return;
+  }
+
+  response_->headers =
+      base::MakeRefCounted<net::HttpResponseHeaders>(base::StringPrintf(
+          "HTTP/%d.%d %d %s", handshake_response_->http_version.major_value(),
+          handshake_response_->http_version.minor_value(),
+          handshake_response_->status_code,
+          handshake_response_->status_text.c_str()));
+  for (const auto& header : handshake_response_->headers)
+    response_->headers->AddHeader(header->name + ": " + header->value);
+
+  ContinueToHeadersReceived();
+}
+
+void ProxyingWebSocket::ContinueToCompleted() {
+  DCHECK(forwarding_handshake_client_);
+  DCHECK(is_done_);
+  web_request_api_->OnCompleted(&info_, request_, net::ERR_WS_UPGRADE);
+  forwarding_handshake_client_->OnConnectionEstablished(
+      std::move(websocket_), std::move(client_receiver_),
+      std::move(handshake_response_), std::move(readable_));
+
+  // Deletes |this|.
+  delete this;
+}
+
+void ProxyingWebSocket::OnAuthRequired(
+    const net::AuthChallengeInfo& auth_info,
+    const scoped_refptr<net::HttpResponseHeaders>& headers,
+    const net::IPEndPoint& remote_endpoint,
+    OnAuthRequiredCallback callback) {
+  if (!callback) {
+    OnError(net::ERR_FAILED);
+    return;
+  }
+
+  response_->headers = headers;
+  response_->remote_endpoint = remote_endpoint;
+  auth_required_callback_ = std::move(callback);
+
+  auto continuation =
+      base::BindRepeating(&ProxyingWebSocket::OnHeadersReceivedCompleteForAuth,
+                          weak_factory_.GetWeakPtr(), auth_info);
+  info_.AddResponseInfoFromResourceResponse(*response_);
+  int result = web_request_api_->OnHeadersReceived(
+      &info_, request_, continuation, response_->headers.get(),
+      &override_headers_, &redirect_url_);
+
+  if (result == net::ERR_BLOCKED_BY_CLIENT) {
+    OnError(result);
+    return;
+  }
+
+  PauseIncomingMethodCallProcessing();
+  if (result == net::ERR_IO_PENDING)
+    return;
+
+  DCHECK_EQ(net::OK, result);
+  OnHeadersReceivedCompleteForAuth(auth_info, net::OK);
+}
+
+void ProxyingWebSocket::OnBeforeSendHeaders(
+    const net::HttpRequestHeaders& headers,
+    OnBeforeSendHeadersCallback callback) {
+  DCHECK(receiver_as_header_client_.is_bound());
+
+  request_headers_ = headers;
+  on_before_send_headers_callback_ = std::move(callback);
+  OnBeforeRequestComplete(net::OK);
+}
+
+void ProxyingWebSocket::OnHeadersReceived(const std::string& headers,
+                                          const net::IPEndPoint& endpoint,
+                                          OnHeadersReceivedCallback callback) {
+  DCHECK(receiver_as_header_client_.is_bound());
+
+  on_headers_received_callback_ = std::move(callback);
+  response_->headers = base::MakeRefCounted<net::HttpResponseHeaders>(headers);
+
+  ContinueToHeadersReceived();
+}
+
+void ProxyingWebSocket::StartProxying(
+    WebRequestAPI* web_request_api,
+    WebSocketFactory factory,
+    const GURL& url,
+    const GURL& site_for_cookies,
+    const base::Optional<std::string>& user_agent,
+    mojo::PendingRemote<network::mojom::WebSocketHandshakeClient>
+        handshake_client,
+    bool has_extra_headers,
+    int process_id,
+    int render_frame_id,
+    const url::Origin& origin,
+    content::BrowserContext* browser_context,
+    uint64_t* request_id_generator) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+  network::ResourceRequest request;
+  request.url = url;
+  request.site_for_cookies = net::SiteForCookies::FromUrl(site_for_cookies);
+  if (user_agent) {
+    request.headers.SetHeader(net::HttpRequestHeaders::kUserAgent, *user_agent);
+  }
+  request.request_initiator = origin;
+
+  auto* proxy = new ProxyingWebSocket(
+      web_request_api, std::move(factory), request, std::move(handshake_client),
+      has_extra_headers, process_id, render_frame_id, browser_context,
+      request_id_generator);
+  proxy->Start();
+}
+
+void ProxyingWebSocket::OnBeforeRequestComplete(int error_code) {
+  DCHECK(receiver_as_header_client_.is_bound() ||
+         !receiver_as_handshake_client_.is_bound());
+  DCHECK(info_.url.SchemeIsWSOrWSS());
+  if (error_code != net::OK) {
+    OnError(error_code);
+    return;
+  }
+
+  auto continuation =
+      base::BindRepeating(&ProxyingWebSocket::OnBeforeSendHeadersComplete,
+                          weak_factory_.GetWeakPtr());
+
+  info_.AddResponseInfoFromResourceResponse(*response_);
+  int result = web_request_api_->OnBeforeSendHeaders(
+      &info_, request_, continuation, &request_headers_);
+
+  if (result == net::ERR_BLOCKED_BY_CLIENT) {
+    OnError(result);
+    return;
+  }
+
+  if (result == net::ERR_IO_PENDING)
+    return;
+
+  DCHECK_EQ(net::OK, result);
+  OnBeforeSendHeadersComplete(std::set<std::string>(), std::set<std::string>(),
+                              net::OK);
+}
+
+void ProxyingWebSocket::OnBeforeSendHeadersComplete(
+    const std::set<std::string>& removed_headers,
+    const std::set<std::string>& set_headers,
+    int error_code) {
+  DCHECK(receiver_as_header_client_.is_bound() ||
+         !receiver_as_handshake_client_.is_bound());
+  if (error_code != net::OK) {
+    OnError(error_code);
+    return;
+  }
+
+  if (receiver_as_header_client_.is_bound()) {
+    CHECK(on_before_send_headers_callback_);
+    std::move(on_before_send_headers_callback_)
+        .Run(error_code, request_headers_);
+  }
+
+  info_.AddResponseInfoFromResourceResponse(*response_);
+  web_request_api_->OnSendHeaders(&info_, request_, request_headers_);
+
+  if (!receiver_as_header_client_.is_bound())
+    ContinueToStartRequest(net::OK);
+}
+
+void ProxyingWebSocket::ContinueToStartRequest(int error_code) {
+  if (error_code != net::OK) {
+    OnError(error_code);
+    return;
+  }
+
+  base::flat_set<std::string> used_header_names;
+  std::vector<network::mojom::HttpHeaderPtr> additional_headers;
+  for (net::HttpRequestHeaders::Iterator it(request_headers_); it.GetNext();) {
+    additional_headers.push_back(
+        network::mojom::HttpHeader::New(it.name(), it.value()));
+    used_header_names.insert(base::ToLowerASCII(it.name()));
+  }
+  for (const auto& header : additional_headers_) {
+    if (!used_header_names.contains(base::ToLowerASCII(header->name))) {
+      additional_headers.push_back(
+          network::mojom::HttpHeader::New(header->name, header->value));
+    }
+  }
+
+  mojo::PendingRemote<network::mojom::TrustedHeaderClient>
+      trusted_header_client = mojo::NullRemote();
+  if (has_extra_headers_) {
+    trusted_header_client =
+        receiver_as_header_client_.BindNewPipeAndPassRemote();
+  }
+
+  std::move(factory_).Run(
+      info_.url, std::move(additional_headers),
+      receiver_as_handshake_client_.BindNewPipeAndPassRemote(),
+      receiver_as_auth_handler_.BindNewPipeAndPassRemote(),
+      std::move(trusted_header_client));
+
+  // Here we detect mojo connection errors on |receiver_as_handshake_client_|.
+  // See also CreateWebSocket in
+  // //network/services/public/mojom/network_context.mojom.
+  receiver_as_handshake_client_.set_disconnect_with_reason_handler(
+      base::BindOnce(&ProxyingWebSocket::OnMojoConnectionErrorWithCustomReason,
+                     base::Unretained(this)));
+  forwarding_handshake_client_.set_disconnect_handler(base::BindOnce(
+      &ProxyingWebSocket::OnMojoConnectionError, base::Unretained(this)));
+}
+
+void ProxyingWebSocket::OnHeadersReceivedComplete(int error_code) {
+  if (error_code != net::OK) {
+    OnError(error_code);
+    return;
+  }
+
+  if (on_headers_received_callback_) {
+    base::Optional<std::string> headers;
+    if (override_headers_)
+      headers = override_headers_->raw_headers();
+    std::move(on_headers_received_callback_)
+        .Run(net::OK, headers, base::nullopt);
+  }
+
+  if (override_headers_) {
+    response_->headers = override_headers_;
+    override_headers_ = nullptr;
+  }
+
+  ResumeIncomingMethodCallProcessing();
+  info_.AddResponseInfoFromResourceResponse(*response_);
+  web_request_api_->OnResponseStarted(&info_, request_);
+
+  if (!receiver_as_header_client_.is_bound())
+    ContinueToCompleted();
+}
+
+void ProxyingWebSocket::OnAuthRequiredComplete(
+    extensions::ExtensionWebRequestEventRouter::AuthRequiredResponse rv) {
+  CHECK(auth_required_callback_);
+  ResumeIncomingMethodCallProcessing();
+  switch (rv) {
+    case extensions::ExtensionWebRequestEventRouter::AuthRequiredResponse::
+        AUTH_REQUIRED_RESPONSE_NO_ACTION:
+    case extensions::ExtensionWebRequestEventRouter::AuthRequiredResponse::
+        AUTH_REQUIRED_RESPONSE_CANCEL_AUTH:
+      std::move(auth_required_callback_).Run(base::nullopt);
+      break;
+
+    case extensions::ExtensionWebRequestEventRouter::AuthRequiredResponse::
+        AUTH_REQUIRED_RESPONSE_SET_AUTH:
+      std::move(auth_required_callback_).Run(auth_credentials_);
+      break;
+    case extensions::ExtensionWebRequestEventRouter::AuthRequiredResponse::
+        AUTH_REQUIRED_RESPONSE_IO_PENDING:
+      NOTREACHED();
+      break;
+  }
+}
+
+void ProxyingWebSocket::OnHeadersReceivedCompleteForAuth(
+    const net::AuthChallengeInfo& auth_info,
+    int rv) {
+  if (rv != net::OK) {
+    OnError(rv);
+    return;
+  }
+  ResumeIncomingMethodCallProcessing();
+  info_.AddResponseInfoFromResourceResponse(*response_);
+
+  auto continuation = base::BindRepeating(
+      &ProxyingWebSocket::OnAuthRequiredComplete, weak_factory_.GetWeakPtr());
+  auto auth_rv = extensions::ExtensionWebRequestEventRouter::
+      AuthRequiredResponse::AUTH_REQUIRED_RESPONSE_IO_PENDING;
+  PauseIncomingMethodCallProcessing();
+
+  OnAuthRequiredComplete(auth_rv);
+}
+
+void ProxyingWebSocket::PauseIncomingMethodCallProcessing() {
+  receiver_as_handshake_client_.Pause();
+  receiver_as_auth_handler_.Pause();
+  if (receiver_as_header_client_.is_bound())
+    receiver_as_header_client_.Pause();
+}
+
+void ProxyingWebSocket::ResumeIncomingMethodCallProcessing() {
+  receiver_as_handshake_client_.Resume();
+  receiver_as_auth_handler_.Resume();
+  if (receiver_as_header_client_.is_bound())
+    receiver_as_header_client_.Resume();
+}
+
+void ProxyingWebSocket::OnError(int error_code) {
+  if (!is_done_) {
+    is_done_ = true;
+    web_request_api_->OnErrorOccurred(&info_, request_, error_code);
+  }
+
+  // Deletes |this|.
+  delete this;
+}
+
+void ProxyingWebSocket::OnMojoConnectionErrorWithCustomReason(
+    uint32_t custom_reason,
+    const std::string& description) {
+  // Here we want to nofiy the custom reason to the client, which is why
+  // we reset |forwarding_handshake_client_| manually.
+  forwarding_handshake_client_.ResetWithReason(custom_reason, description);
+  OnError(net::ERR_FAILED);
+  // Deletes |this|.
+}
+
+void ProxyingWebSocket::OnMojoConnectionError() {
+  OnError(net::ERR_FAILED);
+  // Deletes |this|.
+}
+
+}  // namespace electron

shell/browser/net/proxying_websocket.h

@@ -0,0 +1,165 @@
+// 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 SHELL_BROWSER_NET_PROXYING_WEBSOCKET_H_
+#define SHELL_BROWSER_NET_PROXYING_WEBSOCKET_H_
+
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "base/optional.h"
+#include "components/keyed_service/core/keyed_service_shutdown_notifier.h"
+#include "extensions/browser/api/web_request/web_request_api.h"
+#include "extensions/browser/api/web_request/web_request_info.h"
+#include "mojo/public/cpp/bindings/pending_receiver.h"
+#include "mojo/public/cpp/bindings/receiver.h"
+#include "mojo/public/cpp/bindings/remote.h"
+#include "services/network/public/cpp/resource_request.h"
+#include "services/network/public/mojom/network_context.mojom.h"
+#include "services/network/public/mojom/websocket.mojom.h"
+#include "shell/browser/net/web_request_api_interface.h"
+#include "url/gurl.h"
+#include "url/origin.h"
+
+namespace electron {
+
+// A ProxyingWebSocket proxies a WebSocket connection and dispatches
+// WebRequest API events.
+//
+// The code is referenced from the
+// extensions::WebRequestProxyingWebSocket class.
+class ProxyingWebSocket : public network::mojom::WebSocketHandshakeClient,
+                          public network::mojom::AuthenticationHandler,
+                          public network::mojom::TrustedHeaderClient {
+ public:
+  using WebSocketFactory = content::ContentBrowserClient::WebSocketFactory;
+
+  ProxyingWebSocket(
+      WebRequestAPI* web_request_api,
+      WebSocketFactory factory,
+      const network::ResourceRequest& request,
+      mojo::PendingRemote<network::mojom::WebSocketHandshakeClient>
+          handshake_client,
+      bool has_extra_headers,
+      int process_id,
+      int render_frame_id,
+      content::BrowserContext* browser_context,
+      uint64_t* request_id_generator);
+  ~ProxyingWebSocket() override;
+
+  void Start();
+
+  // network::mojom::WebSocketHandshakeClient methods:
+  void OnOpeningHandshakeStarted(
+      network::mojom::WebSocketHandshakeRequestPtr request) override;
+  void OnConnectionEstablished(
+      mojo::PendingRemote<network::mojom::WebSocket> websocket,
+      mojo::PendingReceiver<network::mojom::WebSocketClient> client_receiver,
+      network::mojom::WebSocketHandshakeResponsePtr response,
+      mojo::ScopedDataPipeConsumerHandle readable) override;
+
+  // network::mojom::AuthenticationHandler method:
+  void OnAuthRequired(const net::AuthChallengeInfo& auth_info,
+                      const scoped_refptr<net::HttpResponseHeaders>& headers,
+                      const net::IPEndPoint& remote_endpoint,
+                      OnAuthRequiredCallback callback) override;
+
+  // network::mojom::TrustedHeaderClient methods:
+  void OnBeforeSendHeaders(const net::HttpRequestHeaders& headers,
+                           OnBeforeSendHeadersCallback callback) override;
+  void OnHeadersReceived(const std::string& headers,
+                         const net::IPEndPoint& endpoint,
+                         OnHeadersReceivedCallback callback) override;
+
+  static void StartProxying(
+      WebRequestAPI* web_request_api,
+      WebSocketFactory factory,
+      const GURL& url,
+      const GURL& site_for_cookies,
+      const base::Optional<std::string>& user_agent,
+      mojo::PendingRemote<network::mojom::WebSocketHandshakeClient>
+          handshake_client,
+      bool has_extra_headers,
+      int process_id,
+      int render_frame_id,
+      const url::Origin& origin,
+      content::BrowserContext* browser_context,
+      uint64_t* request_id_generator);
+
+  WebRequestAPI* web_request_api() { return web_request_api_; }
+
+ private:
+  void OnBeforeRequestComplete(int error_code);
+  void OnBeforeSendHeadersComplete(const std::set<std::string>& removed_headers,
+                                   const std::set<std::string>& set_headers,
+                                   int error_code);
+  void ContinueToStartRequest(int error_code);
+  void OnHeadersReceivedComplete(int error_code);
+  void ContinueToHeadersReceived();
+  void OnAuthRequiredComplete(
+      extensions::ExtensionWebRequestEventRouter::AuthRequiredResponse rv);
+  void OnHeadersReceivedCompleteForAuth(const net::AuthChallengeInfo& auth_info,
+                                        int rv);
+  void ContinueToCompleted();
+
+  void PauseIncomingMethodCallProcessing();
+  void ResumeIncomingMethodCallProcessing();
+  void OnError(int result);
+  // This is used for detecting errors on mojo connection with the network
+  // service.
+  void OnMojoConnectionErrorWithCustomReason(uint32_t custom_reason,
+                                             const std::string& description);
+  // This is used for detecting errors on mojo connection with original client
+  // (i.e., renderer).
+  void OnMojoConnectionError();
+
+  // Passed from api::WebRequest.
+  WebRequestAPI* web_request_api_;
+
+  // Saved to feed the api::WebRequest.
+  network::ResourceRequest request_;
+
+  WebSocketFactory factory_;
+  mojo::Remote<network::mojom::WebSocketHandshakeClient>
+      forwarding_handshake_client_;
+  mojo::Receiver<network::mojom::WebSocketHandshakeClient>
+      receiver_as_handshake_client_{this};
+  mojo::Receiver<network::mojom::AuthenticationHandler>
+      receiver_as_auth_handler_{this};
+  mojo::Receiver<network::mojom::TrustedHeaderClient>
+      receiver_as_header_client_{this};
+
+  net::HttpRequestHeaders request_headers_;
+  network::mojom::URLResponseHeadPtr response_;
+  net::AuthCredentials auth_credentials_;
+  OnAuthRequiredCallback auth_required_callback_;
+  scoped_refptr<net::HttpResponseHeaders> override_headers_;
+  std::vector<network::mojom::HttpHeaderPtr> additional_headers_;
+
+  OnBeforeSendHeadersCallback on_before_send_headers_callback_;
+  OnHeadersReceivedCallback on_headers_received_callback_;
+
+  GURL redirect_url_;
+  bool is_done_ = false;
+  bool has_extra_headers_;
+  mojo::PendingRemote<network::mojom::WebSocket> websocket_;
+  mojo::PendingReceiver<network::mojom::WebSocketClient> client_receiver_;
+  network::mojom::WebSocketHandshakeResponsePtr handshake_response_ = nullptr;
+  mojo::ScopedDataPipeConsumerHandle readable_;
+
+  extensions::WebRequestInfo info_;
+
+  // Notifies the proxy that the browser context has been shutdown.
+  std::unique_ptr<KeyedServiceShutdownNotifier::Subscription>
+      shutdown_notifier_;
+
+  base::WeakPtrFactory<ProxyingWebSocket> weak_factory_{this};
+  DISALLOW_COPY_AND_ASSIGN(ProxyingWebSocket);
+};
+
+}  // namespace electron
+
+#endif  // SHELL_BROWSER_NET_PROXYING_WEBSOCKET_H_

shell/browser/net/web_request_api_interface.h

@@ -0,0 +1,61 @@
+// Copyright (c) 2020 GitHub, Inc.
+// Use of this source code is governed by the MIT license that can be
+// found in the LICENSE file.
+
+#ifndef SHELL_BROWSER_NET_WEB_REQUEST_API_INTERFACE_H_
+#define SHELL_BROWSER_NET_WEB_REQUEST_API_INTERFACE_H_
+
+#include <set>
+#include <string>
+
+#include "extensions/browser/api/web_request/web_request_info.h"
+#include "services/network/public/cpp/resource_request.h"
+
+namespace electron {
+
+// Defines the interface for WebRequest API, implemented by api::WebRequestNS.
+class WebRequestAPI {
+ public:
+  virtual ~WebRequestAPI() {}
+
+  using BeforeSendHeadersCallback =
+      base::OnceCallback<void(const std::set<std::string>& removed_headers,
+                              const std::set<std::string>& set_headers,
+                              int error_code)>;
+
+  virtual bool HasListener() const = 0;
+  virtual int OnBeforeRequest(extensions::WebRequestInfo* info,
+                              const network::ResourceRequest& request,
+                              net::CompletionOnceCallback callback,
+                              GURL* new_url) = 0;
+  virtual int OnBeforeSendHeaders(extensions::WebRequestInfo* info,
+                                  const network::ResourceRequest& request,
+                                  BeforeSendHeadersCallback callback,
+                                  net::HttpRequestHeaders* headers) = 0;
+  virtual int OnHeadersReceived(
+      extensions::WebRequestInfo* info,
+      const network::ResourceRequest& request,
+      net::CompletionOnceCallback callback,
+      const net::HttpResponseHeaders* original_response_headers,
+      scoped_refptr<net::HttpResponseHeaders>* override_response_headers,
+      GURL* allowed_unsafe_redirect_url) = 0;
+  virtual void OnSendHeaders(extensions::WebRequestInfo* info,
+                             const network::ResourceRequest& request,
+                             const net::HttpRequestHeaders& headers) = 0;
+  virtual void OnBeforeRedirect(extensions::WebRequestInfo* info,
+                                const network::ResourceRequest& request,
+                                const GURL& new_location) = 0;
+  virtual void OnResponseStarted(extensions::WebRequestInfo* info,
+                                 const network::ResourceRequest& request) = 0;
+  virtual void OnErrorOccurred(extensions::WebRequestInfo* info,
+                               const network::ResourceRequest& request,
+                               int net_error) = 0;
+  virtual void OnCompleted(extensions::WebRequestInfo* info,
+                           const network::ResourceRequest& request,
+                           int net_error) = 0;
+  virtual void OnRequestWillBeDestroyed(extensions::WebRequestInfo* info) = 0;
+};
+
+}  // namespace electron
+
+#endif  // SHELL_BROWSER_NET_WEB_REQUEST_API_INTERFACE_H_