@@ -33,8 +33,9 @@ use crate::{
3333    config:: { ConfigBuilder ,  TopicMeshConfig } , 
3434    protocol:: GossipsubCodec , 
3535    rpc:: Receiver , 
36+     rpc_proto:: proto, 
3637    subscription_filter:: WhitelistSubscriptionFilter , 
37-     types:: RpcIn , 
38+     types:: { ControlAction ,   Extensions ,   RpcIn ,   RpcOut } , 
3839    IdentTopic  as  Topic , 
3940} ; 
4041
@@ -248,6 +249,7 @@ where
248249            topics :  Default :: default ( ) , 
249250            sender, 
250251            dont_send :  LinkedHashMap :: new ( ) , 
252+             extensions :  None , 
251253        } , 
252254    ) ; 
253255
@@ -644,6 +646,7 @@ fn test_join() {
644646                topics :  Default :: default ( ) , 
645647                sender, 
646648                dont_send :  LinkedHashMap :: new ( ) , 
649+                 extensions :  None , 
647650            } , 
648651        ) ; 
649652        receivers. insert ( random_peer,  receiver) ; 
@@ -1041,6 +1044,7 @@ fn test_get_random_peers() {
10411044                topics :  topics. clone ( ) , 
10421045                sender :  Sender :: new ( gs. config . connection_handler_queue_len ( ) ) , 
10431046                dont_send :  LinkedHashMap :: new ( ) , 
1047+                 extensions :  None , 
10441048            } , 
10451049        ) ; 
10461050    } 
@@ -5595,6 +5599,7 @@ fn test_all_queues_full() {
55955599            topics :  topics. clone ( ) , 
55965600            sender :  Sender :: new ( 2 ) , 
55975601            dont_send :  LinkedHashMap :: new ( ) , 
5602+             extensions :  None , 
55985603        } , 
55995604    ) ; 
56005605
@@ -5631,6 +5636,7 @@ fn test_slow_peer_returns_failed_publish() {
56315636            topics :  topics. clone ( ) , 
56325637            sender :  Sender :: new ( 2 ) , 
56335638            dont_send :  LinkedHashMap :: new ( ) , 
5639+             extensions :  None , 
56345640        } , 
56355641    ) ; 
56365642    let  peer_id = PeerId :: random ( ) ; 
@@ -5644,6 +5650,7 @@ fn test_slow_peer_returns_failed_publish() {
56445650            topics :  topics. clone ( ) , 
56455651            sender :  Sender :: new ( gs. config . connection_handler_queue_len ( ) ) , 
56465652            dont_send :  LinkedHashMap :: new ( ) , 
5653+             extensions :  None , 
56475654        } , 
56485655    ) ; 
56495656
@@ -5705,6 +5712,7 @@ fn test_slow_peer_returns_failed_ihave_handling() {
57055712            topics :  topics. clone ( ) , 
57065713            sender :  Sender :: new ( 2 ) , 
57075714            dont_send :  LinkedHashMap :: new ( ) , 
5715+             extensions :  None , 
57085716        } , 
57095717    ) ; 
57105718    peers. push ( slow_peer_id) ; 
@@ -5722,6 +5730,7 @@ fn test_slow_peer_returns_failed_ihave_handling() {
57225730            topics :  topics. clone ( ) , 
57235731            sender :  Sender :: new ( gs. config . connection_handler_queue_len ( ) ) , 
57245732            dont_send :  LinkedHashMap :: new ( ) , 
5733+             extensions :  None , 
57255734        } , 
57265735    ) ; 
57275736
@@ -5819,6 +5828,7 @@ fn test_slow_peer_returns_failed_iwant_handling() {
58195828            topics :  topics. clone ( ) , 
58205829            sender :  Sender :: new ( 2 ) , 
58215830            dont_send :  LinkedHashMap :: new ( ) , 
5831+             extensions :  None , 
58225832        } , 
58235833    ) ; 
58245834    peers. push ( slow_peer_id) ; 
@@ -5836,6 +5846,7 @@ fn test_slow_peer_returns_failed_iwant_handling() {
58365846            topics :  topics. clone ( ) , 
58375847            sender :  Sender :: new ( gs. config . connection_handler_queue_len ( ) ) , 
58385848            dont_send :  LinkedHashMap :: new ( ) , 
5849+             extensions :  None , 
58395850        } , 
58405851    ) ; 
58415852
@@ -5913,6 +5924,7 @@ fn test_slow_peer_returns_failed_forward() {
59135924            topics :  topics. clone ( ) , 
59145925            sender :  Sender :: new ( 2 ) , 
59155926            dont_send :  LinkedHashMap :: new ( ) , 
5927+             extensions :  None , 
59165928        } , 
59175929    ) ; 
59185930    peers. push ( slow_peer_id) ; 
@@ -5930,6 +5942,7 @@ fn test_slow_peer_returns_failed_forward() {
59305942            topics :  topics. clone ( ) , 
59315943            sender :  Sender :: new ( gs. config . connection_handler_queue_len ( ) ) , 
59325944            dont_send :  LinkedHashMap :: new ( ) , 
5945+             extensions :  None , 
59335946        } , 
59345947    ) ; 
59355948
@@ -6012,6 +6025,7 @@ fn test_slow_peer_is_downscored_on_publish() {
60126025            topics :  topics. clone ( ) , 
60136026            sender :  Sender :: new ( 2 ) , 
60146027            dont_send :  LinkedHashMap :: new ( ) , 
6028+             extensions :  None , 
60156029        } , 
60166030    ) ; 
60176031    gs. as_peer_score_mut ( ) . add_peer ( slow_peer_id) ; 
@@ -6026,6 +6040,7 @@ fn test_slow_peer_is_downscored_on_publish() {
60266040            topics :  topics. clone ( ) , 
60276041            sender :  Sender :: new ( gs. config . connection_handler_queue_len ( ) ) , 
60286042            dont_send :  LinkedHashMap :: new ( ) , 
6043+             extensions :  None , 
60296044        } , 
60306045    ) ; 
60316046
@@ -6787,3 +6802,76 @@ fn test_validation_message_size_within_topic_specific() {
67876802        _ => panic ! ( "Unexpected event" ) , 
67886803    } 
67896804} 
6805+ 
6806+ #[ test]  
6807+ fn  test_extensions_message_creation ( )  { 
6808+     let  extensions_rpc = RpcOut :: Extensions ( Extensions  { } ) ; 
6809+     let  proto_rpc:  proto:: RPC  = extensions_rpc. into ( ) ; 
6810+ 
6811+     assert ! ( proto_rpc. control. is_some( ) ) ; 
6812+     let  control = proto_rpc. control . unwrap ( ) ; 
6813+     assert ! ( control. extensions. is_some( ) ) ; 
6814+     assert ! ( control. ihave. is_empty( ) ) ; 
6815+     assert ! ( control. iwant. is_empty( ) ) ; 
6816+     assert ! ( control. graft. is_empty( ) ) ; 
6817+     assert ! ( control. prune. is_empty( ) ) ; 
6818+     assert ! ( control. idontwant. is_empty( ) ) ; 
6819+ } 
6820+ 
6821+ #[ test]  
6822+ fn  test_handle_extensions_message ( )  { 
6823+     let  mut  gs:  Behaviour  = Behaviour :: new ( 
6824+         MessageAuthenticity :: Anonymous , 
6825+         ConfigBuilder :: default ( ) 
6826+             . validation_mode ( ValidationMode :: None ) 
6827+             . build ( ) 
6828+             . unwrap ( ) , 
6829+     ) 
6830+     . unwrap ( ) ; 
6831+ 
6832+     let  peer_id = PeerId :: random ( ) ; 
6833+     let  sender = Sender :: new ( gs. config . connection_handler_queue_len ( ) ) ; 
6834+ 
6835+     // Add peer without extensions 
6836+     gs. connected_peers . insert ( 
6837+         peer_id, 
6838+         PeerDetails  { 
6839+             kind :  PeerKind :: Gossipsubv1_3 , 
6840+             connections :  vec ! [ ConnectionId :: new_unchecked( 0 ) ] , 
6841+             outbound :  false , 
6842+             topics :  BTreeSet :: new ( ) , 
6843+             sender, 
6844+             dont_send :  LinkedHashMap :: new ( ) , 
6845+             extensions :  None , 
6846+         } , 
6847+     ) ; 
6848+ 
6849+     // Simulate receiving extensions message 
6850+     let  extensions = Extensions  { } ; 
6851+     gs. handle_extensions ( & peer_id,  extensions) ; 
6852+ 
6853+     // Verify extensions were stored 
6854+     let  peer_details = gs. connected_peers . get ( & peer_id) . unwrap ( ) ; 
6855+     assert ! ( peer_details. extensions. is_some( ) ) ; 
6856+ 
6857+     // Simulate receiving duplicate extensions message from another peer 
6858+     // TODO: when more extensions are added, we should test that they are not overridden. 
6859+     let  duplicate_rpc = RpcIn  { 
6860+         messages :  vec ! [ ] , 
6861+         subscriptions :  vec ! [ ] , 
6862+         control_msgs :  vec ! [ ControlAction :: Extensions ( None ) ] , 
6863+     } ; 
6864+ 
6865+     gs. on_connection_handler_event ( 
6866+         peer_id, 
6867+         ConnectionId :: new_unchecked ( 0 ) , 
6868+         HandlerEvent :: Message  { 
6869+             rpc :  duplicate_rpc, 
6870+             invalid_messages :  vec ! [ ] , 
6871+         } , 
6872+     ) ; 
6873+ 
6874+     // Extensions should still be present (not cleared or changed) 
6875+     let  peer_details = gs. connected_peers . get ( & peer_id) . unwrap ( ) ; 
6876+     assert ! ( peer_details. extensions. is_some( ) ) ; 
6877+ } 
0 commit comments