Fix headers + add test

2 files changed, 24 insertions, 5 deletions

diff --git a/src/bun.js/api/server.zig b/src/bun.js/api/server.zig
index c939d4de1..ea4625f73 100644
--- a/src/bun.js/api/server.zig
+++ b/src/bun.js/api/server.zig
@@ -3307,6 +3307,8 @@ pub fn NewServer(comptime ssl_enabled_: bool, comptime debug_mode_: bool) type {
                 }
             }
 
+            
+
             if (optional) |opts| {
                 getter: {
                     if (opts.isEmptyOrUndefinedOrNull()) {
@@ -3328,6 +3330,9 @@ pub fn NewServer(comptime ssl_enabled_: bool, comptime debug_mode_: bool) type {
                                 sec_websocket_extensions = protocol;
                             }
 
+                            // we must write the status first so that 200 OK isn't written
+                            upgrader.resp.writeStatus("101 Switching Protocols");
+
                             fetch_headers.toUWSResponse(comptime ssl_enabled, upgrader.resp);
                             break :getter;
                         } else if (headers_value.isObject()) {
@@ -3340,6 +3345,9 @@ pub fn NewServer(comptime ssl_enabled_: bool, comptime debug_mode_: bool) type {
                                     sec_websocket_extensions = protocol;
                                 }
 
+                                // we must write the status first so that 200 OK isn't written
+                                upgrader.resp.writeStatus("101 Switching Protocols");
+
                                 fetch_headers.toUWSResponse(comptime ssl_enabled, upgrader.resp);
                                 fetch_headers_to_deref = fetch_headers;
                             }
diff --git a/test/bun.js/websocket-server.test.ts b/test/bun.js/websocket-server.test.ts
index 0d188fcce..0e41d3186 100644
--- a/test/bun.js/websocket-server.test.ts
+++ b/test/bun.js/websocket-server.test.ts
@@ -18,13 +18,22 @@ describe("websocket server", () => {
     var server = serve({
       port: getPort(),
       websocket: {
-        open(ws) {
-          ws.send(ws.data);
+        open(ws) {},
+        message(ws, msg) {
+          ws.send("hello world");
         },
-        message(ws, msg) {},
       },
       fetch(req, server) {
-        if (server.upgrade(req, { data: "hello world" })) {
+        if (
+          server.upgrade(req, {
+            data: "hello world",
+
+            // check that headers works
+            headers: {
+              "x-a": "text/plain",
+            },
+          })
+        ) {
           if (server.upgrade(req)) {
             throw new Error("should not upgrade twice");
           }
@@ -37,7 +46,9 @@ describe("websocket server", () => {
 
     await new Promise((resolve, reject) => {
       const websocket = new WebSocket(`ws://localhost:${server.port}`);
-
+      websocket.onopen = () => {
+        websocket.send("hello world");
+      };
       websocket.onmessage = (e) => {
         try {
           expect(e.data).toBe("hello world");