fix: handle close code on ws client properly

Author: peter-jerry-ye

examples/websocket_client/main.mbt

@@ -13,58 +13,102 @@
 // limitations under the License.
 
 ///|
-/// WebSocket client example
+/// WebSocket client example - Autobahn Test Suite Client
 /// 
-/// This demonstrates how to connect to a WebSocket server,
-/// send messages, and receive responses.
+/// This implements the Autobahn test suite client logic:
+/// 1. Get case count
+/// 2. Run each test case (echo all messages)
+/// 3. Update reports
 async fn main {
-  connect_to_echo_server()
+  run_autobahn_tests()
 }
 
 ///|
-pub async fn connect_to_echo_server() -> Unit {
-  println("Connecting to WebSocket echo server at localhost:8080")
+/// Run Autobahn WebSocket test suite
+pub async fn run_autobahn_tests() -> Unit {
+  let host = @sys.get_env_vars()
+    .get("WS_TEST_HOST")
+    .unwrap_or("127.0.0.1".to_string())
+  let port = 9001
+  let agent = "moonbit-async-websocket"
 
-  // Connect to the server
-  let client = @websocket.Client::connect("0.0.0.0", "/", port=8080)
-  println("Connected successfully!")
+  // Step 1: Get case count
+  @stdio.stdout.write("Getting case count from \{host}:\{port}...\n")
+  let case_count = get_case_count(host, port)
+  @stdio.stdout.write("Ok, will run \{case_count} cases\n\n")
 
-  // Send some test messages
-  let test_messages = [
-    "Hello, WebSocket!", "This is a test message", "MoonBit WebSocket client works!",
-    "Final message",
-  ]
-  for message in test_messages {
-    // Send text message
-    println("Sending: \{message}")
-    client.send_text(message)
+  // Step 2: Run each test case
+  for case_id = 1; case_id <= case_count; case_id = case_id + 1 {
+    @stdio.stdout.write(
+      "Running test case \{case_id}/\{case_count} as user agent \{agent}\n",
+    )
+    run_test_case(host, port, case_id, agent)
+  }
+
+  // Step 3: Update reports
+  @stdio.stdout.write("\nUpdating reports...\n")
+  update_reports(host, port, agent)
+  @stdio.stdout.write("All tests completed!\n")
+}
 
-    // Receive echo response
-    let response = client.receive()
-    match response {
-      @websocket.Text(text) => println("Received: \{text}")
-      @websocket.Binary(data) =>
-        println("Received binary data (\{data.length()} bytes)")
+///|
+/// Get the total number of test cases
+async fn get_case_count(host : String, port : Int) -> Int {
+  let client = @websocket.Client::connect(host, "/getCaseCount", port~)
+  let message = client.receive()
+  client.close()
+  match message {
+    @websocket.Message::Text(text) => {
+      let count_str = text.to_string()
+      @strconv.parse_int(count_str)
+    }
+    _ => {
+      @stdio.stdout.write("Error: Expected text message with case count\n")
+      0
     }
   }
+}
 
-  // Small delay between messages
-  // Note: In a real implementation, you might want to add a sleep function
-  // For now, we'll just continue immediately
-
-  // Test binary message
-  println("Sending binary data...")
-  let binary_data = @encoding/utf8.encode("Binary test data")
-  client.send_binary(binary_data)
-  let binary_response = client.receive()
-  match binary_response {
-    @websocket.Text(text) => println("Received text response: \{text}")
-    @websocket.Binary(data) =>
-      println("Received binary response (\{data.length()} bytes)")
+///|
+/// Run a single test case - echo all messages back to server
+async fn run_test_case(
+  host : String,
+  port : Int,
+  case_id : Int,
+  agent : String,
+) -> Unit {
+  let path = "/runCase?case=\{case_id}&agent=\{agent}"
+  let client = @websocket.Client::connect(host, path, port~)
+  for {
+    let message = client.receive() catch {
+      @websocket.ConnectionClosed(_, _) =>
+        // Test case completed
+        break
+      err => {
+        @stdio.stdout.write("Error in case \{case_id}: \{err}\n")
+        break
+      }
+    }
+    // Echo the message back (core test logic)
+    match message {
+      @websocket.Message::Text(text) => client.send_text(text)
+      @websocket.Message::Binary(data) => client.send_binary(data)
+    }
   }
+  client.close()
+}
 
-  // Close the connection
-  println("Closing connection...")
-  client.send_close()
-  println("Client example completed")
+///|
+/// Update test reports on the server
+async fn update_reports(host : String, port : Int, agent : String) -> Unit {
+  let path = "/updateReports?agent=\{agent}"
+  let client = @websocket.Client::connect(host, path, port~)
+  // Wait for server to close the connection
+  ignore(
+    client.receive() catch {
+      @websocket.ConnectionClosed(_, _) => @websocket.Message::Text("")
+      _ => @websocket.Message::Text("")
+    },
+  )
+  client.close()
 }

examples/websocket_client/moon.pkg.json

@@ -1,7 +1,10 @@
 {
   "import": [
     "moonbitlang/async/websocket",
-    "moonbitlang/async"
+    "moonbitlang/async",
+    "moonbitlang/async/stdio",
+    "moonbitlang/async/io",
+    "moonbitlang/x/sys"
   ],
   "is-main": true
 }

src/websocket/client.mbt

@@ -58,7 +58,12 @@ pub async fn Client::connect(
   let key = base64_encode(nonce.unsafe_reinterpret_as_bytes())
   let request = "GET \{path} HTTP/1.1\r\n"
   conn.write(request)
-  conn.write("Host: \{host}\r\n")
+  let host_header = if port == 80 || port == 443 {
+    host
+  } else {
+    "\{host}:\{port}"
+  }
+  conn.write("Host: \{host_header}\r\n")
   conn.write("Upgrade: websocket\r\n")
   conn.write("Connection: Upgrade\r\n")
   conn.write("Sec-WebSocket-Key: \{key}\r\n")
@@ -149,17 +154,17 @@ pub async fn Client::send_close(
   code? : CloseCode = Normal,
   reason? : String,
 ) -> Unit {
-  if self.closed is Some(_) {
-    return
+  if self.closed is Some(e) {
+    raise e
   }
-  let mut payload = FixedArray::make(0, b'\x00')
   let code_int = code.to_int()
-  let reason_bytes = if reason is Some(r) {
-    @encoding/utf8.encode(r)
-  } else {
-    b""
+  let reason_bytes = @encoding/utf8.encode(reason.unwrap_or(""))
+  if reason_bytes.length() > 123 {
+    // Close reason too long
+    // TODO: should we close the connection anyway?
+    fail("Close reason too long")
   }
-  payload = FixedArray::make(2 + reason_bytes.length(), b'\x00')
+  let payload = FixedArray::make(2 + reason_bytes.length(), b'\x00')
   payload.unsafe_write_uint16_be(0, code_int.to_uint16())
   payload.blit_from_bytesview(2, reason_bytes)
   write_frame(
@@ -169,12 +174,16 @@ pub async fn Client::send_close(
     payload.unsafe_reinterpret_as_bytes(),
     self.rand.int().to_be_bytes(),
   )
+  // Wait until the server acknowledges the close
+  ignore(read_frame(self.conn)) catch {
+    _ => ()
+  }
   self.closed = Some(ConnectionClosed(code, reason))
 }
 
 ///|
 /// Send a text message
-pub async fn Client::send_text(self : Client, text : String) -> Unit {
+pub async fn Client::send_text(self : Client, text : StringView) -> Unit {
   if self.closed is Some(code) {
     raise code
   }
@@ -190,7 +199,7 @@ pub async fn Client::send_text(self : Client, text : String) -> Unit {
 
 ///|
 /// Send a binary message
-pub async fn Client::send_binary(self : Client, data : Bytes) -> Unit {
+pub async fn Client::send_binary(self : Client, data : BytesView) -> Unit {
   if self.closed is Some(code) {
     raise code
   }
@@ -247,67 +256,127 @@ pub async fn Client::receive(self : Client) -> Message {
   }
   let frames : Array[Frame] = []
   let mut first_opcode : OpCode? = None
-  for {
+  while self.closed is None {
     let frame = read_frame(self.conn)
 
     // Handle control frames immediately
     match frame.opcode {
-      OpCode::Close => {
+      Close => {
         // Parse close code and reason
+        // Ref: https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.1
         let mut close_code = Normal
         let mut reason : String? = None
-        if frame.payload is [u16be(code), .. data] {
+        if frame.payload is [u16be(code), .. rest] {
           close_code = CloseCode::from_int(code.reinterpret_as_int())
-          reason = Some(
-            @encoding/utf8.decode(data) catch {
-              _ => {
-                close_code = ProtocolError
-                ""
-              }
-            },
-          )
+          reason = Some(@encoding/utf8.decode(rest)) catch {
+            _ => {
+              // Invalid reason, fail fast
+              close_code = ProtocolError
+              None
+            }
+          }
+        } else {
+          guard frame.payload is [] else {
+            // Invalid close payload
+            close_code = ProtocolError
+          }
         }
-        self.closed = Some(ConnectionClosed(close_code, reason))
-        raise ConnectionClosed(close_code, reason)
+        // If we didn't send close first, respond with close
+        if self.closed is None {
+          // Echo the close frame back and close
+          self.send_close(code=close_code, reason?) catch {
+            _ => ()
+          }
+        }
+        continue
       }
-      OpCode::Ping => {
+      Ping =>
         // Auto-respond to ping with pong
         self.pong(data=frame.payload)
-        continue
-      }
-      OpCode::Pong =>
+      Pong =>
         // Ignore pong frames
-        continue
-      _ => ()
-    }
-
-    // Track the first opcode for message type
-    if first_opcode is None {
-      first_opcode = Some(frame.opcode)
-    }
-    frames.push(frame)
-
-    // If this is the final frame, assemble the message
-    if frame.fin {
-      break
-    }
-  }
-
-  // Assemble message from frames
-  let total_size = frames.fold(init=0, fn(acc, f) { acc + f.payload.length() })
-  let data = FixedArray::make(total_size, b'\x00')
-  let mut offset = 0
-  for frame in frames {
-    let payload_arr = frame.payload.to_fixedarray()
-    for i = 0; i < payload_arr.length(); i = i + 1 {
-      data[offset + i] = payload_arr[i]
+        ()
+      Text =>
+        if first_opcode is Some(_) {
+          // Ref https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
+          // We don't have extensions, so fragments MUST NOT be interleaved
+          self.send_close(code=ProtocolError) catch {
+            _ => ()
+          }
+        } else if frame.fin {
+          // Single-frame text message
+          return Message::Text(@encoding/utf8.decode(frame.payload)) catch {
+            _ => {
+              // Ref https://datatracker.ietf.org/doc/html/rfc6455#section-8.1
+              // We MUST Fail the WebSocket Connection if the payload is not
+              // valid UTF-8
+              self.send_close(code=InvalidFramePayload) catch {
+                _ => ()
+              }
+              continue
+            }
+          }
+        } else {
+          first_opcode = Some(Text)
+          // Start of fragmented text message
+          frames.push(frame)
+        }
+      Binary =>
+        if first_opcode is Some(_) {
+          // Ref https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
+          // We don't have extensions, so fragments MUST NOT be interleaved
+          self.send_close(code=ProtocolError) catch {
+            _ => ()
+          }
+        } else if frame.fin {
+          // Single-frame binary message
+          return Message::Binary(frame.payload)
+        } else {
+          first_opcode = Some(Binary)
+          frames.push(frame)
+        }
+      Continuation => {
+        if first_opcode is None {
+          // Continuation frame without a starting frame
+          self.send_close(code=ProtocolError) catch {
+            _ => ()
+          }
+          continue
+        }
+        frames.push(frame)
+        if frame.fin {
+          // Final fragment received, assemble message
+          let total_size = frames.fold(init=0, fn(acc, f) {
+            acc + f.payload.length()
+          })
+          let data = FixedArray::make(total_size, b'\x00')
+          let mut offset = 0
+          for f in frames {
+            data.blit_from_bytes(offset, f.payload, 0, f.payload.length())
+            offset += f.payload.length()
+          }
+          let message_data = data.unsafe_reinterpret_as_bytes()
+          match first_opcode {
+            Some(Text) => {
+              let text = @encoding/utf8.decode(message_data) catch {
+                _ => {
+                  self.send_close(code=InvalidFramePayload) catch {
+                    _ => ()
+                  }
+                  continue
+                }
+              }
+              return Message::Text(text)
+            }
+            Some(Binary) => return Message::Binary(message_data)
+            _ => panic()
+          }
+          // Reset for next message
+          frames.clear()
+          first_opcode = None
+        }
+      }
     }
-    offset += payload_arr.length()
-  }
-  let message_data = data.unsafe_reinterpret_as_bytes()
-  match first_opcode {
-    Some(OpCode::Text) => Text(@encoding/utf8.decode_lossy(message_data))
-    Some(OpCode::Binary) => Binary(message_data)
-    _ => Binary(message_data)
   }
+  raise self.closed.unwrap()
 }