18
18
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19
19
// DEALINGS IN THE SOFTWARE.
20
20
21
- use libp2p_core:: { Endpoint , Multiaddr } ;
21
+ use libp2p_core:: { ConnectedPoint , Endpoint , Multiaddr } ;
22
22
use libp2p_identity:: PeerId ;
23
23
use libp2p_swarm:: {
24
+ behaviour:: { ConnectionEstablished , DialFailure , ListenFailure } ,
24
25
dummy, ConnectionClosed , ConnectionDenied , ConnectionId , FromSwarm , NetworkBehaviour ,
25
26
PollParameters , THandler , THandlerInEvent , THandlerOutEvent , ToSwarm ,
26
27
} ;
@@ -249,12 +250,6 @@ impl NetworkBehaviour for Behaviour {
249
250
Kind :: EstablishedTotal ,
250
251
) ?;
251
252
252
- self . established_inbound_connections . insert ( connection_id) ;
253
- self . established_per_peer
254
- . entry ( peer)
255
- . or_default ( )
256
- . insert ( connection_id) ;
257
-
258
253
Ok ( dummy:: ConnectionHandler )
259
254
}
260
255
@@ -305,12 +300,6 @@ impl NetworkBehaviour for Behaviour {
305
300
Kind :: EstablishedTotal ,
306
301
) ?;
307
302
308
- self . established_outbound_connections . insert ( connection_id) ;
309
- self . established_per_peer
310
- . entry ( peer)
311
- . or_default ( )
312
- . insert ( connection_id) ;
313
-
314
303
Ok ( dummy:: ConnectionHandler )
315
304
}
316
305
@@ -328,10 +317,33 @@ impl NetworkBehaviour for Behaviour {
328
317
. or_default ( )
329
318
. remove ( & connection_id) ;
330
319
}
331
- FromSwarm :: ConnectionEstablished ( _) => { }
320
+ FromSwarm :: ConnectionEstablished ( ConnectionEstablished {
321
+ peer_id,
322
+ endpoint,
323
+ connection_id,
324
+ ..
325
+ } ) => {
326
+ match endpoint {
327
+ ConnectedPoint :: Listener { .. } => {
328
+ self . established_inbound_connections . insert ( connection_id) ;
329
+ }
330
+ ConnectedPoint :: Dialer { .. } => {
331
+ self . established_outbound_connections . insert ( connection_id) ;
332
+ }
333
+ }
334
+
335
+ self . established_per_peer
336
+ . entry ( peer_id)
337
+ . or_default ( )
338
+ . insert ( connection_id) ;
339
+ }
340
+ FromSwarm :: DialFailure ( DialFailure { connection_id, .. } ) => {
341
+ self . pending_outbound_connections . remove ( & connection_id) ;
342
+ }
332
343
FromSwarm :: AddressChange ( _) => { }
333
- FromSwarm :: DialFailure ( _) => { }
334
- FromSwarm :: ListenFailure ( _) => { }
344
+ FromSwarm :: ListenFailure ( ListenFailure { connection_id, .. } ) => {
345
+ self . pending_inbound_connections . remove ( & connection_id) ;
346
+ }
335
347
FromSwarm :: NewListener ( _) => { }
336
348
FromSwarm :: NewListenAddr ( _) => { }
337
349
FromSwarm :: ExpiredListenAddr ( _) => { }
@@ -364,7 +376,9 @@ impl NetworkBehaviour for Behaviour {
364
376
#[ cfg( test) ]
365
377
mod tests {
366
378
use super :: * ;
367
- use libp2p_swarm:: { dial_opts:: DialOpts , DialError , ListenError , Swarm , SwarmEvent } ;
379
+ use libp2p_swarm:: {
380
+ behaviour:: toggle:: Toggle , dial_opts:: DialOpts , DialError , ListenError , Swarm , SwarmEvent ,
381
+ } ;
368
382
use libp2p_swarm_test:: SwarmExt ;
369
383
use quickcheck:: * ;
370
384
@@ -466,19 +480,125 @@ mod tests {
466
480
quickcheck ( prop as fn ( _) ) ;
467
481
}
468
482
483
+ /// Another sibling [`NetworkBehaviour`] implementation might deny established connections in
484
+ /// [`handle_established_outbound_connection`] or [`handle_established_inbound_connection`].
485
+ /// [`Behaviour`] must not increase the established counters in
486
+ /// [`handle_established_outbound_connection`] or [`handle_established_inbound_connection`], but
487
+ /// in [`SwarmEvent::ConnectionEstablished`] as the connection might still be denied by a
488
+ /// sibling [`NetworkBehaviour`] in the former case. Only in the latter case
489
+ /// ([`SwarmEvent::ConnectionEstablished`]) can the connection be seen as established.
490
+ #[ test]
491
+ fn support_other_behaviour_denying_connection ( ) {
492
+ let mut swarm1 = Swarm :: new_ephemeral ( |_| {
493
+ Behaviour :: new_with_connection_denier ( ConnectionLimits :: default ( ) )
494
+ } ) ;
495
+ let mut swarm2 = Swarm :: new_ephemeral ( |_| Behaviour :: new ( ConnectionLimits :: default ( ) ) ) ;
496
+
497
+ async_std:: task:: block_on ( async {
498
+ // Have swarm2 dial swarm1.
499
+ let ( listen_addr, _) = swarm1. listen ( ) . await ;
500
+ swarm2. dial ( listen_addr) . unwrap ( ) ;
501
+ async_std:: task:: spawn ( swarm2. loop_on_next ( ) ) ;
502
+
503
+ // Wait for the ConnectionDenier of swarm1 to deny the established connection.
504
+ let cause = swarm1
505
+ . wait ( |event| match event {
506
+ SwarmEvent :: IncomingConnectionError {
507
+ error : ListenError :: Denied { cause } ,
508
+ ..
509
+ } => Some ( cause) ,
510
+ _ => None ,
511
+ } )
512
+ . await ;
513
+
514
+ cause. downcast :: < std:: io:: Error > ( ) . unwrap ( ) ;
515
+
516
+ assert_eq ! (
517
+ 0 ,
518
+ swarm1
519
+ . behaviour_mut( )
520
+ . limits
521
+ . established_inbound_connections
522
+ . len( ) ,
523
+ "swarm1 connection limit behaviour to not count denied established connection as established connection"
524
+ )
525
+ } ) ;
526
+ }
527
+
469
528
#[ derive( libp2p_swarm_derive:: NetworkBehaviour ) ]
470
529
#[ behaviour( prelude = "libp2p_swarm::derive_prelude" ) ]
471
530
struct Behaviour {
472
531
limits : super :: Behaviour ,
473
532
keep_alive : libp2p_swarm:: keep_alive:: Behaviour ,
533
+ connection_denier : Toggle < ConnectionDenier > ,
474
534
}
475
535
476
536
impl Behaviour {
477
537
fn new ( limits : ConnectionLimits ) -> Self {
478
538
Self {
479
539
limits : super :: Behaviour :: new ( limits) ,
480
540
keep_alive : libp2p_swarm:: keep_alive:: Behaviour ,
541
+ connection_denier : None . into ( ) ,
481
542
}
482
543
}
544
+ fn new_with_connection_denier ( limits : ConnectionLimits ) -> Self {
545
+ Self {
546
+ limits : super :: Behaviour :: new ( limits) ,
547
+ keep_alive : libp2p_swarm:: keep_alive:: Behaviour ,
548
+ connection_denier : Some ( ConnectionDenier { } ) . into ( ) ,
549
+ }
550
+ }
551
+ }
552
+
553
+ struct ConnectionDenier { }
554
+
555
+ impl NetworkBehaviour for ConnectionDenier {
556
+ type ConnectionHandler = dummy:: ConnectionHandler ;
557
+ type ToSwarm = Void ;
558
+
559
+ fn handle_established_inbound_connection (
560
+ & mut self ,
561
+ _connection_id : ConnectionId ,
562
+ _peer : PeerId ,
563
+ _local_addr : & Multiaddr ,
564
+ _remote_addr : & Multiaddr ,
565
+ ) -> Result < THandler < Self > , ConnectionDenied > {
566
+ Err ( ConnectionDenied :: new ( std:: io:: Error :: new (
567
+ std:: io:: ErrorKind :: Other ,
568
+ "ConnectionDenier" ,
569
+ ) ) )
570
+ }
571
+
572
+ fn handle_established_outbound_connection (
573
+ & mut self ,
574
+ _connection_id : ConnectionId ,
575
+ _peer : PeerId ,
576
+ _addr : & Multiaddr ,
577
+ _role_override : Endpoint ,
578
+ ) -> Result < THandler < Self > , ConnectionDenied > {
579
+ Err ( ConnectionDenied :: new ( std:: io:: Error :: new (
580
+ std:: io:: ErrorKind :: Other ,
581
+ "ConnectionDenier" ,
582
+ ) ) )
583
+ }
584
+
585
+ fn on_swarm_event ( & mut self , _event : FromSwarm < Self :: ConnectionHandler > ) { }
586
+
587
+ fn on_connection_handler_event (
588
+ & mut self ,
589
+ _peer_id : PeerId ,
590
+ _connection_id : ConnectionId ,
591
+ event : THandlerOutEvent < Self > ,
592
+ ) {
593
+ void:: unreachable ( event)
594
+ }
595
+
596
+ fn poll (
597
+ & mut self ,
598
+ _cx : & mut Context < ' _ > ,
599
+ _params : & mut impl PollParameters ,
600
+ ) -> Poll < ToSwarm < Self :: ToSwarm , THandlerInEvent < Self > > > {
601
+ Poll :: Pending
602
+ }
483
603
}
484
604
}
0 commit comments