pool: allow to set event limits per kind

Add `max_size_per_kind` and `max_num_tags_per_kind` hashmaps, to set specific limits per kind. Add methods to get limits per kind and automatically fallback to generic ones. Increase default kind 3 event limit to `840000` bytes and `10000` tags. Closes #430 Signed-off-by: Yuki Kishimoto <[email protected]>
rust-nostr · May 23, 2024 · d6561f1 · d6561f1
1 parent e658a1b
commit d6561f1
Show file tree

Hide file tree

Showing 8 changed files with 134 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,9 +7,18 @@
 
 ## [Unreleased]
 
+### Summary
+
+### Changed
+
 ### Added
 
 * nostr: add `Tag::is_root` method ([Xiao Yu])
+* pool: allow to set event limits per kind ([Yuki Kishimoto])
+
+### Fixed
+
+### Removed
 
 ## [v0.31.0]
 

diff --git a/bindings/nostr-sdk-ffi/src/client/options.rs b/bindings/nostr-sdk-ffi/src/client/options.rs
@@ -110,7 +110,7 @@ impl Options {
     /// Set custom relay limits
     pub fn relay_limits(self: Arc<Self>, limits: &RelayLimits) -> Self {
         let mut builder = unwrap_or_clone_arc(self);
-        builder.inner = builder.inner.relay_limits(**limits);
+        builder.inner = builder.inner.relay_limits(limits.deref().clone());
         builder
     }
 }

diff --git a/bindings/nostr-sdk-ffi/src/relay/limits.rs b/bindings/nostr-sdk-ffi/src/relay/limits.rs
@@ -6,6 +6,7 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use nostr_ffi::helper::unwrap_or_clone_arc;
+use nostr_ffi::Kind;
 use nostr_sdk::pool::relay;
 use uniffi::Object;
 
@@ -56,10 +57,31 @@ impl RelayLimits {
         builder
     }
 
+    /// Maximum size per kind of normalised JSON, in bytes.
+    pub fn event_max_size_per_kind(self: Arc<Self>, kind: &Kind, max_size: Option<u32>) -> Self {
+        let mut builder = unwrap_or_clone_arc(self);
+        builder.inner.events = builder.inner.events.set_max_size_per_kind(**kind, max_size);
+        builder
+    }
+
     /// Maximum number of tags allowed (default: 2_000)
     pub fn event_max_num_tags(self: Arc<Self>, max_num_tags: Option<u16>) -> Self {
         let mut builder = unwrap_or_clone_arc(self);
         builder.inner.events.max_num_tags = max_num_tags;
         builder
     }
+
+    /// Maximum number of tags allowed per kind
+    pub fn event_max_num_tags_per_kind(
+        self: Arc<Self>,
+        kind: &Kind,
+        max_num_tags: Option<u16>,
+    ) -> Self {
+        let mut builder = unwrap_or_clone_arc(self);
+        builder.inner.events = builder
+            .inner
+            .events
+            .set_max_num_tags_per_kind(**kind, max_num_tags);
+        builder
+    }
 }
diff --git a/bindings/nostr-sdk-js/src/client/options.rs b/bindings/nostr-sdk-js/src/client/options.rs
@@ -92,6 +92,6 @@ impl JsOptions {
     /// Set custom relay limits
     #[wasm_bindgen(js_name = relayLimits)]
     pub fn relay_limits(self, limits: &JsRelayLimits) -> Self {
-        self.inner.relay_limits(**limits).into()
+        self.inner.relay_limits(limits.deref().clone()).into()
     }
 }
diff --git a/bindings/nostr-sdk-js/src/relay/limits.rs b/bindings/nostr-sdk-js/src/relay/limits.rs
@@ -52,10 +52,30 @@ impl JsRelayLimits {
         self
     }
 
+    /// Maximum size per kind of normalised JSON, in bytes
+    #[wasm_bindgen(js_name = eventMaxSizePerKind)]
+    pub fn event_max_size_per_kind(mut self, kind: u16, max_size: Option<u32>) -> Self {
+        self.inner.events = self
+            .inner
+            .events
+            .set_max_size_per_kind(Kind::from(kind), max_size);
+        self
+    }
+
     /// Maximum number of tags allowed (default: 2_000)
     #[wasm_bindgen(js_name = eventMaxNumTags)]
     pub fn event_max_num_tags(mut self, max_num_tags: Option<u16>) -> Self {
         self.inner.events.max_num_tags = max_num_tags;
         self
     }
+
+    /// Maximum number of tags per kind allowed
+    #[wasm_bindgen(js_name = eventMaxNumTagsPerKind)]
+    pub fn event_max_num_tags_per_kind(mut self, kind: u16, max_num_tags: Option<u16>) -> Self {
+        self.inner.events = self
+            .inner
+            .events
+            .set_max_num_tags_per_kind(Kind::from(kind), max_num_tags);
+        self
+    }
 }
diff --git a/crates/nostr-relay-pool/src/relay/internal.rs b/crates/nostr-relay-pool/src/relay/internal.rs
@@ -23,7 +23,7 @@ use nostr::nips::nip01::Coordinate;
 use nostr::nips::nip11::RelayInformationDocument;
 use nostr::secp256k1::rand::{self, Rng};
 use nostr::{
-    ClientMessage, Event, EventId, Filter, JsonUtil, Keys, MissingPartialEvent, PartialEvent,
+    ClientMessage, Event, EventId, Filter, JsonUtil, Keys, Kind, MissingPartialEvent, PartialEvent,
     RawRelayMessage, RelayMessage, SubscriptionId, Timestamp, Url,
 };
 use nostr_database::{DynNostrDatabase, Order};
@@ -733,8 +733,10 @@ impl InternalRelay {
                 subscription_id,
                 event,
             } => {
+                let kind: Kind = Kind::from(event.kind);
+
                 // Check event size
-                if let Some(max_size) = self.opts.limits.events.max_size {
+                if let Some(max_size) = self.opts.limits.events.get_max_size(&kind) {
                     let size: usize = event.as_json().as_bytes().len();
                     let max_size: usize = max_size as usize;
                     if size > max_size {
@@ -743,7 +745,7 @@ impl InternalRelay {
                 }
 
                 // Check tags limit
-                if let Some(max_num_tags) = self.opts.limits.events.max_num_tags {
+                if let Some(max_num_tags) = self.opts.limits.events.get_max_num_tags(&kind) {
                     let size: usize = event.tags.len();
                     let max_num_tags: usize = max_num_tags as usize;
                     if size > max_num_tags {

diff --git a/crates/nostr-relay-pool/src/relay/limits.rs b/crates/nostr-relay-pool/src/relay/limits.rs
@@ -4,8 +4,15 @@
 
 //! Relay limits
 
+use std::collections::HashMap;
+
+use nostr::Kind;
+
+const MAX_EVENT_SIZE: u32 = 70_000; // bytes
+const MAX_CONTACT_LIST_EVENT_SIZE: u32 = 840_000; // bytes
+
 /// Relay limits
-#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
 pub struct RelayLimits {
     /// Message limits
     pub messages: RelayMessageLimits,
@@ -48,19 +55,31 @@ impl RelayMessageLimits {
 }
 
 /// Events limits
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RelayEventLimits {
     /// Maximum size of normalised JSON, in bytes (default: 70_000)
     pub max_size: Option<u32>,
+    /// Maximum size of normalized JSON per [Kind], in bytes.
+    pub max_size_per_kind: HashMap<Kind, Option<u32>>,
     /// Maximum number of tags allowed (default: 2_000)
     pub max_num_tags: Option<u16>,
+    /// Maximum number of tags allowed per [Kind].
+    pub max_num_tags_per_kind: HashMap<Kind, Option<u16>>,
 }
 
 impl Default for RelayEventLimits {
     fn default() -> Self {
+        let mut max_size_per_kind: HashMap<Kind, Option<u32>> = HashMap::with_capacity(1);
+        max_size_per_kind.insert(Kind::ContactList, Some(MAX_CONTACT_LIST_EVENT_SIZE));
+
+        let mut max_num_tags_per_kind: HashMap<Kind, Option<u16>> = HashMap::with_capacity(1);
+        max_num_tags_per_kind.insert(Kind::ContactList, Some(10_000));
+
         Self {
-            max_size: Some(70_000),
+            max_size: Some(MAX_EVENT_SIZE),
+            max_size_per_kind,
             max_num_tags: Some(2_000),
+            max_num_tags_per_kind,
         }
     }
 }
@@ -71,7 +90,60 @@ impl RelayEventLimits {
     pub fn disable() -> Self {
         Self {
             max_size: None,
+            max_size_per_kind: HashMap::new(),
             max_num_tags: None,
+            max_num_tags_per_kind: HashMap::new(),
         }
     }
+
+    /// Add/Edit max size per [Kind]
+    pub fn set_max_size_per_kind(mut self, kind: Kind, max_size: Option<u32>) -> Self {
+        self.max_size_per_kind.insert(kind, max_size);
+        self
+    }
+
+    /// Add/Edit max number of tags per [Kind]
+    pub fn set_max_num_tags_per_kind(mut self, kind: Kind, max_num_tags: Option<u16>) -> Self {
+        self.max_num_tags_per_kind.insert(kind, max_num_tags);
+        self
+    }
+
+    /// Get max size for [Kind]
+    ///
+    /// Fallback to `max_size` if no limit is specified for [Kind]
+    pub fn get_max_size(&self, kind: &Kind) -> Option<u32> {
+        match self.max_size_per_kind.get(kind).copied() {
+            Some(limit) => limit,
+            None => self.max_size,
+        }
+    }
+
+    /// Get max number of tags allowed for [Kind]
+    ///
+    /// Fallback to `max_num_tags` if no limit is specified for [Kind]
+    pub fn get_max_num_tags(&self, kind: &Kind) -> Option<u16> {
+        match self.max_num_tags_per_kind.get(kind).copied() {
+            Some(limit) => limit,
+            None => self.max_num_tags,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_event_limits_get_max_size() {
+        let limits = RelayLimits::default();
+
+        assert_eq!(
+            limits.events.get_max_size(&Kind::TextNote),
+            Some(MAX_EVENT_SIZE)
+        );
+        assert_eq!(
+            limits.events.get_max_size(&Kind::ContactList),
+            Some(MAX_CONTACT_LIST_EVENT_SIZE)
+        );
+    }
 }
diff --git a/crates/nostr-sdk/src/client/mod.rs b/crates/nostr-sdk/src/client/mod.rs
@@ -408,7 +408,7 @@ impl Client {
         // Set min POW difficulty and limits
         let opts: RelayOptions = opts
             .pow(self.opts.get_min_pow_difficulty())
-            .limits(self.opts.relay_limits);
+            .limits(self.opts.relay_limits.clone());
 
         // Add relay
         self.add_relay_with_opts::<Url>(url, opts).await