From 55bb6bd4405a7d8994eac18f3dd7d58e504190a7 Mon Sep 17 00:00:00 2001
From: Greg Heartsfield <scsibug@imap.cc>
Date: Sun, 19 Dec 2021 16:26:32 -0600
Subject: [PATCH] feat: add resource limits for websocket messages

---
 src/conn.rs |  2 +-
 src/main.rs | 10 +++++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/conn.rs b/src/conn.rs
index 6dc5ca2..1ab643b 100644
--- a/src/conn.rs
+++ b/src/conn.rs
@@ -53,7 +53,7 @@ impl ClientConn {
                 v.push(id);
             }
         }
-        return v;
+        v
     }
 
     /// Add a new subscription for this connection.
diff --git a/src/main.rs b/src/main.rs
index 53728a9..d094fae 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,6 +18,7 @@ use tokio::sync::broadcast;
 use tokio::sync::broadcast::{Receiver, Sender};
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
+use tungstenite::protocol::WebSocketConfig;
 
 /// Start running a Nostr relay server.
 fn main() -> Result<(), Error> {
@@ -93,8 +94,15 @@ async fn nostr_server(
 ) {
     // get a broadcast channel for clients to communicate on
     let mut bcast_rx = broadcast.subscribe();
+    // websocket configuration / limits
+    let config = WebSocketConfig {
+        max_send_queue: None,
+        max_message_size: Some(2 << 19), // 512K
+        max_frame_size: Some(2 << 19),   // 512k
+        accept_unmasked_frames: false,   // follow the spec
+    };
     // upgrade the TCP connection to WebSocket
-    let conn = tokio_tungstenite::accept_async(stream).await;
+    let conn = tokio_tungstenite::accept_async_with_config(stream, Some(config)).await;
     let ws_stream = conn.expect("websocket handshake error");
     // wrap websocket into a stream & sink of Nostr protocol messages
     let mut nostr_stream = protostream::wrap_ws_in_nostr(ws_stream);