From 0d8d39ad223d42ab10e7478d1024fc2f2abf2e85 Mon Sep 17 00:00:00 2001
From: Greg Heartsfield <scsibug@imap.cc>
Date: Sat, 17 Dec 2022 09:27:29 -0600
Subject: [PATCH] feat: add rate limiting setting for subscription creation

---
 config.toml   |  5 +++++
 src/config.rs |  4 +++-
 src/server.rs | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/config.toml b/config.toml
index e179e89..28ad0eb 100644
--- a/config.toml
+++ b/config.toml
@@ -65,6 +65,11 @@ reject_future_seconds = 1800
 # an integer.  If not set (or set to 0), defaults to unlimited.
 #messages_per_sec = 0
 
+# Limit client subscriptions created per second, averaged over one
+# minute.  Must be an integer.  If not set (or set to 0), defaults to
+# unlimited.
+#subscriptions_per_min = 0
+
 # Limit the maximum size of an EVENT message.  Defaults to 128 KB.
 # Set to 0 for unlimited.
 #max_event_bytes = 131072
diff --git a/src/config.rs b/src/config.rs
index 383da24..53fa7c3 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -52,7 +52,8 @@ pub struct Retention {
 #[allow(unused)]
 pub struct Limits {
     pub messages_per_sec: Option<u32>, // Artificially slow down event writing to limit disk consumption (averaged over 1 minute)
-    pub max_event_bytes: Option<usize>, // Maximum size of an EVENT message
+    pub subscriptions_per_min: Option<u32>, // Artificially slow down request (db query) creation to prevent abuse (averaged over 1 minute)
+    pub max_event_bytes: Option<usize>,     // Maximum size of an EVENT message
     pub max_ws_message_bytes: Option<usize>,
     pub max_ws_frame_bytes: Option<usize>,
     pub broadcast_buffer: usize, // events to buffer for subscribers (prevents slow readers from consuming memory)
@@ -214,6 +215,7 @@ impl Default for Settings {
             },
             limits: Limits {
                 messages_per_sec: None,
+                subscriptions_per_min: None,
                 max_event_bytes: Some(2 << 17),      // 128K
                 max_ws_message_bytes: Some(2 << 17), // 128K
                 max_ws_frame_bytes: Some(2 << 17),   // 128K
diff --git a/src/server.rs b/src/server.rs
index b8d72c2..47a2be3 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -14,6 +14,7 @@ use crate::notice::Notice;
 use crate::subscription::Subscription;
 use futures::SinkExt;
 use futures::StreamExt;
+use governor::{Jitter, Quota, RateLimiter};
 use http::header::HeaderMap;
 use hyper::header::ACCEPT;
 use hyper::service::{make_service_fn, service_fn};
@@ -438,6 +439,19 @@ async fn nostr_server(
     let mut bcast_rx = broadcast.subscribe();
     // Track internal client state
     let mut conn = conn::ClientConn::new(client_info.remote_ip);
+    // subscription creation rate limiting
+    let mut sub_lim_opt = None;
+    // 100ms jitter when the rate limiter returns
+    let jitter = Jitter::up_to(Duration::from_millis(100));
+    let sub_per_min_setting = settings.limits.subscriptions_per_min;
+    if let Some(sub_per_min) = sub_per_min_setting {
+        if sub_per_min > 0 {
+            trace!("Rate limits for sub creation ({}/min)", sub_per_min);
+            let quota_time = core::num::NonZeroU32::new(sub_per_min).unwrap();
+            let quota = Quota::per_minute(quota_time);
+            sub_lim_opt = Some(RateLimiter::direct(quota));
+        }
+    }
     // Use the remote IP as the client identifier
     let cid = conn.get_client_prefix();
     // Create a channel for receiving query results from the database.
@@ -606,11 +620,15 @@ async fn nostr_server(
                     Ok(NostrMessage::SubMsg(s)) => {
                         debug!("subscription requested (cid: {}, sub: {:?})", cid, s.id);
                         // subscription handling consists of:
+                        // * check for rate limits
                         // * registering the subscription so future events can be matched
                         // * making a channel to cancel to request later
                         // * sending a request for a SQL query
             // Do nothing if the sub already exists.
             if !current_subs.contains(&s) {
+                if let Some(ref lim) = sub_lim_opt {
+                    lim.until_ready_with_jitter(jitter).await;
+                }
                 current_subs.push(s.clone());
                             let (abandon_query_tx, abandon_query_rx) = oneshot::channel::<()>();
                             match conn.subscribe(s.clone()) {