Implement abort for WebSocketStream handshake

The WebSocketStream explainer describes being able to pass an
AbortSignal object to the WebSocketStream constructor to abort the
handshake. This wasn't implemented yet in Blink. Implement it.

Bug: 983030
Change-Id: I9d96bf145a4978867dd8ae98a13af9ad9cff752c
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1827116Reviewed-by: Yutaka Hirano <yhirano@chromium.org>
Commit-Queue: Adam Rice <ricea@chromium.org>
Cr-Commit-Position: refs/heads/master@{#701971}

third_party/blink/renderer/modules/websockets/mock_websocket_channel.h

@@ -52,6 +52,7 @@ class MockWebSocketChannel : public WebSocketChannel {
     FailMock(reason, level, location.get());
   }
   MOCK_METHOD0(Disconnect, void());
+  MOCK_METHOD0(CancelHandshake, void());
   MOCK_METHOD0(ApplyBackpressure, void());
   MOCK_METHOD0(RemoveBackpressure, void());
 };

third_party/blink/renderer/modules/websockets/websocket_channel.h

@@ -102,6 +102,10 @@ class MODULES_EXPORT WebSocketChannel
   // Do not call any methods after calling this method.
   virtual void Disconnect() = 0;  // Will suppress didClose().
 
+  // Cancel the WebSocket handshake. Does nothing if the connection is already
+  // established. Do not call any other methods after this one.
+  virtual void CancelHandshake() = 0;
+
   // Clients can call ApplyBackpressure() to indicate that they want to stop
   // receiving new messages. WebSocketChannelClient::DidReceive*Message() may
   // still be called after this, until existing flow control quota is used up.

third_party/blink/renderer/modules/websockets/websocket_channel_impl.cc

@@ -417,6 +417,18 @@ void WebSocketChannelImpl::Disconnect() {
   Dispose();
 }
 
+void WebSocketChannelImpl::CancelHandshake() {
+  NETWORK_DVLOG(1) << this << " CancelHandshake()";
+  if (GetState() != State::kConnecting)
+    return;
+
+  // This may still disconnect even if the handshake is complete if we haven't
+  // got the message yet.
+  // TODO(ricea): Plumb it through to the network stack to fix the race
+  // condition.
+  Disconnect();
+}
+
 void WebSocketChannelImpl::ApplyBackpressure() {
   backpressure_ = true;
 }

third_party/blink/renderer/modules/websockets/websocket_channel_impl.h

@@ -107,6 +107,7 @@ class MODULES_EXPORT WebSocketChannelImpl final
             mojom::ConsoleMessageLevel,
             std::unique_ptr<SourceLocation>) override;
   void Disconnect() override;
+  void CancelHandshake() override;
   void ApplyBackpressure() override;
   void RemoveBackpressure() override;
 

third_party/blink/renderer/modules/websockets/websocket_stream.cc

@@ -14,6 +14,7 @@
 #include "third_party/blink/renderer/bindings/core/v8/v8_binding_for_core.h"
 #include "third_party/blink/renderer/bindings/core/v8/v8_throw_dom_exception.h"
 #include "third_party/blink/renderer/bindings/modules/v8/v8_websocket_close_info.h"
+#include "third_party/blink/renderer/core/dom/abort_signal.h"
 #include "third_party/blink/renderer/core/execution_context/execution_context.h"
 #include "third_party/blink/renderer/core/streams/readable_stream.h"
 #include "third_party/blink/renderer/core/streams/readable_stream_default_controller_interface.h"
@@ -598,7 +599,21 @@ void WebSocketStream::Connect(ScriptState* script_state,
   DVLOG(1) << "WebSocketStream " << this << " Connect() url=" << url
            << " options=" << options;
 
-  // TODO(ricea): Support AbortSignal.
+  auto* signal = options->signal();
+  if (signal && signal->aborted()) {
+    auto exception = V8ThrowDOMException::CreateOrEmpty(
+        script_state->GetIsolate(), DOMExceptionCode::kAbortError,
+        "WebSocket handshake was aborted");
+    connection_resolver_->Reject(exception);
+    closed_resolver_->Reject(exception);
+    return;
+  }
+
+  if (signal) {
+    signal->AddAlgorithm(
+        WTF::Bind(&WebSocketStream::OnAbort, WrapWeakPersistent(this)));
+  }
+
   auto result = common_.Connect(
       ExecutionContext::From(script_state), url,
       options->hasProtocols() ? options->protocols() : Vector<String>(),
@@ -680,6 +695,22 @@ v8::Local<v8::Value> WebSocketStream::CreateNetworkErrorDOMException() {
                                             "A network error occurred");
 }
 
+void WebSocketStream::OnAbort() {
+  DVLOG(1) << "WebSocketStream " << this << " OnAbort()";
+
+  if (was_ever_connected_ || !channel_)
+    return;
+
+  channel_->CancelHandshake();
+  channel_ = nullptr;
+
+  auto exception = V8ThrowDOMException::CreateOrEmpty(
+      script_state_->GetIsolate(), DOMExceptionCode::kAbortError,
+      "WebSocket handshake was aborted");
+  connection_resolver_->Reject(exception);
+  closed_resolver_->Reject(exception);
+}
+
 WebSocketCloseInfo* WebSocketStream::MakeCloseInfo(uint16_t code,
                                                    const String& reason) {
   DVLOG(1) << "WebSocketStream MakeCloseInfo() code=" << code

third_party/blink/renderer/modules/websockets/websocket_stream.h

@@ -108,6 +108,7 @@ class MODULES_EXPORT WebSocketStream final
   void CloseInternal(int code,
                      const String& reason,
                      ExceptionState& exception_state);
+  void OnAbort();
 
   v8::Local<v8::Value> CreateNetworkErrorDOMException();
   static WebSocketCloseInfo* MakeCloseInfo(uint16_t code, const String& reason);

third_party/blink/web_tests/external/wpt/websockets/stream-tentative/abort.any.js

+// META: script=../websocket.sub.js
+// META: script=resources/url-constants.js
+// META: script=/common/utils.js
+// META: global=window,worker
+
+promise_test(async t => {
+  const controller = new AbortController();
+  controller.abort();
+  const key = token();
+  const wsUrl = new URL(
+      `/fetch/api/resources/stash-put.py?key=${key}&value=connected`,
+      location.href);
+  wsUrl.protocol = wsUrl.protocol.replace('http', 'ws');
+  // We intentionally use the port for the HTTP server, not the WebSocket
+  // server, because we don't expect the connection to be performed.
+  const wss = new WebSocketStream(wsUrl, { signal: controller.signal });
+  await promise_rejects(t, 'AbortError', wss.connection,
+                        'connection should reject');
+  await promise_rejects(t, 'AbortError', wss.closed, 'closed should reject');
+  // An incorrect implementation could pass this test due a race condition,
+  // but it is hard to completely eliminate the possibility.
+  const response = await fetch(`/fetch/api/resources/stash-take.py?key=${key}`);
+  assert_equals(await response.text(), 'null', 'response should be null');
+}, 'abort before constructing should prevent connection');
+
+promise_test(async t => {
+  const controller = new AbortController();
+  const wss = new WebSocketStream(`${BASEURL}/handshake_sleep_2`,
+                                  { signal: controller.signal });
+  // Give the connection a chance to start.
+  await new Promise(resolve => t.step_timeout(resolve, 0));
+  controller.abort();
+  await promise_rejects(t, 'AbortError', wss.connection,
+                        'connection should reject');
+  await promise_rejects(t, 'AbortError', wss.closed, 'closed should reject');
+}, 'abort during handshake should work');
+
+promise_test(async t => {
+  const controller = new AbortController();
+  const wss = new WebSocketStream(ECHOURL, { signal: controller.signal });
+  const { readable, writable } = await wss.connection;
+  controller.abort();
+  writable.getWriter().write('connected');
+  const { value } = await readable.getReader().read();
+  assert_equals(value, 'connected', 'value should match');
+}, 'abort after connect should do nothing');