@@ -27,6 +27,7 @@ use std::{
2727
2828use bimap:: BiMap ;
2929use futures:: { FutureExt , StreamExt , future:: BoxFuture , stream:: FuturesUnordered } ;
30+ use hashlink:: LruCache ;
3031use libp2p_core:: { Endpoint , Multiaddr , transport:: PortUse } ;
3132use libp2p_identity:: PeerId ;
3233use libp2p_request_response:: ProtocolSupport ;
@@ -40,6 +41,13 @@ use crate::{
4041 codec:: { Cookie , ErrorCode , Message , Namespace , NewRegistration , Registration , Ttl } ,
4142} ;
4243
44+ /// Default maximum active registrations per peer.
45+ pub const MAX_REGISTRATION_PEER : usize = 32 ;
46+ /// Default maximum active registrations in total.
47+ pub const MAX_REGISTRATIONS_TOTAL : usize = 10_000 ;
48+ /// Default size of the cache that stores client cookies.
49+ pub const COOKIES_CACHE_SIZE : usize = 10_000 ;
50+
4351pub struct Behaviour {
4452 inner : libp2p_request_response:: Behaviour < crate :: codec:: Codec > ,
4553
@@ -49,6 +57,9 @@ pub struct Behaviour {
4957pub struct Config {
5058 min_ttl : Ttl ,
5159 max_ttl : Ttl ,
60+ max_registrations_per_peer : usize ,
61+ max_registrations_total : usize ,
62+ max_cookies : usize ,
5263}
5364
5465impl Config {
@@ -61,13 +72,31 @@ impl Config {
6172 self . max_ttl = max_ttl;
6273 self
6374 }
75+
76+ pub fn with_max_registration_per_peer ( mut self , max : usize ) -> Self {
77+ self . max_registrations_per_peer = max;
78+ self
79+ }
80+
81+ pub fn with_max_registration_total ( mut self , max : usize ) -> Self {
82+ self . max_registrations_total = max;
83+ self
84+ }
85+
86+ pub fn with_max_stored_cookies ( mut self , max : usize ) -> Self {
87+ self . max_cookies = max;
88+ self
89+ }
6490}
6591
6692impl Default for Config {
6793 fn default ( ) -> Self {
6894 Self {
6995 min_ttl : MIN_TTL ,
7096 max_ttl : MAX_TTL ,
97+ max_registrations_per_peer : MAX_REGISTRATION_PEER ,
98+ max_registrations_total : MAX_REGISTRATIONS_TOTAL ,
99+ max_cookies : COOKIES_CACHE_SIZE ,
71100 }
72101 }
73102}
@@ -279,9 +308,7 @@ fn handle_request(
279308
280309 Some ( ( event, Some ( response) ) )
281310 }
282- Err ( TtlOutOfRange :: TooLong { .. } ) | Err ( TtlOutOfRange :: TooShort { .. } ) => {
283- let error = ErrorCode :: InvalidTtl ;
284-
311+ Err ( error) => {
285312 let response = Message :: RegisterResponse ( Err ( error) ) ;
286313
287314 let event = Event :: PeerNotRegistered {
@@ -351,69 +378,53 @@ impl RegistrationId {
351378#[ derive( Debug , PartialEq ) ]
352379struct ExpiredRegistration ( Registration ) ;
353380
354- pub struct Registrations {
381+ struct Registrations {
382+ config : Config ,
355383 registrations_for_peer : BiMap < ( PeerId , Namespace ) , RegistrationId > ,
356384 registrations : HashMap < RegistrationId , Registration > ,
357- cookies : HashMap < Cookie , HashSet < RegistrationId > > ,
358- min_ttl : Ttl ,
359- max_ttl : Ttl ,
385+ cookies : LruCache < Cookie , HashSet < RegistrationId > > ,
360386 next_expiry : FuturesUnordered < BoxFuture < ' static , RegistrationId > > ,
361387}
362388
363- #[ derive( Debug , thiserror:: Error ) ]
364- pub enum TtlOutOfRange {
365- #[ error( "Requested TTL ({requested}s) is too long; max {bound}s" ) ]
366- TooLong { bound : Ttl , requested : Ttl } ,
367- #[ error( "Requested TTL ({requested}s) is too short; min {bound}s" ) ]
368- TooShort { bound : Ttl , requested : Ttl } ,
369- }
370-
371389impl Default for Registrations {
372390 fn default ( ) -> Self {
373391 Registrations :: with_config ( Config :: default ( ) )
374392 }
375393}
376394
377395impl Registrations {
378- pub fn with_config ( config : Config ) -> Self {
396+ fn with_config ( config : Config ) -> Self {
379397 Self {
380398 registrations_for_peer : Default :: default ( ) ,
381399 registrations : Default :: default ( ) ,
382- min_ttl : config. min_ttl ,
383- max_ttl : config. max_ttl ,
384- cookies : Default :: default ( ) ,
400+ cookies : LruCache :: new ( config. max_cookies ) ,
401+ config,
385402 next_expiry : FuturesUnordered :: from_iter ( vec ! [ futures:: future:: pending( ) . boxed( ) ] ) ,
386403 }
387404 }
388405
389- pub fn add (
390- & mut self ,
391- new_registration : NewRegistration ,
392- ) -> Result < Registration , TtlOutOfRange > {
406+ fn add ( & mut self , new_registration : NewRegistration ) -> Result < Registration , ErrorCode > {
393407 let ttl = new_registration. effective_ttl ( ) ;
394- if ttl > self . max_ttl {
395- return Err ( TtlOutOfRange :: TooLong {
396- bound : self . max_ttl ,
397- requested : ttl,
398- } ) ;
399- }
400- if ttl < self . min_ttl {
401- return Err ( TtlOutOfRange :: TooShort {
402- bound : self . min_ttl ,
403- requested : ttl,
404- } ) ;
408+ if ttl > self . config . max_ttl || ttl < self . config . min_ttl {
409+ return Err ( ErrorCode :: InvalidTtl ) ;
405410 }
406411
407- let namespace = new_registration. namespace ;
408- let registration_id = RegistrationId :: new ( ) ;
412+ let peer = new_registration. record . peer_id ( ) ;
409413
410- if let Some ( old_registration ) = self
414+ if self
411415 . registrations_for_peer
412- . get_by_left ( & ( new_registration. record . peer_id ( ) , namespace. clone ( ) ) )
416+ . left_values ( )
417+ . filter ( |( p, _) | p == & peer)
418+ . count ( )
419+ >= self . config . max_registrations_per_peer
420+ || self . registrations_for_peer . len ( ) > self . config . max_registrations_total
413421 {
414- self . registrations . remove ( old_registration ) ;
422+ return Err ( ErrorCode :: Unavailable ) ;
415423 }
416424
425+ let namespace = new_registration. namespace ;
426+ let registration_id = RegistrationId :: new ( ) ;
427+
417428 self . registrations_for_peer . insert (
418429 ( new_registration. record . peer_id ( ) , namespace. clone ( ) ) ,
419430 registration_id,
@@ -436,7 +447,7 @@ impl Registrations {
436447 Ok ( registration)
437448 }
438449
439- pub fn remove ( & mut self , namespace : Namespace , peer_id : PeerId ) {
450+ fn remove ( & mut self , namespace : Namespace , peer_id : PeerId ) {
440451 let reggo_to_remove = self
441452 . registrations_for_peer
442453 . remove_by_left ( & ( peer_id, namespace) ) ;
@@ -446,7 +457,7 @@ impl Registrations {
446457 }
447458 }
448459
449- pub fn get (
460+ fn get (
450461 & mut self ,
451462 discover_namespace : Option < Namespace > ,
452463 cookie : Option < Cookie > ,
@@ -536,9 +547,8 @@ impl Registrations {
536547 }
537548}
538549
539- #[ derive( Debug , thiserror:: Error , Eq , PartialEq ) ]
540- #[ error( "The provided cookie is not valid for a DISCOVER request for the given namespace" ) ]
541- pub struct CookieNamespaceMismatch ;
550+ #[ derive( Debug , Eq , PartialEq ) ]
551+ struct CookieNamespaceMismatch ;
542552
543553#[ cfg( test) ]
544554mod tests {
@@ -646,10 +656,8 @@ mod tests {
646656
647657 #[ tokio:: test]
648658 async fn given_two_registration_ttls_one_expires_one_lives ( ) {
649- let mut registrations = Registrations :: with_config ( Config {
650- min_ttl : 0 ,
651- max_ttl : 4 ,
652- } ) ;
659+ let mut registrations =
660+ Registrations :: with_config ( Config :: default ( ) . with_min_ttl ( 0 ) . with_max_ttl ( 4 ) ) ;
653661
654662 let start_time = SystemTime :: now ( ) ;
655663
@@ -682,10 +690,8 @@ mod tests {
682690
683691 #[ tokio:: test]
684692 async fn given_peer_unregisters_before_expiry_do_not_emit_registration_expired ( ) {
685- let mut registrations = Registrations :: with_config ( Config {
686- min_ttl : 1 ,
687- max_ttl : 10 ,
688- } ) ;
693+ let mut registrations =
694+ Registrations :: with_config ( Config :: default ( ) . with_min_ttl ( 1 ) . with_max_ttl ( 10 ) ) ;
689695 let dummy_registration = new_dummy_registration_with_ttl ( "foo" , 2 ) ;
690696 let namespace = dummy_registration. namespace . clone ( ) ;
691697 let peer_id = dummy_registration. record . peer_id ( ) ;
@@ -704,10 +710,8 @@ mod tests {
704710 #[ tokio:: test]
705711 async fn given_all_registrations_expired_then_successfully_handle_new_registration_and_expiry ( )
706712 {
707- let mut registrations = Registrations :: with_config ( Config {
708- min_ttl : 0 ,
709- max_ttl : 10 ,
710- } ) ;
713+ let mut registrations =
714+ Registrations :: with_config ( Config :: default ( ) . with_min_ttl ( 0 ) . with_max_ttl ( 10 ) ) ;
711715 let dummy_registration = new_dummy_registration_with_ttl ( "foo" , 1 ) ;
712716
713717 registrations. add ( dummy_registration. clone ( ) ) . unwrap ( ) ;
@@ -721,10 +725,8 @@ mod tests {
721725
722726 #[ tokio:: test]
723727 async fn cookies_are_cleaned_up_if_registrations_expire ( ) {
724- let mut registrations = Registrations :: with_config ( Config {
725- min_ttl : 1 ,
726- max_ttl : 10 ,
727- } ) ;
728+ let mut registrations =
729+ Registrations :: with_config ( Config :: default ( ) . with_min_ttl ( 1 ) . with_max_ttl ( 10 ) ) ;
728730
729731 registrations
730732 . add ( new_dummy_registration_with_ttl ( "foo" , 2 ) )
0 commit comments