@@ -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