live.rs (10309B)
1 #![forbid(unsafe_code)] 2 3 use crate::{errors::BaseRelayError, relay::filter::BaseRelayMatchedFilterContext}; 4 use std::collections::BTreeMap; 5 use tangle_groups::GroupAuthContext; 6 use tangle_protocol::SubscriptionId; 7 use tangle_store_pocket::{PocketEvent, PocketFilter, PocketOwnedFilter}; 8 9 #[derive(Debug, Clone, PartialEq, Eq)] 10 pub(crate) struct LiveSubscriptionSet { 11 subscriptions: BTreeMap<SubscriptionId, LiveSubscription>, 12 max_subscriptions: usize, 13 } 14 15 #[derive(Debug, Clone, PartialEq, Eq)] 16 struct LiveSubscription { 17 filters: Vec<PocketOwnedFilter>, 18 } 19 20 #[derive(Debug, Clone, PartialEq, Eq)] 21 pub(crate) struct LiveSubscriptionMatch { 22 subscription_id: SubscriptionId, 23 matched_filters: Vec<(usize, PocketOwnedFilter)>, 24 } 25 26 impl LiveSubscriptionMatch { 27 fn new( 28 subscription_id: SubscriptionId, 29 matched_filters: Vec<(usize, PocketOwnedFilter)>, 30 ) -> Self { 31 Self { 32 subscription_id, 33 matched_filters, 34 } 35 } 36 37 pub(crate) fn subscription_id(&self) -> &SubscriptionId { 38 &self.subscription_id 39 } 40 41 pub(crate) fn into_subscription_id(self) -> SubscriptionId { 42 self.subscription_id 43 } 44 45 #[cfg(test)] 46 pub(crate) fn matched_filter_context(&self) -> BaseRelayMatchedFilterContext { 47 let (filter_index, filter) = &self.matched_filters[0]; 48 BaseRelayMatchedFilterContext::from_filter(*filter_index, filter) 49 } 50 51 pub(crate) fn matched_filter_contexts(&self) -> Vec<BaseRelayMatchedFilterContext> { 52 self.matched_filters 53 .iter() 54 .map(|(filter_index, filter)| { 55 BaseRelayMatchedFilterContext::from_filter(*filter_index, filter) 56 }) 57 .collect() 58 } 59 60 pub(crate) fn filters(&self) -> impl Iterator<Item = &PocketFilter> { 61 self.matched_filters 62 .iter() 63 .map(|(_, filter)| -> &PocketFilter { filter }) 64 } 65 } 66 67 impl LiveSubscriptionSet { 68 pub(crate) fn new( 69 max_pending_events: usize, 70 max_subscriptions: usize, 71 ) -> Result<Self, BaseRelayError> { 72 if max_pending_events == 0 { 73 return Err(BaseRelayError::invalid( 74 "live subscription pending event limit must be greater than zero", 75 )); 76 } 77 if max_subscriptions == 0 { 78 return Err(BaseRelayError::invalid( 79 "live subscription count limit must be greater than zero", 80 )); 81 } 82 Ok(Self { 83 subscriptions: BTreeMap::new(), 84 max_subscriptions, 85 }) 86 } 87 88 pub(crate) fn subscribe( 89 &mut self, 90 subscription_id: SubscriptionId, 91 filters: Vec<PocketOwnedFilter>, 92 ) -> Result<(), BaseRelayError> { 93 self.ensure_can_subscribe(&subscription_id, &filters)?; 94 self.subscriptions 95 .insert(subscription_id, LiveSubscription { filters }); 96 Ok(()) 97 } 98 99 pub(crate) fn ensure_can_subscribe( 100 &self, 101 subscription_id: &SubscriptionId, 102 filters: &[PocketOwnedFilter], 103 ) -> Result<(), BaseRelayError> { 104 if filters.is_empty() { 105 return Err(BaseRelayError::invalid( 106 "subscription must include at least one filter", 107 )); 108 } 109 if !self.subscriptions.contains_key(subscription_id) 110 && self.subscriptions.len() >= self.max_subscriptions 111 { 112 return Err(BaseRelayError::invalid( 113 "connection subscription limit exceeded", 114 )); 115 } 116 Ok(()) 117 } 118 119 pub(crate) fn close(&mut self, subscription_id: &SubscriptionId) -> CloseResult { 120 if self.subscriptions.remove(subscription_id).is_some() { 121 CloseResult::Closed 122 } else { 123 CloseResult::NotFound 124 } 125 } 126 127 pub(crate) fn contains(&self, subscription_id: &SubscriptionId) -> bool { 128 self.subscriptions.contains_key(subscription_id) 129 } 130 131 pub(crate) fn close_all(&mut self) -> usize { 132 let closed = self.subscriptions.len(); 133 self.subscriptions.clear(); 134 closed 135 } 136 137 pub(crate) fn fanout( 138 &self, 139 event: &PocketEvent, 140 auth: &GroupAuthContext, 141 visible_to_auth: impl Fn(&PocketEvent, &GroupAuthContext) -> bool, 142 ) -> Result<Vec<LiveSubscriptionMatch>, BaseRelayError> { 143 self.subscriptions.iter().try_fold( 144 Vec::new(), 145 |mut matched, (subscription_id, subscription)| { 146 let mut matched_filters = Vec::new(); 147 for (filter_index, filter) in subscription.filters.iter().enumerate() { 148 if filter 149 .event_matches(event) 150 .map_err(|error| BaseRelayError::error(error.to_string()))? 151 { 152 matched_filters.push((filter_index, filter.clone())); 153 } 154 } 155 if matched_filters.is_empty() { 156 return Ok(matched); 157 } 158 if visible_to_auth(event, auth) { 159 matched.push(LiveSubscriptionMatch::new( 160 subscription_id.clone(), 161 matched_filters, 162 )); 163 } 164 Ok(matched) 165 }, 166 ) 167 } 168 169 pub(crate) fn active_count(&self) -> usize { 170 self.subscriptions.len() 171 } 172 } 173 174 #[derive(Debug, Clone, Copy, PartialEq, Eq)] 175 pub enum CloseResult { 176 Closed, 177 NotFound, 178 } 179 180 #[cfg(test)] 181 mod tests { 182 use super::{CloseResult, LiveSubscriptionSet}; 183 use tangle_groups::GroupAuthContext; 184 use tangle_protocol::{SubscriptionId, filter_from_value}; 185 use tangle_test_support::{FixtureKey, tangle_v2_event}; 186 187 #[test] 188 fn live_subscription_fanout_keeps_healthy_subscriptions_open() { 189 let mut subscriptions = LiveSubscriptionSet::new(1, 1).expect("subscriptions"); 190 let subscription_id = SubscriptionId::new("live").expect("subscription"); 191 subscriptions 192 .subscribe( 193 subscription_id.clone(), 194 vec![pocket_filter(serde_json::json!({"kinds":[1]}))], 195 ) 196 .expect("subscribe"); 197 let first = tangle_v2_event(FixtureKey::Member, 1_714_124_433, 1, Vec::new(), "first") 198 .expect("first"); 199 let second = tangle_v2_event(FixtureKey::Member, 1_714_124_434, 1, Vec::new(), "second") 200 .expect("second"); 201 let third = tangle_v2_event(FixtureKey::Member, 1_714_124_435, 1, Vec::new(), "third") 202 .expect("third"); 203 204 assert_live_match( 205 subscriptions 206 .fanout( 207 &pocket_event(&first), 208 &GroupAuthContext::unauthenticated(), 209 |_, _| true, 210 ) 211 .expect("fanout") 212 .as_slice(), 213 &subscription_id, 214 0, 215 ); 216 assert_live_match( 217 subscriptions 218 .fanout( 219 &pocket_event(&second), 220 &GroupAuthContext::unauthenticated(), 221 |_, _| true, 222 ) 223 .expect("fanout") 224 .as_slice(), 225 &subscription_id, 226 0, 227 ); 228 assert_live_match( 229 subscriptions 230 .fanout( 231 &pocket_event(&third), 232 &GroupAuthContext::unauthenticated(), 233 |_, _| true, 234 ) 235 .expect("fanout") 236 .as_slice(), 237 &subscription_id, 238 0, 239 ); 240 assert_eq!(subscriptions.close(&subscription_id), CloseResult::Closed); 241 } 242 243 #[test] 244 fn live_subscription_fanout_uses_pocket_filter_matching_and_auth_gate() { 245 let mut subscriptions = LiveSubscriptionSet::new(4, 4).expect("subscriptions"); 246 let event = tangle_v2_event( 247 FixtureKey::Member, 248 1_714_124_433, 249 1, 250 vec![tangle_protocol::Tag::from_parts("t", &["market"]).expect("tag")], 251 "first", 252 ) 253 .expect("event"); 254 let matched = SubscriptionId::new("matched").expect("subscription"); 255 let mismatched = SubscriptionId::new("mismatched").expect("subscription"); 256 subscriptions 257 .subscribe( 258 matched.clone(), 259 vec![pocket_filter(serde_json::json!({ 260 "ids": [event.id().as_str()], 261 "authors": [event.unsigned().pubkey().as_str()], 262 "kinds": [1], 263 "#t": ["market"], 264 "since": 1_714_124_433, 265 "until": 1_714_124_434 266 }))], 267 ) 268 .expect("matched subscribe"); 269 subscriptions 270 .subscribe( 271 mismatched, 272 vec![pocket_filter(serde_json::json!({"kinds":[2]}))], 273 ) 274 .expect("mismatched subscribe"); 275 let event = pocket_event(&event); 276 277 let output = subscriptions 278 .fanout(&event, &GroupAuthContext::unauthenticated(), |_, _| true) 279 .expect("fanout"); 280 assert_live_match(output.as_slice(), &matched, 0); 281 assert!( 282 subscriptions 283 .fanout(&event, &GroupAuthContext::unauthenticated(), |_, _| false) 284 .expect("auth gated fanout") 285 .is_empty() 286 ); 287 } 288 289 fn pocket_filter(value: serde_json::Value) -> tangle_store_pocket::PocketOwnedFilter { 290 let filter = filter_from_value(&value).expect("filter"); 291 crate::pocket_conversion::tangle_filter_to_pocket(&filter).expect("pocket filter") 292 } 293 294 fn assert_live_match( 295 matches: &[super::LiveSubscriptionMatch], 296 subscription_id: &SubscriptionId, 297 filter_index: usize, 298 ) { 299 assert!(matches!( 300 matches, 301 [delivered] if delivered.subscription_id() == subscription_id 302 && delivered.matched_filter_context().filter_index() == filter_index 303 )); 304 } 305 306 fn pocket_event(event: &tangle_protocol::Event) -> tangle_store_pocket::PocketOwnedEvent { 307 crate::pocket_conversion::tangle_event_to_pocket(event).expect("pocket event") 308 } 309 }