@@ -24,6 +24,7 @@ use crate::{
2424
2525use bytes:: Bytes ;
2626use codec:: { DecodeAll , Encode } ;
27+ use futures:: { channel:: oneshot, stream:: FuturesUnordered , StreamExt } ;
2728use libp2p:: {
2829 core:: connection:: ConnectionId ,
2930 swarm:: {
@@ -35,11 +36,14 @@ use libp2p::{
3536use log:: { debug, error, warn} ;
3637
3738use sc_network_common:: { role:: Roles , sync:: message:: BlockAnnouncesHandshake } ;
39+ use sc_utils:: mpsc:: TracingUnboundedSender ;
3840use sp_runtime:: traits:: Block as BlockT ;
3941
4042use std:: {
4143 collections:: { HashMap , HashSet , VecDeque } ,
44+ future:: Future ,
4245 iter,
46+ pin:: Pin ,
4347 task:: Poll ,
4448} ;
4549
@@ -68,6 +72,9 @@ mod rep {
6872 pub const BAD_MESSAGE : Rep = Rep :: new ( -( 1 << 12 ) , "Bad message" ) ;
6973}
7074
75+ type PendingSyncSubstreamValidation =
76+ Pin < Box < dyn Future < Output = Result < ( PeerId , Roles ) , PeerId > > + Send > > ;
77+
7178// Lock must always be taken in order declared here.
7279pub struct Protocol < B : BlockT > {
7380 /// Pending list of messages to return from `poll` as a priority.
@@ -87,6 +94,8 @@ pub struct Protocol<B: BlockT> {
8794 bad_handshake_substreams : HashSet < ( PeerId , sc_peerset:: SetId ) > ,
8895 /// Connected peers.
8996 peers : HashMap < PeerId , Roles > ,
97+ sync_substream_validations : FuturesUnordered < PendingSyncSubstreamValidation > ,
98+ tx : TracingUnboundedSender < crate :: event:: SyncEvent < B > > ,
9099 _marker : std:: marker:: PhantomData < B > ,
91100}
92101
@@ -96,6 +105,7 @@ impl<B: BlockT> Protocol<B> {
96105 roles : Roles ,
97106 network_config : & config:: NetworkConfiguration ,
98107 block_announces_protocol : config:: NonDefaultSetConfig ,
108+ tx : TracingUnboundedSender < crate :: event:: SyncEvent < B > > ,
99109 ) -> error:: Result < ( Self , sc_peerset:: PeersetHandle , Vec < ( PeerId , Multiaddr ) > ) > {
100110 let mut known_addresses = Vec :: new ( ) ;
101111
@@ -179,6 +189,8 @@ impl<B: BlockT> Protocol<B> {
179189 . collect ( ) ,
180190 bad_handshake_substreams : Default :: default ( ) ,
181191 peers : HashMap :: new ( ) ,
192+ sync_substream_validations : FuturesUnordered :: new ( ) ,
193+ tx,
182194 // TODO: remove when `BlockAnnouncesHandshake` is moved away from `Protocol`
183195 _marker : Default :: default ( ) ,
184196 } ;
@@ -418,6 +430,23 @@ impl<B: BlockT> NetworkBehaviour for Protocol<B> {
418430 return Poll :: Ready ( NetworkBehaviourAction :: CloseConnection { peer_id, connection } ) ,
419431 } ;
420432
433+ while let Poll :: Ready ( Some ( validation_result) ) =
434+ self . sync_substream_validations . poll_next_unpin ( cx)
435+ {
436+ match validation_result {
437+ Ok ( ( peer, roles) ) => {
438+ self . peers . insert ( peer, roles) ;
439+ } ,
440+ Err ( peer) => {
441+ log:: debug!(
442+ target: "sub-libp2p" ,
443+ "`SyncingEngine` rejected stream"
444+ ) ;
445+ self . behaviour . disconnect_peer ( & peer, HARDCODED_PEERSETS_SYNC ) ;
446+ } ,
447+ }
448+ }
449+
421450 let outcome = match event {
422451 NotificationsOut :: CustomProtocolOpen {
423452 peer_id,
@@ -440,16 +469,29 @@ impl<B: BlockT> NetworkBehaviour for Protocol<B> {
440469 best_hash : handshake. best_hash ,
441470 genesis_hash : handshake. genesis_hash ,
442471 } ;
443- self . peers . insert ( peer_id, roles) ;
444472
445- CustomMessageOutcome :: NotificationStreamOpened {
446- remote : peer_id,
447- protocol : self . notification_protocols [ usize:: from ( set_id) ] . clone ( ) ,
448- negotiated_fallback,
449- received_handshake : handshake. encode ( ) ,
450- roles,
451- notifications_sink,
452- }
473+ let ( tx, rx) = oneshot:: channel ( ) ;
474+ let _ = self . tx . unbounded_send (
475+ crate :: SyncEvent :: NotificationStreamOpened {
476+ remote : peer_id,
477+ received_handshake : handshake,
478+ sink : notifications_sink,
479+ tx,
480+ } ,
481+ ) ;
482+ self . sync_substream_validations . push ( Box :: pin ( async move {
483+ match rx. await {
484+ Ok ( accepted) =>
485+ if accepted {
486+ Ok ( ( peer_id, roles) )
487+ } else {
488+ Err ( peer_id)
489+ } ,
490+ Err ( _) => Err ( peer_id) ,
491+ }
492+ } ) ) ;
493+
494+ CustomMessageOutcome :: None
453495 } ,
454496 Ok ( msg) => {
455497 debug ! (
@@ -469,15 +511,27 @@ impl<B: BlockT> NetworkBehaviour for Protocol<B> {
469511 let roles = handshake. roles ;
470512 self . peers . insert ( peer_id, roles) ;
471513
472- CustomMessageOutcome :: NotificationStreamOpened {
473- remote : peer_id,
474- protocol : self . notification_protocols [ usize:: from ( set_id) ]
475- . clone ( ) ,
476- negotiated_fallback,
477- received_handshake,
478- roles,
479- notifications_sink,
480- }
514+ let ( tx, rx) = oneshot:: channel ( ) ;
515+ let _ = self . tx . unbounded_send (
516+ crate :: SyncEvent :: NotificationStreamOpened {
517+ remote : peer_id,
518+ received_handshake : handshake,
519+ sink : notifications_sink,
520+ tx,
521+ } ,
522+ ) ;
523+ self . sync_substream_validations . push ( Box :: pin ( async move {
524+ match rx. await {
525+ Ok ( accepted) =>
526+ if accepted {
527+ Ok ( ( peer_id, roles) )
528+ } else {
529+ Err ( peer_id)
530+ } ,
531+ Err ( _) => Err ( peer_id) ,
532+ }
533+ } ) ) ;
534+ CustomMessageOutcome :: None
481535 } ,
482536 Err ( err2) => {
483537 log:: debug!(
@@ -535,6 +589,12 @@ impl<B: BlockT> NetworkBehaviour for Protocol<B> {
535589 NotificationsOut :: CustomProtocolReplaced { peer_id, notifications_sink, set_id } =>
536590 if self . bad_handshake_substreams . contains ( & ( peer_id, set_id) ) {
537591 CustomMessageOutcome :: None
592+ } else if set_id == HARDCODED_PEERSETS_SYNC {
593+ let _ = self . tx . unbounded_send ( crate :: SyncEvent :: NotificationSinkReplaced {
594+ remote : peer_id,
595+ sink : notifications_sink,
596+ } ) ;
597+ CustomMessageOutcome :: None
538598 } else {
539599 CustomMessageOutcome :: NotificationStreamReplaced {
540600 remote : peer_id,
@@ -548,6 +608,12 @@ impl<B: BlockT> NetworkBehaviour for Protocol<B> {
548608 // handshake. The outer layers have never received an opening event about this
549609 // substream, and consequently shouldn't receive a closing event either.
550610 CustomMessageOutcome :: None
611+ } else if set_id == HARDCODED_PEERSETS_SYNC {
612+ let _ = self . tx . unbounded_send ( crate :: SyncEvent :: NotificationStreamClosed {
613+ remote : peer_id,
614+ } ) ;
615+ self . peers . remove ( & peer_id) ;
616+ CustomMessageOutcome :: None
551617 } else {
552618 CustomMessageOutcome :: NotificationStreamClosed {
553619 remote : peer_id,
@@ -558,6 +624,12 @@ impl<B: BlockT> NetworkBehaviour for Protocol<B> {
558624 NotificationsOut :: Notification { peer_id, set_id, message } => {
559625 if self . bad_handshake_substreams . contains ( & ( peer_id, set_id) ) {
560626 CustomMessageOutcome :: None
627+ } else if set_id == HARDCODED_PEERSETS_SYNC {
628+ let _ = self . tx . unbounded_send ( crate :: SyncEvent :: NotificationsReceived {
629+ remote : peer_id,
630+ messages : vec ! [ message. freeze( ) ] ,
631+ } ) ;
632+ CustomMessageOutcome :: None
561633 } else {
562634 let protocol_name = self . notification_protocols [ usize:: from ( set_id) ] . clone ( ) ;
563635 CustomMessageOutcome :: NotificationsReceived {
0 commit comments