chore: adjust implementation

Author: peter-jerry-ye

src/websocket/frame.mbt

@@ -21,17 +21,14 @@ async fn[R : @io.Reader] read_frame(reader : R) -> Frame {
   let byte1 = header[1]
   let fin = (byte0.to_int() & 0x80) != 0
   let opcode_byte = byte0 & b'\x0F'
-  let opcode = match OpCode::from_byte(opcode_byte) {
-    Some(op) => op
-    None => raise FrameError("Invalid opcode: \{opcode_byte}")
-  }
+  let opcode = OpCode::from_byte(opcode_byte)
   let masked = (byte1.to_int() & 0x80) != 0
   let mut payload_len = (byte1.to_int() & 0x7F).to_int64()
 
-  // Validate payload length according to RFC 6455
+  // Validate payload length according to RFC 6455 Section 5.2
   if payload_len == 126L {
     let len_bytes = reader.read_exactly(2)
-    guard len_bytes is [u16be(len), ..]
+    guard len_bytes is [u16be(len)]
     payload_len = len.to_int64()
     if payload_len < 126L {
       raise FrameError(
@@ -40,7 +37,7 @@ async fn[R : @io.Reader] read_frame(reader : R) -> Frame {
     }
   } else if payload_len == 127L {
     let len_bytes = reader.read_exactly(8)
-    guard len_bytes is [u64be(len), ..]
+    guard len_bytes is [u64be(len)]
     payload_len = len.reinterpret_as_int64()
     if payload_len < 65536L {
       raise FrameError(
@@ -49,14 +46,12 @@ async fn[R : @io.Reader] read_frame(reader : R) -> Frame {
     }
     if payload_len < 0L {
       raise FrameError(
-        "Payload length too large (negative when interpreted as signed)",
+        "Invalid payload length: MSB must be 0 for 64-bit length",
       )
     }
   }
-
-  // Check for reasonable payload size limit (1MB for now)
-  if payload_len > 1048576L {
-    raise FrameError("Payload too large: \{payload_len} bytes (max 1MB)")
+  if payload_len > @int.max_value.to_int64() {
+    raise FrameError("Payload too large: \{payload_len} bytes")
   }
 
   // Read masking key if present
@@ -89,13 +84,6 @@ async fn[W : @io.Writer] write_frame(
   masked : Bool,
 ) -> Unit {
   let payload_len = payload.length().to_int64()
-
-  // Validate payload size
-  if payload_len > 1048576L {
-    raise FrameError(
-      "Payload too large for sending: \{payload_len} bytes (max 1MB)",
-    )
-  }
   let mut header_len = 2
 
   // Calculate extended length size
@@ -105,45 +93,37 @@ async fn[W : @io.Writer] write_frame(
     header_len += 8
   }
 
-  // Add mask size if needed
-  if masked {
-    header_len += 4
-  }
-
   // Build header
-  let header = FixedArray::make(header_len, b'\x00')
-  let mut offset = 0
+  let header = if masked {
+    FixedArray::make(header_len + 4, b'\x00')
+  } else {
+    FixedArray::make(header_len, b'\x00')
+  }
 
   // First byte: FIN + opcode
-  header[offset] = if fin {
+  header[0] = if fin {
     (0x80 | opcode.to_byte().to_int()).to_byte()
   } else {
     opcode.to_byte()
   }
-  offset += 1
 
   // Second byte: MASK + payload length
   let mask_bit = if masked { 0x80 } else { 0 }
   if payload_len < 126L {
-    header[offset] = (mask_bit | payload_len.to_int()).to_byte()
-    offset += 1
+    header[1] = (mask_bit | payload_len.to_int()).to_byte()
   } else if payload_len <= 65535L {
-    header[offset] = (mask_bit | 126).to_byte()
-    offset += 1
-    header.unsafe_write_uint16_be(offset, payload_len.to_uint16())
-    offset += 2
+    header[1] = (mask_bit | 126).to_byte()
+    header.unsafe_write_uint16_be(2, payload_len.to_uint16())
   } else {
-    header[offset] = (mask_bit | 127).to_byte()
-    offset += 1
-    header.unsafe_write_uint64_be(offset, payload_len.reinterpret_as_uint64())
-    offset += 8
+    header[1] = (mask_bit | 127).to_byte()
+    header.unsafe_write_uint64_be(2, payload_len.reinterpret_as_uint64())
   }
 
   // Add masking key and mask payload if needed
   let final_payload = if masked {
     let mask = generate_mask()
     for i = 0; i < 4; i = i + 1 {
-      header[offset + i] = mask[i]
+      header[header_len + i] = mask[i]
     }
     let payload_arr = payload.to_fixedarray()
     mask_payload(payload_arr, mask)

src/websocket/server.mbt

@@ -44,6 +44,7 @@ async fn ServerConnection::handshake(conn : @socket.Tcp) -> ServerConnection? {
     if line.is_blank() {
       break
     }
+    let line = line[:-1] // Remove trailing \r
 
     // Parse header line
     if line.contains(":") {
@@ -53,9 +54,9 @@ async fn ServerConnection::handshake(conn : @socket.Tcp) -> ServerConnection? {
         // Join remaining parts in case the value contains colons
         let value_parts = parts[1:]
         let value = if value_parts.length() == 1 {
-          value_parts[0].trim(chars=" \t\r").to_string()
+          value_parts[0].trim(chars=" \t").to_string()
         } else {
-          value_parts.map(_.trim(chars=" \t\r")).join(":").to_string()
+          value_parts.join(":").trim(chars=" \t").to_string()
         }
         // Handle multi-value headers by taking the first value
         if not(headers.contains(key)) {
@@ -85,11 +86,13 @@ async fn ServerConnection::handshake(conn : @socket.Tcp) -> ServerConnection? {
   let accept_key = generate_accept_key(key)
 
   // Send upgrade response
-  let response = "HTTP/1.1 101 Switching Protocols\r\n" +
-    "Upgrade: websocket\r\n" +
-    "Connection: Upgrade\r\n" +
-    "Sec-WebSocket-Accept: \{accept_key}\r\n" +
-    "\r\n"
+  let response =
+    $|HTTP/1.1 101 Switching Protocols\r
+    $|Upgrade: websocket\r
+    $|Connection: Upgrade\r
+    $|Sec-WebSocket-Accept: \{accept_key}\r
+    $|\r
+    $|
   conn.write(@encoding/utf8.encode(response))
   Some({ conn, closed: None })
 }
@@ -157,7 +160,7 @@ pub async fn ServerConnection::pong(
 pub async fn ServerConnection::send_close(
   self : ServerConnection,
   code? : CloseCode = Normal,
-  reason? : String = "",
+  reason? : BytesView = "",
 ) -> Unit {
   if self.closed is Some(c) {
     raise ConnectionClosed(c)
@@ -172,11 +175,7 @@ pub async fn ServerConnection::send_close(
 
   // Encode reason
   if reason != "" {
-    let reason_bytes = @encoding/utf8.encode(reason)
-    let reason_arr = reason_bytes.to_fixedarray()
-    for i = 0; i < reason_arr.length(); i = i + 1 {
-      payload[2 + i] = reason_arr[i]
-    }
+    payload.blit_from_bytesview(2, reason)
   }
   write_frame(
     self.conn,
@@ -205,28 +204,26 @@ pub async fn ServerConnection::receive(self : ServerConnection) -> Message {
       OpCode::Close => {
         // Parse close code and reason
         let mut close_code = Normal
-        let mut reason = ""
         if frame.payload.length() >= 2 {
           let payload_arr = frame.payload.to_fixedarray()
           let code_int = (payload_arr[0].to_int() << 8) |
             payload_arr[1].to_int()
           close_code = CloseCode::from_int(code_int).unwrap_or(Normal)
           if frame.payload.length() > 2 {
-            let reason_bytes = FixedArray::make(
-              frame.payload.length() - 2,
-              b'\x00',
-            )
-            for i = 2; i < frame.payload.length(); i = i + 1 {
-              reason_bytes[i - 2] = payload_arr[i]
-            }
-            reason = @encoding/utf8.decode_lossy(
-              reason_bytes.unsafe_reinterpret_as_bytes(),
-            )
+            // As per spec https://datatracker.ietf.org/doc/html/rfc6455#autoid-27
+            // The data is not guaranteed to be human readable
+            // So we do not decode it here
+            // And we are not using it further
+            let _reason_bytes = payload_arr.unsafe_reinterpret_as_bytes()[2:]
+
           }
         }
-        // Echo the close frame back and close
-        self.send_close(code=close_code, reason~)
-        self.closed = Some(close_code)
+        // 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)
+          self.closed = Some(close_code)
+        }
         raise ConnectionClosed(close_code)
       }
       OpCode::Ping => {

src/websocket/types.mbt

@@ -36,15 +36,15 @@ fn OpCode::to_byte(self : OpCode) -> Byte {
 }
 
 ///|
-fn OpCode::from_byte(byte : Byte) -> OpCode? {
+fn OpCode::from_byte(byte : Byte) -> OpCode raise {
   match byte {
-    b'\x00' => Some(Continuation)
-    b'\x01' => Some(Text)
-    b'\x02' => Some(Binary)
-    b'\x08' => Some(Close)
-    b'\x09' => Some(Ping)
-    b'\x0A' => Some(Pong)
-    _ => None
+    b'\x00' => Continuation
+    b'\x01' => Text
+    b'\x02' => Binary
+    b'\x08' => Close
+    b'\x09' => Ping
+    b'\x0A' => Pong
+    _ => raise FrameError("Invalid opcode byte: \{byte}")
   }
 }