@@ -112,16 +112,31 @@ impl Writeable for TestCustomMessage {
112
112
113
113
struct TestCustomMessageHandler {
114
114
expected_messages : Mutex < VecDeque < TestCustomMessage > > ,
115
+ add_reply_path : Mutex < bool > ,
115
116
}
116
117
117
118
impl TestCustomMessageHandler {
118
119
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
+ }
120
124
}
121
125
122
126
fn expect_message ( & self , message : TestCustomMessage ) {
123
127
self . expected_messages . lock ( ) . unwrap ( ) . push_back ( message) ;
124
128
}
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
+ }
125
140
}
126
141
127
142
impl Drop for TestCustomMessageHandler {
@@ -144,10 +159,14 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
144
159
}
145
160
let response_option = match msg {
146
161
TestCustomMessage :: Ping => Some ( TestCustomMessage :: Pong ) ,
147
- TestCustomMessage :: Pong => None ,
162
+ TestCustomMessage :: Pong => Some ( TestCustomMessage :: Ping ) ,
148
163
} ;
149
164
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
+ }
151
170
} else {
152
171
ResponseInstruction :: NoResponse
153
172
}
@@ -401,6 +420,63 @@ fn async_response_over_one_blinded_hop() {
401
420
pass_along_path ( & nodes) ;
402
421
}
403
422
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
+
404
480
#[ test]
405
481
fn too_big_packet_error ( ) {
406
482
// Make sure we error as expected if a packet is too big to send.
0 commit comments