Implement WebSocket core protocol.

Still need to add handshake separately.

c++/src/kj/compat/http-test.c++

@@ -381,6 +381,26 @@ kj::Promise<void> expectRead(kj::AsyncInputStream& in, kj::StringPtr expected) {
   }));
 }
 
+kj::Promise<void> expectRead(kj::AsyncInputStream& in, kj::ArrayPtr<const byte> expected) {
+  if (expected.size() == 0) return kj::READY_NOW;
+
+  auto buffer = kj::heapArray<byte>(expected.size());
+
+  auto promise = in.tryRead(buffer.begin(), 1, buffer.size());
+  return promise.then(kj::mvCapture(buffer, [&in,expected](kj::Array<byte> buffer, size_t amount) {
+    if (amount == 0) {
+      KJ_FAIL_ASSERT("expected data never sent", expected);
+    }
+
+    auto actual = buffer.slice(0, amount);
+    if (memcmp(actual.begin(), expected.begin(), actual.size()) != 0) {
+      KJ_FAIL_ASSERT("data from stream doesn't match expected", expected, actual);
+    }
+
+    return expectRead(in, expected.slice(amount, expected.size()));
+  }));
+}
+
 void testHttpClientRequest(kj::AsyncIoContext& io, const HttpRequestTestCase& testCase) {
   auto pipe = io.provider->newTwoWayPipe();
 
@@ -1150,6 +1170,298 @@ KJ_TEST("HttpClient <-> HttpServer") {
 
 // -----------------------------------------------------------------------------
 
+KJ_TEST("WebSocket core protocol") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+
+  auto client = newWebSocket(kj::mv(pipe.ends[0]));
+  auto server = newWebSocket(kj::mv(pipe.ends[1]));
+
+  auto mediumString = kj::strArray(kj::repeat(kj::StringPtr("123456789"), 30), "");
+  auto bigString = kj::strArray(kj::repeat(kj::StringPtr("123456789"), 10000), "");
+
+  auto clientTask = client->send(kj::StringPtr("hello"))
+      .then([&]() { return client->send(mediumString); })
+      .then([&]() { return client->send(bigString); })
+      .then([&]() { return client->send(kj::StringPtr("world").asBytes()); })
+      .then([&]() { return client->close(1234, "bored"); });
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "hello");
+  }
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == mediumString);
+  }
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == bigString);
+  }
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::Array<byte>>());
+    KJ_EXPECT(kj::str(message.get<kj::Array<byte>>().asChars()) == "world");
+  }
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<WebSocket::Close>());
+    KJ_EXPECT(message.get<WebSocket::Close>().code == 1234);
+    KJ_EXPECT(message.get<WebSocket::Close>().reason == "bored");
+  }
+
+  auto serverTask = server->close(4321, "whatever");
+
+  {
+    auto message = client->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<WebSocket::Close>());
+    KJ_EXPECT(message.get<WebSocket::Close>().code == 4321);
+    KJ_EXPECT(message.get<WebSocket::Close>().reason == "whatever");
+  }
+
+  clientTask.wait(io.waitScope);
+  serverTask.wait(io.waitScope);
+}
+
+KJ_TEST("WebSocket fragmented") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+
+  auto client = kj::mv(pipe.ends[0]);
+  auto server = newWebSocket(kj::mv(pipe.ends[1]));
+
+  byte DATA[] = {
+    0x01, 0x06, 'h', 'e', 'l', 'l', 'o', ' ',
+
+    0x00, 0x03, 'w', 'o', 'r',
+
+    0x80, 0x02, 'l', 'd',
+  };
+
+  auto clientTask = client->write(DATA, sizeof(DATA));
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "hello world");
+  }
+
+  clientTask.wait(io.waitScope);
+}
+
+class ConstantMaskGenerator final: public WebSocket::MaskKeyGenerator {
+public:
+  void next(byte (&bytes)[4]) override {
+    bytes[0] = 12;
+    bytes[1] = 34;
+    bytes[2] = 56;
+    bytes[3] = 78;
+  }
+};
+
+KJ_TEST("WebSocket masked") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+  ConstantMaskGenerator maskGenerator;
+
+  auto client = kj::mv(pipe.ends[0]);
+  auto server = newWebSocket(kj::mv(pipe.ends[1]), maskGenerator);
+
+  byte DATA[] = {
+    0x81, 0x86, 12, 34, 56, 78, 'h' ^ 12, 'e' ^ 34, 'l' ^ 56, 'l' ^ 78, 'o' ^ 12, ' ' ^ 34,
+  };
+
+  auto clientTask = client->write(DATA, sizeof(DATA));
+  auto serverTask = server->send(kj::StringPtr("hello "));
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "hello ");
+  }
+
+  expectRead(*client, DATA).wait(io.waitScope);
+
+  clientTask.wait(io.waitScope);
+  serverTask.wait(io.waitScope);
+}
+
+KJ_TEST("WebSocket unsolicited pong") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+
+  auto client = kj::mv(pipe.ends[0]);
+  auto server = newWebSocket(kj::mv(pipe.ends[1]));
+
+  byte DATA[] = {
+    0x01, 0x06, 'h', 'e', 'l', 'l', 'o', ' ',
+
+    0x8A, 0x03, 'f', 'o', 'o',
+
+    0x80, 0x05, 'w', 'o', 'r', 'l', 'd',
+  };
+
+  auto clientTask = client->write(DATA, sizeof(DATA));
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "hello world");
+  }
+
+  clientTask.wait(io.waitScope);
+}
+
+KJ_TEST("WebSocket ping") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+
+  auto client = kj::mv(pipe.ends[0]);
+  auto server = newWebSocket(kj::mv(pipe.ends[1]));
+
+  // Be extra-annoying by having the ping arrive between fragments.
+  byte DATA[] = {
+    0x01, 0x06, 'h', 'e', 'l', 'l', 'o', ' ',
+
+    0x89, 0x03, 'f', 'o', 'o',
+
+    0x80, 0x05, 'w', 'o', 'r', 'l', 'd',
+  };
+
+  auto clientTask = client->write(DATA, sizeof(DATA));
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "hello world");
+  }
+
+  auto serverTask = server->send(kj::StringPtr("bar"));
+
+  byte EXPECTED[] = {
+    0x8A, 0x03, 'f', 'o', 'o',  // pong
+    0x81, 0x03, 'b', 'a', 'r',  // message
+  };
+
+  expectRead(*client, EXPECTED).wait(io.waitScope);
+
+  clientTask.wait(io.waitScope);
+  serverTask.wait(io.waitScope);
+}
+
+KJ_TEST("WebSocket ping mid-send") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+
+  auto client = kj::mv(pipe.ends[0]);
+  auto server = newWebSocket(kj::mv(pipe.ends[1]));
+
+  auto bigString = kj::strArray(kj::repeat(kj::StringPtr("12345678"), 65536), "");
+  auto serverTask = server->send(bigString).eagerlyEvaluate(nullptr);
+
+  byte DATA[] = {
+    0x89, 0x03, 'f', 'o', 'o',  // ping
+    0x81, 0x03, 'b', 'a', 'r',  // some other message
+  };
+
+  auto clientTask = client->write(DATA, sizeof(DATA));
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "bar");
+  }
+
+  byte EXPECTED1[] = { 0x81, 0x7f, 0, 0, 0, 0, 0, 8, 0, 0 };
+  expectRead(*client, EXPECTED1).wait(io.waitScope);
+  expectRead(*client, bigString).wait(io.waitScope);
+
+  byte EXPECTED2[] = { 0x8A, 0x03, 'f', 'o', 'o' };
+  expectRead(*client, EXPECTED2).wait(io.waitScope);
+
+  clientTask.wait(io.waitScope);
+  serverTask.wait(io.waitScope);
+}
+
+KJ_TEST("WebSocket double-ping mid-send") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+
+  auto client = kj::mv(pipe.ends[0]);
+  auto server = newWebSocket(kj::mv(pipe.ends[1]));
+
+  auto bigString = kj::strArray(kj::repeat(kj::StringPtr("12345678"), 65536), "");
+  auto serverTask = server->send(bigString).eagerlyEvaluate(nullptr);
+
+  byte DATA[] = {
+    0x89, 0x03, 'f', 'o', 'o',  // ping
+    0x89, 0x03, 'q', 'u', 'x',  // ping2
+    0x81, 0x03, 'b', 'a', 'r',  // some other message
+  };
+
+  auto clientTask = client->write(DATA, sizeof(DATA));
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "bar");
+  }
+
+  byte EXPECTED1[] = { 0x81, 0x7f, 0, 0, 0, 0, 0, 8, 0, 0 };
+  expectRead(*client, EXPECTED1).wait(io.waitScope);
+  expectRead(*client, bigString).wait(io.waitScope);
+
+  byte EXPECTED2[] = { 0x8A, 0x03, 'q', 'u', 'x' };
+  expectRead(*client, EXPECTED2).wait(io.waitScope);
+
+  clientTask.wait(io.waitScope);
+  serverTask.wait(io.waitScope);
+}
+
+KJ_TEST("WebSocket ping received during pong send") {
+  auto io = kj::setupAsyncIo();
+  auto pipe = io.provider->newTwoWayPipe();
+
+  auto client = kj::mv(pipe.ends[0]);
+  auto server = newWebSocket(kj::mv(pipe.ends[1]));
+
+  // Send a very large ping so that sending the pong takes a while. Then send a second ping
+  // immediately after.
+  byte PREFIX[] = { 0x89, 0x7f, 0, 0, 0, 0, 0, 8, 0, 0 };
+  auto bigString = kj::strArray(kj::repeat(kj::StringPtr("12345678"), 65536), "");
+  byte POSTFIX[] = {
+    0x89, 0x03, 'f', 'o', 'o',
+    0x81, 0x03, 'b', 'a', 'r',
+  };
+
+  kj::ArrayPtr<const byte> parts[] = {PREFIX, bigString.asBytes(), POSTFIX};
+  auto clientTask = client->write(parts);
+
+  {
+    auto message = server->receive().wait(io.waitScope);
+    KJ_ASSERT(message.is<kj::String>());
+    KJ_EXPECT(message.get<kj::String>() == "bar");
+  }
+
+  byte EXPECTED1[] = { 0x8A, 0x7f, 0, 0, 0, 0, 0, 8, 0, 0 };
+  expectRead(*client, EXPECTED1).wait(io.waitScope);
+  expectRead(*client, bigString).wait(io.waitScope);
+
+  byte EXPECTED2[] = { 0x8A, 0x03, 'f', 'o', 'o' };
+  expectRead(*client, EXPECTED2).wait(io.waitScope);
+
+  clientTask.wait(io.waitScope);
+}
+
+// -----------------------------------------------------------------------------
+
 KJ_TEST("HttpServer request timeout") {
   auto PIPELINE_TESTS = pipelineTestCases();
 

c++/src/kj/compat/http.h

@@ -369,12 +369,56 @@ private:
 };
 
 class WebSocket {
+  // Interface representincg an open WebSocket session.
+  //
+  // Each side can send and receive data and "close" messages.
+  //
+  // Ping/Pong and message fragmentation are not exposed through this interface. These features of
+  // the underlying WebSocket protocol are not exposed by the browser-level Javascript API either,
+  // and thus applications typically need to implement these features at the applicaiton protocol
+  // level instead. The implementation is, however, expected to reply to Ping messages it receives.
+
 public:
-  WebSocket(kj::Own<kj::AsyncIoStream> stream);
-  // Create a WebSocket wrapping the given I/O stream.
+  virtual kj::Promise<void> send(kj::ArrayPtr<const byte> message) = 0;
+  virtual kj::Promise<void> send(kj::ArrayPtr<const char> message) = 0;
+  // Send a message (binary or text). The underlying buffer must remain valid, and you must not
+  // call send() again, until the returned promise resolves.
+
+  virtual kj::Promise<void> close(uint16_t code, kj::StringPtr reason) = 0;
+  // Send a Close message.
+  //
+  // Note that the returned Promise resolves once the message has been sent -- it does NOT wait
+  // for the other end to send a Close reply. The application should await a reply before dropping
+  // the WebSocket object.
+
+  struct Close {
+    uint16_t code;
+    kj::String reason;
+  };
+
+  typedef kj::OneOf<kj::String, kj::Array<byte>, Close> Message;
+
+  virtual kj::Promise<Message> receive() = 0;
+  // Read one message from the WebSocket and return it. Can only call once at a time. Do not call
+  // again after EndOfStream is received.
+
+  class MaskKeyGenerator {
+    // Class for generating WebSocket packet masks keys. See RFC6455 to understand how masking is
+    // used in WebSockets.
+    //
+    // The RFC insists that mask keys must be crypto-random, but it is not crypto -- it's just a
+    // value to be XOR'd with each four bytes of the data, and the mask itself is transmitted in
+    // plaintext ahead of the message. Apparently the WebSocket designers imagined that a random
+    // mask would make mass surveillance via string matching more difficult, but in practice this
+    // seems like no more than a minor speedbump. The other purpose of the mask is to prevent dumb
+    // proxies and captive portals from getting confused, but even a global constant mask could
+    // accomplish that.
+    //
+    // KJ leaves it up to the application to decide how to generate masks.
 
-  kj::Promise<void> send(kj::ArrayPtr<const byte> message);
-  kj::Promise<void> send(kj::ArrayPtr<const char> message);
+  public:
+    virtual void next(byte (&bytes)[4]) = 0;
+  };
 };
 
 class HttpClient {
@@ -428,10 +472,11 @@ public:
     // `statusText` and `headers` remain valid until `upstreamOrBody` is dropped.
   };
   virtual kj::Promise<WebSocketResponse> openWebSocket(
-      kj::StringPtr url, const HttpHeaders& headers, kj::Own<WebSocket> downstream);
+      kj::StringPtr url, const HttpHeaders& headers);
   // Tries to open a WebSocket. Default implementation calls send() and never returns a WebSocket.
   //
-  // `url` and `headers` are invalidated when the returned promise resolves.
+  // `url` and `headers` need only remain valid until `openWebSocket()` returns (they can be
+  // stack-allocated).
 
   virtual kj::Promise<kj::Own<kj::AsyncIoStream>> connect(kj::StringPtr host);
   // Handles CONNECT requests. Only relevant for proxy clients. Default implementation throws
@@ -478,12 +523,11 @@ public:
 
   class WebSocketResponse: public Response {
   public:
-    kj::Own<WebSocket> startWebSocket(
-        uint statusCode, kj::StringPtr statusText, const HttpHeaders& headers,
-        WebSocket& upstream);
-    // Begin the response.
+    kj::Own<WebSocket> acceptWebSocket(
+        uint statusCode, kj::StringPtr statusText, const HttpHeaders& headers);
+    // Accept and open the WebSocket.
     //
-    // `statusText` and `headers` need only remain valid until startWebSocket() returns (they can
+    // `statusText` and `headers` need only remain valid until acceptWebSocket() returns (they can
     // be stack-allocated).
   };
 
@@ -523,6 +567,15 @@ kj::Own<HttpClient> newHttpClient(HttpService& service);
 kj::Own<HttpService> newHttpService(HttpClient& client);
 // Adapts an HttpClient to an HttpService and vice versa.
 
+kj::Own<WebSocket> newWebSocket(kj::Own<kj::AsyncIoStream> stream,
+                                kj::Maybe<WebSocket::MaskKeyGenerator&> maskKeyGenerator = nullptr);
+// Create a new WebSocket on top of the given stream. It is assumed that the HTTP -> WebSocket
+// upgrade handshake has already occurred (or is not needed), and messages can immediately be
+// sent and received on the stream. Normally applications would not call this directly.
+//
+// `maskKeyGenerator` is optional, but if omitted, the WebSocket frames will not be masked. Refer
+// to RFC6455 to understand when masking is required.
+
 struct HttpServerSettings {
   kj::Duration headerTimeout = 15 * kj::SECONDS;
   // After initial connection open, or after receiving the first byte of a pipelined request,