Add a new overload of Socket::handleProxyStatus for internal use

src/workerd/api/sockets.c++

@@ -323,16 +323,34 @@ void Socket::handleProxyStatus(
         msg = kj::str(msg, ". It looks like you might be trying to connect to a HTTP-based service",
             " — consider using fetch instead");
       }
-      auto exc = kj::Exception(kj::Exception::Type::FAILED, __FILE__, __LINE__,
-        kj::str(JSG_EXCEPTION(Error), msg));
-      resolveFulfiller(js, exc);
-      readable->getController().cancel(js, nullptr).markAsHandled(js);
-      writable->getController().abort(js, nullptr).markAsHandled(js);
+      handleProxyError(js, JSG_KJ_EXCEPTION(FAILED, Error, msg));
     }
   });
   result.markAsHandled(js);
 }
 
+void Socket::handleProxyStatus(jsg::Lock& js, kj::Promise<kj::Maybe<kj::Exception>> connectResult) {
+  // It's kind of weird to take a promise that resolves to a Maybe<Exception> but we can't just use
+  // a Promise<void> and put our logic in the error handler because awaitIo doesn't provide the
+  // jsg::Lock for void promises or to errorFunc implementations, only non-void success callbacks,
+  // but we need the lock in our callback here.
+  // TODO(cleanup): Extend awaitIo to provide the jsg::Lock in more cases.
+  auto& context = IoContext::current();
+  auto result = context.awaitIo(js, kj::mv(connectResult),
+      [this, self = JSG_THIS](jsg::Lock& js, kj::Maybe<kj::Exception> result) -> void {
+    KJ_IF_MAYBE(e, result) {
+      handleProxyError(js, JSG_KJ_EXCEPTION(FAILED, Error, "connection attempt failed"));
+    }
+  });
+  result.markAsHandled(js);
+}
+
+void Socket::handleProxyError(jsg::Lock& js, kj::Exception e) {
+  resolveFulfiller(js, kj::mv(e));
+  readable->getController().cancel(js, nullptr).markAsHandled(js);
+  writable->getController().abort(js, nullptr).markAsHandled(js);
+}
+
 void Socket::handleReadableEof(jsg::Lock& js, jsg::Promise<void> onEof) {
   KJ_ASSERT(!getAllowHalfOpen(options));
   // Listen for EOF on the ReadableStream.

src/workerd/api/sockets.h

@@ -72,7 +72,11 @@ class Socket: public jsg::Object {
 
   void handleProxyStatus(
       jsg::Lock& js, kj::Promise<kj::HttpClient::ConnectRequest::Status> status);
+  void handleProxyStatus(jsg::Lock& js, kj::Promise<kj::Maybe<kj::Exception>> status);
   // Sets up relevant callbacks to handle the case when the proxy rejects our connection.
+  // The first variant is useful for connections established using HTTP connect. The latter is for
+  // connections established any other way, where the lack of an exception indicates we connected
+  // successfully.
 
   void handleReadableEof(jsg::Lock& js, jsg::Promise<void> onEof);
   // Sets up relevant callbacks to handle the case when the readable stream reaches EOF.
@@ -110,6 +114,9 @@ class Socket: public jsg::Object {
   kj::Promise<kj::Own<kj::AsyncIoStream>> processConnection();
   jsg::Promise<void> maybeCloseWriteSide(jsg::Lock& js);
 
+  void handleProxyError(jsg::Lock& js, kj::Exception e);
+  // Helper method for handleProxyStatus implementations.
+
   void resolveFulfiller(jsg::Lock& js, kj::Maybe<kj::Exception> maybeErr) {
     KJ_IF_MAYBE(err, maybeErr) {
       closedResolver.reject(js, kj::cp(*err));