@@ -111,16 +111,37 @@ impl Writeable for TestCustomMessage {
111
111
}
112
112
113
113
struct TestCustomMessageHandler {
114
- expected_messages : Mutex < VecDeque < TestCustomMessage > > ,
114
+ expected_messages : Mutex < VecDeque < Expectations > > ,
115
+ }
116
+
117
+ struct Expectations {
118
+ message : TestCustomMessage ,
119
+ add_reply_path : bool ,
115
120
}
116
121
117
122
impl TestCustomMessageHandler {
118
123
fn new ( ) -> Self {
119
- Self { expected_messages : Mutex :: new ( VecDeque :: new ( ) ) }
124
+ Self {
125
+ expected_messages : Mutex :: new ( VecDeque :: new ( ) ) ,
126
+ }
120
127
}
121
128
122
129
fn expect_message ( & self , message : TestCustomMessage ) {
123
- self . expected_messages . lock ( ) . unwrap ( ) . push_back ( message) ;
130
+ self . expected_messages . lock ( ) . unwrap ( ) . push_back (
131
+ Expectations {
132
+ message,
133
+ add_reply_path : false ,
134
+ }
135
+ ) ;
136
+ }
137
+
138
+ fn expect_message_and_response ( & self , message : TestCustomMessage ) {
139
+ self . expected_messages . lock ( ) . unwrap ( ) . push_back (
140
+ Expectations {
141
+ message,
142
+ add_reply_path : true ,
143
+ }
144
+ ) ;
124
145
}
125
146
}
126
147
@@ -138,16 +159,23 @@ impl Drop for TestCustomMessageHandler {
138
159
impl CustomOnionMessageHandler for TestCustomMessageHandler {
139
160
type CustomMessage = TestCustomMessage ;
140
161
fn handle_custom_message ( & self , msg : Self :: CustomMessage , responder : Option < Responder > ) -> ResponseInstruction < Self :: CustomMessage > {
141
- match self . expected_messages . lock ( ) . unwrap ( ) . pop_front ( ) {
142
- Some ( expected_msg) => assert_eq ! ( expected_msg, msg) ,
162
+ let add_reply_path = match self . expected_messages . lock ( ) . unwrap ( ) . pop_front ( ) {
163
+ Some ( Expectations { message : expected_msg, add_reply_path } ) => {
164
+ assert_eq ! ( expected_msg, msg) ;
165
+ add_reply_path
166
+ } ,
143
167
None => panic ! ( "Unexpected message: {:?}" , msg) ,
144
- }
168
+ } ;
145
169
let response_option = match msg {
146
170
TestCustomMessage :: Ping => Some ( TestCustomMessage :: Pong ) ,
147
- TestCustomMessage :: Pong => None ,
171
+ TestCustomMessage :: Pong => Some ( TestCustomMessage :: Ping ) ,
148
172
} ;
149
173
if let ( Some ( response) , Some ( responder) ) = ( response_option, responder) {
150
- responder. respond ( response)
174
+ if add_reply_path {
175
+ responder. respond_with_reply_path ( response)
176
+ } else {
177
+ responder. respond ( response)
178
+ }
151
179
} else {
152
180
ResponseInstruction :: NoResponse
153
181
}
@@ -401,6 +429,62 @@ fn async_response_over_one_blinded_hop() {
401
429
pass_along_path ( & nodes) ;
402
430
}
403
431
432
+ fn do_test_async_response_with_reply_path_over_one_blinded_hop ( reply_path_succeed : bool ) {
433
+ // Simulate an asynchronous interaction between two nodes, Alice and Bob.
434
+
435
+ let mut nodes = create_nodes ( 2 ) ;
436
+ let alice = & nodes[ 0 ] ;
437
+ let bob = & nodes[ 1 ] ;
438
+
439
+ // Alice receives a message from Bob with a reply path
440
+ let message = TestCustomMessage :: Ping ;
441
+ let path_id = Some ( [ 2 ; 32 ] ) ;
442
+
443
+ let secp_ctx = Secp256k1 :: new ( ) ;
444
+ let reply_path = BlindedPath :: new_for_message ( & [ bob. node_id ] , & * bob. entropy_source , & secp_ctx) . unwrap ( ) ;
445
+
446
+ if reply_path_succeed {
447
+ // Add a channel so that nodes are announced to each other.
448
+ // This will allow creating the reply path by Alice to include in the response.
449
+ add_channel_to_graph ( alice, bob, & secp_ctx, 24 ) ;
450
+ }
451
+
452
+ let responder = Some ( Responder :: new ( reply_path, path_id) ) ;
453
+ alice. custom_message_handler . expect_message_and_response ( message. clone ( ) ) ;
454
+
455
+ // Alice handles the message reponse, and creates the appropriate ResponseInstruction for it.
456
+ let response_instruction = alice. custom_message_handler . handle_custom_message ( message, responder) ;
457
+
458
+ if !reply_path_succeed {
459
+ // Simulate Alice attempting to asynchronously respond back to Bob
460
+ // but failing to create a reply path.
461
+ assert_eq ! (
462
+ alice. messenger. handle_onion_message_response( response_instruction) ,
463
+ Err ( SendError :: PathNotFound ) ,
464
+ ) ;
465
+ } else {
466
+ // Simulate Alice asynchronously responding back to Bob with a response.
467
+ assert_eq ! (
468
+ alice. messenger. handle_onion_message_response( response_instruction) ,
469
+ Ok ( Some ( SendSuccess :: Buffered ) ) ,
470
+ ) ;
471
+
472
+ bob. custom_message_handler . expect_message ( TestCustomMessage :: Pong ) ;
473
+ pass_along_path ( & nodes) ;
474
+
475
+ // Simulate Bob responding back to Alice through the reply path created by her.
476
+ alice. custom_message_handler . expect_message ( TestCustomMessage :: Ping ) ;
477
+ nodes. reverse ( ) ;
478
+ pass_along_path ( & nodes) ;
479
+ }
480
+ }
481
+
482
+ #[ test]
483
+ fn async_response_with_reply_path_over_one_blinded_hop ( ) {
484
+ do_test_async_response_with_reply_path_over_one_blinded_hop ( true ) ;
485
+ do_test_async_response_with_reply_path_over_one_blinded_hop ( false ) ;
486
+ }
487
+
404
488
#[ test]
405
489
fn too_big_packet_error ( ) {
406
490
// Make sure we error as expected if a packet is too big to send.
0 commit comments