fix: unbounded WebSocket message reads allow memory exhaustion

mocket.js.mbt

@@ -461,6 +461,13 @@ pub fn __ws_emit_js_port(
     }
     _ => ()
   }
+  let max = mocket.max_body_size
+  if max > 0 && payload.length() > max {
+    match event_type {
+      "message" | "binary" => return
+      _ => ()
+    }
+  }
   let peer = WebSocketPeer::{ connection_id, subscribed_channels: [] }
   dispatch_ws_event(handler, peer, event_type, payload)
 }

mocket.native.mbt

@@ -434,6 +434,32 @@ async fn read_body_limited(reader : &@io.Reader, max_size : Int) -> Bytes {
   buf.to_bytes()
 }
 
+///|
+async fn read_ws_limited(reader : &@io.Reader, max_size : Int) -> Bytes? {
+  if max_size <= 0 {
+    return Some(reader.read_all().binary())
+  }
+  let buf = Buffer()
+  let chunk_size = 8192
+  let chunk = FixedArray::make(chunk_size, b'\x00')
+  for total = 0 {
+    let n = reader.read(chunk) catch { _ => return None }
+    if n <= 0 {
+      break
+    }
+    if total + n > max_size {
+      return None
+    }
+    let arr : Array[Byte] = []
+    for i = 0; i < n; i = i + 1 {
+      arr.push(chunk[i])
+    }
+    buf.write_bytes(Bytes::from_array(arr))
+    continue total + n
+  }
+  Some(buf.to_bytes())
+}
+
 ///|
 async fn handle_websocket_request(
   port : Int,
@@ -454,11 +480,17 @@ async fn handle_websocket_request(
         for ;; {
           let msg = ws.recv()
           match msg.kind {
-            Text => {
-              let text = msg.read_all().text() catch { _ => "" }
-              handler(Message(peer, Text(text)))
-            }
-            Binary => handler(Message(peer, Binary(msg.read_all().binary())))
+            Text =>
+              match read_ws_limited(msg, mocket.max_body_size) {
+                Some(data) =>
+                  handler(Message(peer, Text(@utf8.decode_lossy(data))))
+                None => ()
+              }
+            Binary =>
+              match read_ws_limited(msg, mocket.max_body_size) {
+                Some(data) => handler(Message(peer, Binary(data)))
+                None => ()
+              }
           }
         }
       } catch {

native/mongoose/mongoose.mbt

@@ -104,6 +104,9 @@ let server_map : Map[Int, @mocket.Mocket] = Map([])
 ///|
 let ws_handler_map : Map[Int, @mocket.WebSocketHandler] = Map([])
 
+///|
+let ws_max_body_size : Ref[Int] = Ref(1048576)
+
 ///|
 fn[T : Show] to_cbytes(s : T) -> Bytes {
   @utf8.encode(s.to_string())
@@ -158,6 +161,7 @@ fn first_ws_handler() -> @mocket.WebSocketHandler {
 
 ///|
 fn register_ws_handlers(mocket : @mocket.Mocket, port : Int) -> Unit {
+  ws_max_body_size.val = mocket.max_body_size
   let mut done = false
   mocket.ws_static_routes.each(fn(_, handler) {
     if !done {
@@ -263,9 +267,19 @@ pub fn __ws_emit(
       )
       handler(Open(peer))
     }
-    "message" => handler(Message(peer, Text(from_cbytes(payload))))
+    "message" => {
+      let max = ws_max_body_size.val
+      if max > 0 && payload.length() > max {
+        return
+      }
+      handler(Message(peer, Text(from_cbytes(payload))))
+    }
     "binary" => {
       let len = ws_msg_body_len()
+      let max = ws_max_body_size.val
+      if max > 0 && len > max {
+        return
+      }
       let arr = Array::make(len, b'\x00')
       let buf = Bytes::from_array(arr)
       let copied = ws_msg_copy(buf, buf.length())