@@ -206,13 +206,12 @@ impl<T: Callback> ServiceProtocol for IdentifyProtocol<T> {
206206 async fn connected ( & mut self , context : ProtocolContextMutRef < ' _ > , version : & str ) {
207207 let session = context. session ;
208208 debug ! ( "IdentifyProtocol connected, session: {:?}" , session) ;
209- if self . callback . register ( & context, version) {
210- let remote_info =
211- RemoteInfo :: new ( session. clone ( ) , Duration :: from_secs ( DEFAULT_TIMEOUT ) ) ;
212- self . remote_infos . insert ( session. id , remote_info) ;
213-
214- let listen_addrs: Vec < Multiaddr > = self
215- . callback
209+ let remote_info = RemoteInfo :: new ( session. clone ( ) , Duration :: from_secs ( DEFAULT_TIMEOUT ) ) ;
210+ self . remote_infos . insert ( session. id , remote_info) ;
211+ let listen_addrs = if self . callback . register ( & context, version) {
212+ Vec :: new ( )
213+ } else {
214+ self . callback
216215 . local_listen_addrs ( )
217216 . iter ( )
218217 . filter ( |addr| {
@@ -222,24 +221,15 @@ impl<T: Callback> ServiceProtocol for IdentifyProtocol<T> {
222221 } )
223222 . take ( MAX_ADDRS )
224223 . cloned ( )
225- . collect ( ) ;
226-
227- let identify = self . callback . identify ( ) ;
228- let data =
229- IdentifyMessage :: new ( listen_addrs, session. address . clone ( ) , identify) . encode ( ) ;
230- let _ = context
231- . quick_send_message ( data)
232- . await
233- . map_err ( |err| error ! ( "IdentifyProtocol quick_send_message, error: {:?}" , err) ) ;
234- } else if let Err ( e) = context
235- . close_protocol ( context. session . id , SupportProtocols :: Identify . protocol_id ( ) )
224+ . collect ( )
225+ } ;
226+
227+ let identify = self . callback . identify ( ) ;
228+ let data = IdentifyMessage :: new ( listen_addrs, session. address . clone ( ) , identify) . encode ( ) ;
229+ let _ = context
230+ . quick_send_message ( data)
236231 . await
237- {
238- error ! (
239- "Block-Relay-Only Session close IdentifyProtocol failed, session: {:?} {}" ,
240- context. session, e
241- ) ;
242- }
232+ . map_err ( |err| error ! ( "IdentifyProtocol quick_send_message, error: {:?}" , err) ) ;
243233 }
244234
245235 async fn disconnected ( & mut self , context : ProtocolContextMutRef < ' _ > ) {
@@ -363,17 +353,12 @@ impl IdentifyCallback {
363353impl Callback for IdentifyCallback {
364354 fn register ( & self , context : & ProtocolContextMutRef , version : & str ) -> bool {
365355 let session_id = context. session . id ;
366- let is_anchor = self
367- . network_state
368- . with_peer_registry ( |reg| reg. is_anchor ( session_id) ) ;
369- if !is_anchor {
370- self . network_state . with_peer_registry_mut ( |reg| {
371- reg. get_peer_mut ( session_id) . map ( |peer| {
372- peer. protocols . insert ( context. proto_id , version. to_owned ( ) ) ;
373- } )
374- } ) ;
375- }
376- !is_anchor
356+ self . network_state . with_peer_registry_mut ( |reg| {
357+ if let Some ( peer) = reg. get_peer_mut ( session_id) {
358+ peer. protocols . insert ( context. proto_id , version. to_owned ( ) ) ;
359+ }
360+ reg. is_anchor ( session_id)
361+ } )
377362 }
378363
379364 fn unregister ( & self , context : & ProtocolContextMutRef ) {
0 commit comments