Merge pull request #788 from capnproto/cancel-http

Make HttpServer detect client disconnect and cancel request processing

c++/src/kj/async-io-test.c++

@@ -2103,5 +2103,98 @@ KJ_TEST("Userland tee buffer size limit") {
   }
 }
 
+KJ_TEST("Userspace OneWayPipe whenWriteDisconnected()") {
+  kj::EventLoop loop;
+  WaitScope ws(loop);
+
+  auto pipe = newOneWayPipe();
+
+  auto abortedPromise = pipe.out->whenWriteDisconnected();
+  KJ_ASSERT(!abortedPromise.poll(ws));
+
+  pipe.in = nullptr;
+
+  KJ_ASSERT(abortedPromise.poll(ws));
+  abortedPromise.wait(ws);
+}
+
+KJ_TEST("Userspace TwoWayPipe whenWriteDisconnected()") {
+  kj::EventLoop loop;
+  WaitScope ws(loop);
+
+  auto pipe = newTwoWayPipe();
+
+  auto abortedPromise = pipe.ends[0]->whenWriteDisconnected();
+  KJ_ASSERT(!abortedPromise.poll(ws));
+
+  pipe.ends[1] = nullptr;
+
+  KJ_ASSERT(abortedPromise.poll(ws));
+  abortedPromise.wait(ws);
+}
+
+#if !_WIN32  // We don't currently support detecting disconnect with IOCP.
+#if !__CYGWIN__  // TODO(soon): Figure out why whenWriteDisconnected() doesn't work on Cygwin.
+
+KJ_TEST("OS OneWayPipe whenWriteDisconnected()") {
+  auto io = setupAsyncIo();
+
+  auto pipe = io.provider->newOneWayPipe();
+
+  pipe.out->write("foo", 3).wait(io.waitScope);
+  auto abortedPromise = pipe.out->whenWriteDisconnected();
+  KJ_ASSERT(!abortedPromise.poll(io.waitScope));
+
+  pipe.in = nullptr;
+
+  KJ_ASSERT(abortedPromise.poll(io.waitScope));
+  abortedPromise.wait(io.waitScope);
+}
+
+KJ_TEST("OS TwoWayPipe whenWriteDisconnected()") {
+  auto io = setupAsyncIo();
+
+  auto pipe = io.provider->newTwoWayPipe();
+
+  pipe.ends[0]->write("foo", 3).wait(io.waitScope);
+  pipe.ends[1]->write("bar", 3).wait(io.waitScope);
+
+  auto abortedPromise = pipe.ends[0]->whenWriteDisconnected();
+  KJ_ASSERT(!abortedPromise.poll(io.waitScope));
+
+  pipe.ends[1] = nullptr;
+
+  KJ_ASSERT(abortedPromise.poll(io.waitScope));
+  abortedPromise.wait(io.waitScope);
+
+  char buffer[4];
+  KJ_ASSERT(pipe.ends[0]->tryRead(&buffer, sizeof(buffer), sizeof(buffer)).wait(io.waitScope) == 3);
+  buffer[3] = '\0';
+  KJ_EXPECT(buffer == "bar"_kj);
+}
+
+KJ_TEST("import socket FD that's already broken") {
+  auto io = setupAsyncIo();
+
+  int fds[2];
+  KJ_SYSCALL(socketpair(AF_UNIX, SOCK_STREAM, 0, fds));
+  KJ_SYSCALL(write(fds[1], "foo", 3));
+  KJ_SYSCALL(close(fds[1]));
+
+  auto stream = io.lowLevelProvider->wrapSocketFd(fds[0], LowLevelAsyncIoProvider::TAKE_OWNERSHIP);
+
+  auto abortedPromise = stream->whenWriteDisconnected();
+  KJ_ASSERT(abortedPromise.poll(io.waitScope));
+  abortedPromise.wait(io.waitScope);
+
+  char buffer[4];
+  KJ_ASSERT(stream->tryRead(&buffer, sizeof(buffer), sizeof(buffer)).wait(io.waitScope) == 3);
+  buffer[3] = '\0';
+  KJ_EXPECT(buffer == "foo"_kj);
+}
+
+#endif  // !__CYGWIN__
+#endif  // !_WIN32
+
 }  // namespace
 }  // namespace kj

c++/src/kj/async-io-unix.c++

@@ -179,6 +179,17 @@ public:
     }
   }
 
+  Promise<void> whenWriteDisconnected() override {
+    KJ_IF_MAYBE(p, writeDisconnectedPromise) {
+      return p->addBranch();
+    } else {
+      auto fork = observer.whenWriteDisconnected().fork();
+      auto result = fork.addBranch();
+      writeDisconnectedPromise = kj::mv(fork);
+      return kj::mv(result);
+    }
+  }
+
   void shutdownWrite() override {
     // There's no legitimate way to get an AsyncStreamFd that isn't a socket through the
     // UnixAsyncIoProvider interface.
@@ -290,6 +301,7 @@ public:
 private:
   UnixEventPort& eventPort;
   UnixEventPort::FdObserver observer;
+  Maybe<ForkedPromise<void>> writeDisconnectedPromise;
 
   Promise<size_t> tryReadInternal(void* buffer, size_t minBytes, size_t maxBytes,
                                   size_t alreadyRead) {

c++/src/kj/async-io-win32.c++

@@ -311,6 +311,23 @@ public:
     });
   }
 
+  Promise<void> whenWriteDisconnected() override {
+    // Windows IOCP does not provide a direct, documented way to detect when the socket disconnects
+    // without actually doing a read or write. However, there is an undocoumented-but-stable
+    // ioctl called IOCTL_AFD_POLL which can be used for this purpose. In fact, select() is
+    // implemented in terms of this ioctl -- performed synchronously -- but it's entirely possible
+    // to put only one socket into the list and perform the ioctl asynchronously. Here's the
+    // source code for select() in Windows 2000 (not sure how this became public...):
+    //
+    //     https://github.com/pustladi/Windows-2000/blob/661d000d50637ed6fab2329d30e31775046588a9/private/net/sockets/winsock2/wsp/msafd/select.c#L59-L655
+    //
+    // And here's an interesting discussion: https://github.com/python-trio/trio/issues/52
+    //
+    // TODO(soon): Implement this with IOCTL_AFD_POLL. For now I'm leaving it unimplemented because
+    //   I added this method for a Linux-only use case.
+    return NEVER_DONE;
+  }
+
   void shutdownWrite() override {
     // There's no legitimate way to get an AsyncStreamFd that isn't a socket through the
     // Win32AsyncIoProvider interface.

c++/src/kj/async-io.c++

@@ -228,6 +228,12 @@ public:
     } else {
       ownState = kj::heap<AbortedRead>();
       state = *ownState;
+
+      readAborted = true;
+      KJ_IF_MAYBE(f, readAbortFulfiller) {
+        f->get()->fulfill();
+        readAbortFulfiller = nullptr;
+      }
     }
   }
 
@@ -268,6 +274,21 @@ public:
     }
   }
 
+  Promise<void> whenWriteDisconnected() override {
+    if (readAborted) {
+      return kj::READY_NOW;
+    } else KJ_IF_MAYBE(p, readAbortPromise) {
+      return p->addBranch();
+    } else {
+      auto paf = newPromiseAndFulfiller<void>();
+      readAbortFulfiller = kj::mv(paf.fulfiller);
+      auto fork = paf.promise.fork();
+      auto result = fork.addBranch();
+      readAbortPromise = kj::mv(fork);
+      return result;
+    }
+  }
+
   void shutdownWrite() override {
     KJ_IF_MAYBE(s, state) {
       s->shutdownWrite();
@@ -285,6 +306,10 @@ private:
 
   kj::Own<AsyncIoStream> ownState;
 
+  bool readAborted = false;
+  Maybe<Own<PromiseFulfiller<void>>> readAbortFulfiller = nullptr;
+  Maybe<ForkedPromise<void>> readAbortPromise = nullptr;
+
   void endState(AsyncIoStream& obj) {
     KJ_IF_MAYBE(s, state) {
       if (s == &obj) {
@@ -443,6 +468,10 @@ private:
       KJ_FAIL_REQUIRE("can't shutdownWrite() until previous write() completes");
     }
 
+    Promise<void> whenWriteDisconnected() override {
+      KJ_FAIL_ASSERT("can't get here -- implemented by AsyncPipe");
+    }
+
   private:
     PromiseFulfiller<void>& fulfiller;
     AsyncPipe& pipe;
@@ -562,6 +591,10 @@ private:
       KJ_FAIL_REQUIRE("can't shutdownWrite() until previous tryPumpFrom() completes");
     }
 
+    Promise<void> whenWriteDisconnected() override {
+      KJ_FAIL_ASSERT("can't get here -- implemented by AsyncPipe");
+    }
+
   private:
     PromiseFulfiller<uint64_t>& fulfiller;
     AsyncPipe& pipe;
@@ -733,6 +766,10 @@ private:
       pipe.shutdownWrite();
     }
 
+    Promise<void> whenWriteDisconnected() override {
+      KJ_FAIL_ASSERT("can't get here -- implemented by AsyncPipe");
+    }
+
   private:
     PromiseFulfiller<size_t>& fulfiller;
     AsyncPipe& pipe;
@@ -901,6 +938,10 @@ private:
       pipe.shutdownWrite();
     }
 
+    Promise<void> whenWriteDisconnected() override {
+      KJ_FAIL_ASSERT("can't get here -- implemented by AsyncPipe");
+    }
+
   private:
     PromiseFulfiller<uint64_t>& fulfiller;
     AsyncPipe& pipe;
@@ -937,6 +978,9 @@ private:
       // ignore -- currently shutdownWrite() actually means that the PipeWriteEnd was dropped,
       // which is not an error even if reads have been aborted.
     }
+    Promise<void> whenWriteDisconnected() override {
+      KJ_FAIL_ASSERT("can't get here -- implemented by AsyncPipe");
+    }
   };
 
   class ShutdownedWrite final: public AsyncIoStream {
@@ -966,6 +1010,9 @@ private:
       // ignore -- currently shutdownWrite() actually means that the PipeWriteEnd was dropped,
       // so it will only be called once anyhow.
     }
+    Promise<void> whenWriteDisconnected() override {
+      KJ_FAIL_ASSERT("can't get here -- implemented by AsyncPipe");
+    }
   };
 };
 
@@ -1013,6 +1060,10 @@ public:
     return pipe->tryPumpFrom(input, amount);
   }
 
+  Promise<void> whenWriteDisconnected() override {
+    return pipe->whenWriteDisconnected();
+  }
+
 private:
   Own<AsyncPipe> pipe;
   UnwindDetector unwind;
@@ -1049,6 +1100,9 @@ public:
       AsyncInputStream& input, uint64_t amount) override {
     return out->tryPumpFrom(input, amount);
   }
+  Promise<void> whenWriteDisconnected() override {
+    return out->whenWriteDisconnected();
+  }
   void shutdownWrite() override {
     out->shutdownWrite();
   }

c++/src/kj/async-io.h

@@ -106,6 +106,20 @@ public:
   // output stream. If it finds one, it performs the pump. Otherwise, it returns null.
   //
   // The default implementation always returns null.
+
+  virtual Promise<void> whenWriteDisconnected() = 0;
+  // Returns a promise that resolves when the stream has become disconnected such that new write()s
+  // will fail with a DISCONNECTED exception. This is particularly useful, for example, to cancel
+  // work early when it is detected that no one will receive the result.
+  //
+  // Note that not all streams are able to detect this condition without actually performing a
+  // write(); such stream implementations may return a promise that never resolves. (In particular,
+  // as of this writing, whenWriteDisconnected() is not implemented on Windows. Also, for TCP
+  // streams, not all disconnects are detectable -- a power or network failure may lead the
+  // connection to hang forever, or until configured socket options lead to a timeout.)
+  //
+  // Unlike most other asynchronous stream methods, it is safe to call whenWriteDisconnected()
+  // multiple times without canceling the previous promises.
 };
 
 class AsyncIoStream: public AsyncInputStream, public AsyncOutputStream {

c++/src/kj/async-unix.c++

@@ -365,6 +365,13 @@ void UnixEventPort::FdObserver::fire(short events) {
     }
   }
 
+  if (events & (EPOLLHUP | EPOLLERR)) {
+    KJ_IF_MAYBE(f, hupFulfiller) {
+      f->get()->fulfill();
+      hupFulfiller = nullptr;
+    }
+  }
+
   if (events & EPOLLPRI) {
     KJ_IF_MAYBE(f, urgentFulfiller) {
       f->get()->fulfill();
@@ -398,6 +405,12 @@ Promise<void> UnixEventPort::FdObserver::whenUrgentDataAvailable() {
   return kj::mv(paf.promise);
 }
 
+Promise<void> UnixEventPort::FdObserver::whenWriteDisconnected() {
+  auto paf = newPromiseAndFulfiller<void>();
+  hupFulfiller = kj::mv(paf.fulfiller);
+  return kj::mv(paf.promise);
+}
+
 bool UnixEventPort::wait() {
   return doEpollWait(
       timerImpl.timeoutToNextEvent(readClock(), MILLISECONDS, int(maxValue))
@@ -652,6 +665,13 @@ void UnixEventPort::FdObserver::fire(short events) {
     }
   }
 
+  if (events & (POLLHUP | POLLERR | POLLNVAL)) {
+    KJ_IF_MAYBE(f, hupFulfiller) {
+      f->get()->fulfill();
+      hupFulfiller = nullptr;
+    }
+  }
+
   if (events & POLLPRI) {
     KJ_IF_MAYBE(f, urgentFulfiller) {
       f->get()->fulfill();
@@ -675,7 +695,16 @@ void UnixEventPort::FdObserver::fire(short events) {
 short UnixEventPort::FdObserver::getEventMask() {
   return (readFulfiller == nullptr ? 0 : (POLLIN | POLLRDHUP)) |
          (writeFulfiller == nullptr ? 0 : POLLOUT) |
-         (urgentFulfiller == nullptr ? 0 : POLLPRI);
+         (urgentFulfiller == nullptr ? 0 : POLLPRI) |
+         // The POSIX standard says POLLHUP and POLLERR will be reported even if not requested.
+         // But on MacOS, if `events` is 0, then POLLHUP apparently will not be reported:
+         //   https://openradar.appspot.com/37537852
+         // It seems that by settingc any non-zero value -- even one documented as ignored -- we
+         // cause POLLHUP to be reported. Both POLLHUP and POLLERR are documented as being ignored.
+         // So, we'll go ahead and set them. This has no effect on non-broken OSs, causes MacOS to
+         // do the right thing, and sort of looks as if we're explicitly requesting notification of
+         // these two conditions, which we do after all want to know about.
+         POLLHUP | POLLERR;
 }
 
 Promise<void> UnixEventPort::FdObserver::whenBecomesReadable() {
@@ -724,6 +753,19 @@ Promise<void> UnixEventPort::FdObserver::whenUrgentDataAvailable() {
   return kj::mv(paf.promise);
 }
 
+Promise<void> UnixEventPort::FdObserver::whenWriteDisconnected() {
+  if (prev == nullptr) {
+    KJ_DASSERT(next == nullptr);
+    prev = eventPort.observersTail;
+    *prev = this;
+    eventPort.observersTail = &next;
+  }
+
+  auto paf = newPromiseAndFulfiller<void>();
+  hupFulfiller = kj::mv(paf.fulfiller);
+  return kj::mv(paf.promise);
+}
+
 class UnixEventPort::PollContext {
 public:
   PollContext(FdObserver* ptr) {

c++/src/kj/async-unix.h

@@ -278,6 +278,9 @@ public:
   // WARNING: This has some known weird behavior on macOS. See
   //   https://github.com/sandstorm-io/capnproto/issues/374.
 
+  Promise<void> whenWriteDisconnected();
+  // Resolves when poll() on the file descriptor reports POLLHUP or POLLERR.
+
 private:
   UnixEventPort& eventPort;
   int fd;
@@ -286,6 +289,7 @@ private:
   kj::Maybe<Own<PromiseFulfiller<void>>> readFulfiller;
   kj::Maybe<Own<PromiseFulfiller<void>>> writeFulfiller;
   kj::Maybe<Own<PromiseFulfiller<void>>> urgentFulfiller;
+  kj::Maybe<Own<PromiseFulfiller<void>>> hupFulfiller;
   // Replaced each time `whenBecomesReadable()` or `whenBecomesWritable()` is called. Reverted to
   // null every time an event is fired.
 

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

@@ -126,6 +126,8 @@ public:
     }
     return kj::READY_NOW;
   }
+
+  Promise<void> whenWriteDisconnected() override { KJ_UNIMPLEMENTED("not used"); }
 };
 
 KJ_TEST("gzip decompression") {

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

@@ -118,6 +118,8 @@ public:
   Promise<void> write(const void* buffer, size_t size) override;
   Promise<void> write(ArrayPtr<const ArrayPtr<const byte>> pieces) override;
 
+  Promise<void> whenWriteDisconnected() override { return inner.whenWriteDisconnected(); }
+
   inline Promise<void> flush() {
     return pump(Z_SYNC_FLUSH);
   }

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

+// Copyright (c) 2019 Cloudflare, Inc. and contributors
+// Licensed under the MIT License:
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+// Run http-test, but use real OS socketpairs to connect rather than using in-process pipes.
+// This is essentially an integration test between KJ HTTP and KJ OS socket handling.
+#define KJ_HTTP_TEST_USE_OS_PIPE 1
+#include "http-test.c++"

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

@@ -489,6 +489,17 @@ public:
   // shutdown, but is sometimes useful when you want the other end to trigger whatever behavior
   // it normally triggers when a connection is dropped.
 
+  virtual void abort() = 0;
+  // Forcefully close this WebSocket, such that the remote end should get a DISCONNECTED error if
+  // it continues to write. This differs from disconnect(), which only closes the sending
+  // direction, but still allows receives.
+
+  virtual kj::Promise<void> whenAborted() = 0;
+  // Resolves when the remote side aborts the connection such that send() would throw DISCONNECTED,
+  // if this can be detected without actually writing a message. (If not, this promise never
+  // resolves, but send() or receive() will throw DISCONNECTED when appropriate. See also
+  // kj::AsyncOutputStream::whenWriteDisconnected().)
+
   struct Close {
     uint16_t code;
     kj::String reason;
@@ -629,6 +640,9 @@ public:
   //
   // `url` and `headers` are invalidated on the first read from `requestBody` or when the returned
   // promise resolves, whichever comes first.
+  //
+  // Request processing can be canceled by dropping the returned promise. HttpServer may do so if
+  // the client disconnects prematurely.
 
   virtual kj::Promise<kj::Own<kj::AsyncIoStream>> connect(kj::StringPtr host);
   // Handles CONNECT requests. Only relevant for proxy services. Default implementation throws

c++/src/kj/compat/tls.c++

@@ -180,6 +180,10 @@ public:
     return writeInternal(pieces[0], pieces.slice(1, pieces.size()));
   }
 
+  Promise<void> whenWriteDisconnected() override {
+    return inner.whenWriteDisconnected();
+  }
+
   void shutdownWrite() override {
     KJ_REQUIRE(shutdownTask == nullptr, "already called shutdownWrite()");