limits.rs (8553B)
1 /// Limitations imposed by the relay 2 pub struct RelayLimitations { 3 // corresponds to NIP-11 `max_subscriptions` 4 pub maximum_subs: usize, 5 6 // corresponds to NIP-11 `max_message_length` 7 pub max_json_bytes: usize, 8 } 9 10 impl Default for RelayLimitations { 11 fn default() -> Self { 12 Self { 13 maximum_subs: 10, 14 max_json_bytes: 400_000, 15 } 16 } 17 } 18 19 pub struct RelayCoordinatorLimits { 20 pub sub_guardian: SubPassGuardian, 21 pub max_json_bytes: usize, 22 } 23 24 impl RelayCoordinatorLimits { 25 pub fn new(limits: RelayLimitations) -> Self { 26 Self { 27 max_json_bytes: limits.max_json_bytes, 28 sub_guardian: SubPassGuardian::new(limits.maximum_subs), 29 } 30 } 31 32 pub fn new_total(&mut self, new_max: usize) -> Option<Vec<SubPassRevocation>> { 33 let old = self.sub_guardian.total_passes; 34 35 if new_max == old { 36 return None; 37 } 38 39 if new_max > old { 40 let add = new_max - old; 41 self.sub_guardian.spawn_passes(add); 42 self.sub_guardian.total_passes = new_max; 43 return None; 44 } 45 46 // new_max < old 47 let remove = old - new_max; 48 self.sub_guardian.total_passes = new_max; 49 50 let mut pending = Vec::new(); 51 52 for _ in 0..remove { 53 let mut revocation = SubPassRevocation::new(); 54 if let Some(pass) = self.sub_guardian.available_passes.pop() { 55 // can revoke immediately -> do NOT return a revocation object for it 56 revocation.revocate(pass); 57 } else { 58 // can't revoke now -> return a revocation object to be fulfilled later 59 pending.push(revocation); 60 } 61 } 62 63 if pending.is_empty() { 64 None 65 } else { 66 Some(pending) 67 } 68 } 69 } 70 71 pub struct SubPassGuardian { 72 total_passes: usize, 73 available_passes: Vec<SubPass>, 74 } 75 76 impl SubPassGuardian { 77 pub(crate) fn new(max_subs: usize) -> Self { 78 Self { 79 available_passes: (0..max_subs) 80 .map(|_| SubPass { _private: () }) 81 .collect::<Vec<_>>(), 82 total_passes: max_subs, 83 } 84 } 85 86 pub fn take_pass(&mut self) -> Option<SubPass> { 87 self.available_passes.pop() 88 } 89 90 pub fn available_passes(&self) -> usize { 91 self.available_passes.len() 92 } 93 94 pub fn total_passes(&self) -> usize { 95 self.total_passes 96 } 97 98 pub fn return_pass(&mut self, pass: SubPass) { 99 self.available_passes.push(pass); 100 tracing::debug!( 101 "Returned pass. Using {} of {} passes", 102 self.total_passes - self.available_passes(), 103 self.total_passes 104 ); 105 } 106 107 pub(crate) fn spawn_passes(&mut self, new_passes: usize) { 108 for _ in 0..new_passes { 109 self.available_passes.push(SubPass { _private: () }); 110 } 111 } 112 } 113 114 /// Annihilates an existing `SubPass`. These should only be generated from the `RelayCoordinatorLimits` 115 /// when there is a new total subs which is less than the existing amount 116 pub struct SubPassRevocation { 117 revoked: bool, 118 } 119 120 impl SubPassRevocation { 121 pub fn revocate(&mut self, _: SubPass) { 122 self.revoked = true; 123 } 124 125 pub(crate) fn new() -> Self { 126 Self { revoked: false } 127 } 128 } 129 130 /// It completely breaks subscription management if we don't have strict accounting, so we crash if we fail to revocate 131 impl Drop for SubPassRevocation { 132 fn drop(&mut self) { 133 if !self.revoked { 134 panic!("The subscription pass revocator did not revoke the SubPass"); 135 } 136 } 137 } 138 139 pub struct SubPass { 140 _private: (), 141 } 142 143 #[cfg(test)] 144 mod tests { 145 use super::*; 146 147 // ==================== SubPassGuardian tests ==================== 148 149 #[test] 150 fn guardian_starts_with_correct_passes() { 151 let guardian = SubPassGuardian::new(10); 152 assert_eq!(guardian.available_passes(), 10); 153 } 154 155 #[test] 156 fn guardian_take_pass_decrements() { 157 let mut guardian = SubPassGuardian::new(5); 158 let pass = guardian.take_pass(); 159 assert!(pass.is_some()); 160 assert_eq!(guardian.available_passes(), 4); 161 } 162 163 #[test] 164 fn guardian_take_pass_returns_none_when_empty() { 165 let mut guardian = SubPassGuardian::new(1); 166 let _pass = guardian.take_pass(); 167 assert!(guardian.take_pass().is_none()); 168 assert_eq!(guardian.available_passes(), 0); 169 } 170 171 #[test] 172 fn guardian_return_pass_increments() { 173 let mut guardian = SubPassGuardian::new(1); 174 let pass = guardian.take_pass().unwrap(); 175 assert_eq!(guardian.available_passes(), 0); 176 guardian.return_pass(pass); 177 assert_eq!(guardian.available_passes(), 1); 178 } 179 180 #[test] 181 fn guardian_spawn_passes_adds_new_passes() { 182 let mut guardian = SubPassGuardian::new(2); 183 assert_eq!(guardian.available_passes(), 2); 184 guardian.spawn_passes(3); 185 assert_eq!(guardian.available_passes(), 5); 186 } 187 188 #[test] 189 fn guardian_multiple_take_and_return() { 190 let mut guardian = SubPassGuardian::new(3); 191 192 let pass1 = guardian.take_pass().unwrap(); 193 let pass2 = guardian.take_pass().unwrap(); 194 assert_eq!(guardian.available_passes(), 1); 195 196 guardian.return_pass(pass1); 197 assert_eq!(guardian.available_passes(), 2); 198 199 let _pass3 = guardian.take_pass().unwrap(); 200 assert_eq!(guardian.available_passes(), 1); 201 202 guardian.return_pass(pass2); 203 assert_eq!(guardian.available_passes(), 2); 204 } 205 206 // ==================== SubPassRevocation tests ==================== 207 208 #[test] 209 #[should_panic(expected = "did not revoke")] 210 fn revocation_panics_if_not_revoked() { 211 let _revocation = SubPassRevocation::new(); 212 // drop triggers panic 213 } 214 215 #[test] 216 fn revocation_does_not_panic_when_revoked() { 217 let mut guardian = SubPassGuardian::new(1); 218 let pass = guardian.take_pass().unwrap(); 219 let mut revocation = SubPassRevocation::new(); 220 revocation.revocate(pass); 221 // drop should not panic since revoked is true 222 } 223 224 #[test] 225 fn revocation_marks_as_revoked_after_revocate() { 226 let mut guardian = SubPassGuardian::new(1); 227 let pass = guardian.take_pass().unwrap(); 228 let mut revocation = SubPassRevocation::new(); 229 230 assert!(!revocation.revoked); 231 revocation.revocate(pass); 232 assert!(revocation.revoked); 233 } 234 235 // ==================== RelayCoordinatorLimits tests ==================== 236 237 #[test] 238 fn new_total_returns_none_when_same() { 239 let mut limits = RelayCoordinatorLimits::new(RelayLimitations { 240 maximum_subs: 5, 241 max_json_bytes: 400_000, 242 }); 243 244 let revocations = limits.new_total(5); 245 assert!(revocations.is_none()); 246 assert_eq!(limits.sub_guardian.available_passes(), 5); 247 } 248 249 #[test] 250 fn new_total_spawns_passes_when_increasing() { 251 let mut limits = RelayCoordinatorLimits::new(RelayLimitations { 252 maximum_subs: 5, 253 max_json_bytes: 400_000, 254 }); 255 256 let revocations = limits.new_total(10); 257 assert!(revocations.is_none()); 258 assert_eq!(limits.sub_guardian.available_passes(), 10); 259 } 260 261 #[test] 262 fn new_total_returns_revocations_when_decreasing() { 263 let mut limits = RelayCoordinatorLimits::new(RelayLimitations { 264 maximum_subs: 10, 265 max_json_bytes: 400_000, 266 }); 267 268 let revocations = limits.new_total(5); 269 assert!(revocations.is_none()); 270 } 271 272 #[test] 273 fn new_total_partial_revocations_when_passes_in_use() { 274 let mut limits = RelayCoordinatorLimits::new(RelayLimitations { 275 maximum_subs: 5, 276 max_json_bytes: 400_000, 277 }); 278 279 // Take 3 passes (simulate them being in use) 280 let pass = limits.sub_guardian.take_pass().unwrap(); 281 limits.sub_guardian.take_pass(); 282 limits.sub_guardian.take_pass(); 283 assert_eq!(limits.sub_guardian.available_passes(), 2); 284 285 // Now reduce to 2 total (need to remove 3) 286 let revocations = limits.new_total(2); 287 288 assert!(revocations.is_some()); 289 290 let mut revs = revocations.unwrap(); 291 // since there were two available passes, the guardian used those, but there is still one pass unaccounted for 292 assert_eq!(revs.len(), 1); 293 294 revs.pop().unwrap().revocate(pass); 295 } 296 }