@@ -112,16 +112,31 @@ impl Writeable for TestCustomMessage {
112112
113113struct TestCustomMessageHandler {
114114 expected_messages : Mutex < VecDeque < TestCustomMessage > > ,
115+ add_reply_path : Mutex < bool > ,
115116}
116117
117118impl TestCustomMessageHandler {
118119 fn new ( ) -> Self {
119- Self { expected_messages : Mutex :: new ( VecDeque :: new ( ) ) }
120+ Self {
121+ expected_messages : Mutex :: new ( VecDeque :: new ( ) ) ,
122+ add_reply_path : Mutex :: new ( false ) ,
123+ }
120124 }
121125
122126 fn expect_message ( & self , message : TestCustomMessage ) {
123127 self . expected_messages . lock ( ) . unwrap ( ) . push_back ( message) ;
124128 }
129+
130+ fn include_reply_path ( & self ) {
131+ * self . add_reply_path . lock ( ) . unwrap ( ) = true ;
132+ }
133+
134+ fn read_and_reset_add_reply_path ( & self ) -> bool {
135+ let mut add_reply_path = self . add_reply_path . lock ( ) . unwrap ( ) ;
136+ let was_set = * add_reply_path;
137+ * add_reply_path = false ;
138+ was_set
139+ }
125140}
126141
127142impl Drop for TestCustomMessageHandler {
@@ -144,10 +159,14 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
144159 }
145160 let response_option = match msg {
146161 TestCustomMessage :: Ping => Some ( TestCustomMessage :: Pong ) ,
147- TestCustomMessage :: Pong => None ,
162+ TestCustomMessage :: Pong => Some ( TestCustomMessage :: Ping ) ,
148163 } ;
149164 if let ( Some ( response) , Some ( responder) ) = ( response_option, responder) {
150- responder. respond ( response)
165+ if self . read_and_reset_add_reply_path ( ) {
166+ responder. respond_with_reply_path ( response)
167+ } else {
168+ responder. respond ( response)
169+ }
151170 } else {
152171 ResponseInstruction :: NoResponse
153172 }
@@ -401,6 +420,63 @@ fn async_response_over_one_blinded_hop() {
401420 pass_along_path ( & nodes) ;
402421}
403422
423+ fn do_test_async_response_with_reply_path_over_one_blinded_hop ( reply_path_succeed : bool ) {
424+ // Simulate an asynchronous interaction between two nodes, Alice and Bob.
425+
426+ let mut nodes = create_nodes ( 2 ) ;
427+ let alice = & nodes[ 0 ] ;
428+ let bob = & nodes[ 1 ] ;
429+
430+ // Alice receives a message from Bob with a reply path
431+ let message = TestCustomMessage :: Ping ;
432+ let path_id = Some ( [ 2 ; 32 ] ) ;
433+
434+ let secp_ctx = Secp256k1 :: new ( ) ;
435+ let reply_path = BlindedPath :: new_for_message ( & [ bob. node_id ] , & * bob. entropy_source , & secp_ctx) . unwrap ( ) ;
436+
437+ if reply_path_succeed {
438+ // Add a channel so that nodes are announced to each other.
439+ // This will allow creating the reply path by Alice to include in the response.
440+ add_channel_to_graph ( alice, bob, & secp_ctx, 24 ) ;
441+ }
442+
443+ let responder = Some ( Responder :: new ( reply_path, path_id) ) ;
444+ alice. custom_message_handler . expect_message ( message. clone ( ) ) ;
445+ alice. custom_message_handler . include_reply_path ( ) ;
446+
447+ // Alice handles the message reponse, and creates the appropriate ResponseInstruction for it.
448+ let response_instruction = alice. custom_message_handler . handle_custom_message ( message, responder) ;
449+
450+ if !reply_path_succeed {
451+ // Simulate Alice attempting to asynchronously respond back to Bob
452+ // but failing to create a reply path.
453+ assert_eq ! (
454+ alice. messenger. handle_onion_message_response( response_instruction) ,
455+ Err ( SendError :: PathNotFound ) ,
456+ ) ;
457+ } else {
458+ // Simulate Alice asynchronously responding back to Bob with a response.
459+ assert_eq ! (
460+ alice. messenger. handle_onion_message_response( response_instruction) ,
461+ Ok ( Some ( SendSuccess :: Buffered ) ) ,
462+ ) ;
463+
464+ bob. custom_message_handler . expect_message ( TestCustomMessage :: Pong ) ;
465+ pass_along_path ( & nodes) ;
466+
467+ // Simulate Bob responding back to Alice through the reply path created by her.
468+ alice. custom_message_handler . expect_message ( TestCustomMessage :: Ping ) ;
469+ nodes. reverse ( ) ;
470+ pass_along_path ( & nodes) ;
471+ }
472+ }
473+
474+ #[ test]
475+ fn async_response_with_reply_path_over_one_blinded_hop ( ) {
476+ do_test_async_response_with_reply_path_over_one_blinded_hop ( true ) ;
477+ do_test_async_response_with_reply_path_over_one_blinded_hop ( false ) ;
478+ }
479+
404480#[ test]
405481fn too_big_packet_error ( ) {
406482 // Make sure we error as expected if a packet is too big to send.
0 commit comments