fix(ext/websocket): fix closing of WebSocketStream with unread messages (#15632)

1 files changed, 31 insertions, 1 deletions

diff --git a/ext/websocket/02_websocketstream.js b/ext/websocket/02_websocketstream.js
index 5266c8dfb..cf83fe4c7 100644
--- a/ext/websocket/02_websocketstream.js
+++ b/ext/websocket/02_websocketstream.js
@@ -64,11 +64,14 @@
     ],
   );
 
+  const CLOSE_RESPONSE_TIMEOUT = 5000;
+
   const _rid = Symbol("[[rid]]");
   const _url = Symbol("[[url]]");
   const _connection = Symbol("[[connection]]");
   const _closed = Symbol("[[closed]]");
   const _earlyClose = Symbol("[[earlyClose]]");
+  const _closeSent = Symbol("[[closeSent]]");
   class WebSocketStream {
     [_rid];
 
@@ -268,6 +271,21 @@
                     break;
                   }
                 }
+
+                if (
+                  this[_closeSent].state === "fulfilled" &&
+                  this[_closed].state === "pending"
+                ) {
+                  if (
+                    new Date().getTime() - await this[_closeSent].promise <=
+                      CLOSE_RESPONSE_TIMEOUT
+                  ) {
+                    return pull(controller);
+                  }
+
+                  this[_closed].resolve(value);
+                  core.tryClose(this[_rid]);
+                }
               };
               const readable = new ReadableStream({
                 start: (controller) => {
@@ -286,6 +304,12 @@
                       // needed to ignore warnings & assertions
                     }
                   });
+
+                  PromisePrototypeThen(this[_closeSent].promise, () => {
+                    if (this[_closed].state === "pending") {
+                      return pull(controller);
+                    }
+                  });
                 },
                 pull,
                 cancel: async (reason) => {
@@ -328,6 +352,7 @@
 
     [_earlyClose] = false;
     [_closed] = new Deferred();
+    [_closeSent] = new Deferred();
     get closed() {
       webidl.assertBranded(this, WebSocketStreamPrototype);
       return this[_closed].promise;
@@ -369,8 +394,13 @@
       if (this[_connection].state === "pending") {
         this[_earlyClose] = true;
       } else if (this[_closed].state === "pending") {
-        PromisePrototypeCatch(
+        PromisePrototypeThen(
           core.opAsync("op_ws_close", this[_rid], code, closeInfo.reason),
+          () => {
+            setTimeout(() => {
+              this[_closeSent].resolve(new Date().getTime());
+            }, 0);
+          },
           (err) => {
             this[_rid] && core.tryClose(this[_rid]);
             this[_closed].reject(err);