diff --git a/mocket.js.mbt b/mocket.js.mbt index 33de5b5..ebc5d1e 100644 --- a/mocket.js.mbt +++ b/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) } diff --git a/mocket.native.mbt b/mocket.native.mbt index b0c6ff0..fb3af10 100644 --- a/mocket.native.mbt +++ b/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 { diff --git a/native/mongoose/mongoose.mbt b/native/mongoose/mongoose.mbt index bb655ba..9fb9759 100644 --- a/native/mongoose/mongoose.mbt +++ b/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())