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

1 files changed, 132 insertions, 0 deletions

diff --git a/cli/tests/testdata/websocketstream_test.ts b/cli/tests/testdata/websocketstream_test.ts
index db9d1a094..71969314e 100644
--- a/cli/tests/testdata/websocketstream_test.ts
+++ b/cli/tests/testdata/websocketstream_test.ts
@@ -201,3 +201,135 @@ Deno.test("forbidden headers", async () => {
   await ws.closed;
   listener.close();
 });
+
+Deno.test("sync close with empty stream", async () => {
+  const listener = Deno.listen({ port: 4512 });
+  const promise = (async () => {
+    const conn = await listener.accept();
+    const httpConn = Deno.serveHttp(conn);
+    const { request, respondWith } = (await httpConn.nextRequest())!;
+    const { response, socket } = Deno.upgradeWebSocket(request);
+    const p = new Promise<void>((resolve) => {
+      socket.onopen = () => {
+        socket.send("first message");
+        socket.send("second message");
+      };
+      socket.onclose = () => resolve();
+    });
+    await respondWith(response);
+    await p;
+  })();
+
+  const ws = new WebSocketStream("ws://localhost:4512");
+  const { readable } = await ws.connection;
+  const reader = readable.getReader();
+  const firstMessage = await reader.read();
+  assertEquals(firstMessage.value, "first message");
+  const secondMessage = await reader.read();
+  assertEquals(secondMessage.value, "second message");
+  ws.close({ code: 1000 });
+  await ws.closed;
+  await promise;
+  listener.close();
+});
+
+Deno.test("sync close with unread messages in stream", async () => {
+  const listener = Deno.listen({ port: 4512 });
+  const promise = (async () => {
+    const conn = await listener.accept();
+    const httpConn = Deno.serveHttp(conn);
+    const { request, respondWith } = (await httpConn.nextRequest())!;
+    const { response, socket } = Deno.upgradeWebSocket(request);
+    const p = new Promise<void>((resolve) => {
+      socket.onopen = () => {
+        socket.send("first message");
+        socket.send("second message");
+        socket.send("third message");
+        socket.send("fourth message");
+      };
+      socket.onclose = () => resolve();
+    });
+    await respondWith(response);
+    await p;
+  })();
+
+  const ws = new WebSocketStream("ws://localhost:4512");
+  const { readable } = await ws.connection;
+  const reader = readable.getReader();
+  const firstMessage = await reader.read();
+  assertEquals(firstMessage.value, "first message");
+  const secondMessage = await reader.read();
+  assertEquals(secondMessage.value, "second message");
+  ws.close({ code: 1000 });
+  await ws.closed;
+  await promise;
+  listener.close();
+});
+
+Deno.test("async close with empty stream", async () => {
+  const listener = Deno.listen({ port: 4512 });
+  const promise = (async () => {
+    const conn = await listener.accept();
+    const httpConn = Deno.serveHttp(conn);
+    const { request, respondWith } = (await httpConn.nextRequest())!;
+    const { response, socket } = Deno.upgradeWebSocket(request);
+    const p = new Promise<void>((resolve) => {
+      socket.onopen = () => {
+        socket.send("first message");
+        socket.send("second message");
+      };
+      socket.onclose = () => resolve();
+    });
+    await respondWith(response);
+    await p;
+  })();
+
+  const ws = new WebSocketStream("ws://localhost:4512");
+  const { readable } = await ws.connection;
+  const reader = readable.getReader();
+  const firstMessage = await reader.read();
+  assertEquals(firstMessage.value, "first message");
+  const secondMessage = await reader.read();
+  assertEquals(secondMessage.value, "second message");
+  setTimeout(() => {
+    ws.close({ code: 1000 });
+  }, 0);
+  await ws.closed;
+  await promise;
+  listener.close();
+});
+
+Deno.test("async close with unread messages in stream", async () => {
+  const listener = Deno.listen({ port: 4512 });
+  const promise = (async () => {
+    const conn = await listener.accept();
+    const httpConn = Deno.serveHttp(conn);
+    const { request, respondWith } = (await httpConn.nextRequest())!;
+    const { response, socket } = Deno.upgradeWebSocket(request);
+    const p = new Promise<void>((resolve) => {
+      socket.onopen = () => {
+        socket.send("first message");
+        socket.send("second message");
+        socket.send("third message");
+        socket.send("fourth message");
+      };
+      socket.onclose = () => resolve();
+    });
+    await respondWith(response);
+    await p;
+  })();
+
+  const ws = new WebSocketStream("ws://localhost:4512");
+  const { readable } = await ws.connection;
+  const reader = readable.getReader();
+  const firstMessage = await reader.read();
+  assertEquals(firstMessage.value, "first message");
+  const secondMessage = await reader.read();
+  assertEquals(secondMessage.value, "second message");
+  setTimeout(() => {
+    ws.close({ code: 1000 });
+  }, 0);
+  await ws.closed;
+  await promise;
+  listener.close();
+});