Skip to content

Commit d5573ed

Browse files
committed
refactor: Use a real struct instead of conduit_and_remote_pubkey tuple
This makes the process_act() interface a bit cleaner on a successfully completed handshake.
1 parent 2b8d24d commit d5573ed

File tree

5 files changed

+72
-50
lines changed

5 files changed

+72
-50
lines changed

fuzz/src/peer_crypt.rs

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -177,15 +177,15 @@ fn do_completed_handshake_test(generator: &mut FuzzGen) -> Result<(), GeneratorF
177177
// The handshake should complete with any valid private keys
178178
let act2 = test_ctx.responder_handshake.process_act(&test_ctx.act1).unwrap().0.unwrap();
179179
let (act3, (mut initiator_conduit, responder_pubkey)) = match test_ctx.initiator_handshake.process_act(&act2) {
180-
Ok((Some(act3), Some((conduit, remote_pubkey)))) => {
181-
(act3, (conduit, remote_pubkey))
180+
Ok((Some(act3), Some(completed_handshake_info))) => {
181+
(act3, (completed_handshake_info.conduit, completed_handshake_info.their_node_id))
182182
}
183183
_ => panic!("handshake failed")
184184
};
185185

186186
let (mut responder_conduit, initiator_pubkey) = match test_ctx.responder_handshake.process_act(&act3) {
187-
Ok((None, Some((conduit, remote_pubkey)))) => {
188-
(conduit, remote_pubkey)
187+
Ok((None, Some(completed_handshake_info))) => {
188+
(completed_handshake_info.conduit, completed_handshake_info.their_node_id)
189189
}
190190
_ => panic!("handshake failed")
191191
};
@@ -276,12 +276,12 @@ fn do_handshake_test(generator: &mut FuzzGen) -> Result<(), GeneratorFinishedErr
276276
let act2_chunks = split_vec(generator, &act2)?;
277277

278278
let mut act3_option = None;
279-
let mut conduit_and_remote_pubkey_option = None;
279+
let mut initiator_completed_handshake_info_option = None;
280280
for partial_act2 in act2_chunks {
281281
match test_ctx.initiator_handshake.process_act(&partial_act2) {
282-
Ok((Some(act3), Some(conduit_and_remote_pubkey_option_inner))) => {
282+
Ok((Some(act3), Some(completed_handshake_info_option_inner))) => {
283283
act3_option = Some(act3);
284-
conduit_and_remote_pubkey_option = Some(conduit_and_remote_pubkey_option_inner);
284+
initiator_completed_handshake_info_option = Some(completed_handshake_info_option_inner);
285285

286286
// Valid conduit and pubkey indicates handshake is over
287287
break;
@@ -299,7 +299,7 @@ fn do_handshake_test(generator: &mut FuzzGen) -> Result<(), GeneratorFinishedErr
299299

300300
// Ensure we actually received act3 bytes, conduit, and remote pubkey from process_act()
301301
let act3 = act3_option.unwrap();
302-
let (mut initiator_conduit, responder_pubkey) = conduit_and_remote_pubkey_option.unwrap();
302+
let mut initiator_completed_handshake_info = initiator_completed_handshake_info_option.unwrap();
303303

304304
// Possibly generate bad data for act3
305305
let (mut act3, is_bad_data) = maybe_generate_bad_act(generator, act3)?;
@@ -313,11 +313,11 @@ fn do_handshake_test(generator: &mut FuzzGen) -> Result<(), GeneratorFinishedErr
313313
// Split act3 into between 1..7 chunks
314314
let act3_chunks = split_vec(generator, &act3)?;
315315

316-
let mut conduit_and_remote_pubkey_option = None;
316+
let mut responder_completed_handshake_info = None;
317317
for partial_act3 in act3_chunks {
318318
match test_ctx.responder_handshake.process_act(&partial_act3) {
319-
Ok((None, Some(conduit_and_remote_pubkey_option_inner))) => {
320-
conduit_and_remote_pubkey_option = Some(conduit_and_remote_pubkey_option_inner);
319+
Ok((None, Some(completed_handshake_info_inner))) => {
320+
responder_completed_handshake_info = Some(completed_handshake_info_inner);
321321

322322
// Valid conduit and pubkey indicates handshake is over
323323
break;
@@ -333,20 +333,20 @@ fn do_handshake_test(generator: &mut FuzzGen) -> Result<(), GeneratorFinishedErr
333333
};
334334
}
335335
// Ensure we actually received conduit and remote pubkey from process_act()
336-
let (mut responder_conduit, initiator_pubkey) = conduit_and_remote_pubkey_option.unwrap();
336+
let mut responder_completed_handshake_info = responder_completed_handshake_info.unwrap();
337337

338338
// The handshake should complete with each peer knowing the static_public_key of the remote peer
339-
if initiator_pubkey != test_ctx.initiator_static_public_key {
339+
if responder_completed_handshake_info.their_node_id != test_ctx.initiator_static_public_key {
340340
assert!(used_generated_data);
341341
return Ok(());
342342
}
343-
if responder_pubkey != test_ctx.responder_static_public_key {
343+
if initiator_completed_handshake_info.their_node_id != test_ctx.responder_static_public_key {
344344
assert!(used_generated_data);
345345
return Ok(());
346346
}
347347

348348
// The nodes should be able to communicate over the conduit
349-
do_conduit_tests(generator, &mut initiator_conduit, &mut responder_conduit, used_generated_data)
349+
do_conduit_tests(generator, &mut initiator_completed_handshake_info.conduit, &mut responder_completed_handshake_info.conduit, used_generated_data)
350350
}
351351

352352
#[inline]

lightning/src/ln/peers/handshake/mod.rs

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,12 @@ pub struct PeerHandshake {
3636
ready_to_process: bool,
3737
}
3838

39+
/// Container for the information returned from a successfully completed handshake
40+
pub struct CompletedHandshakeInfo {
41+
pub conduit: Conduit,
42+
pub their_node_id: PublicKey,
43+
}
44+
3945
impl IPeerHandshake for PeerHandshake {
4046
/// Instantiate a handshake given the peer's static public key. The ephemeral private key MUST
4147
/// generate a new session with strong cryptographic randomness.
@@ -75,8 +81,8 @@ impl IPeerHandshake for PeerHandshake {
7581
/// # Return values
7682
/// Returns a tuple with the following components:
7783
/// `.0`: Byte vector containing the next act to send back to the peer per the handshake protocol
78-
/// `.1`: Some(Conduit, PublicKey) if the handshake was just processed to completion and messages can now be encrypted and decrypted
79-
fn process_act(&mut self, input: &[u8]) -> Result<(Option<Vec<u8>>, Option<(Conduit, PublicKey)>), String> {
84+
/// `.1`: Some(CompleteHandshakeInfo) if the handshake was just processed to completion and messages can now be encrypted and decrypted
85+
fn process_act(&mut self, input: &[u8]) -> Result<(Option<Vec<u8>>, Option<CompletedHandshakeInfo>), String> {
8086
assert!(self.ready_to_process);
8187
let cur_state = self.state.take().unwrap();
8288

@@ -86,8 +92,8 @@ impl IPeerHandshake for PeerHandshake {
8692
};
8793

8894
let result = match next_state {
89-
HandshakeState::Complete(ref mut conduit_and_pubkey) => {
90-
Ok((act_option, conduit_and_pubkey.take()))
95+
HandshakeState::Complete(ref mut complete_handshake_info) => {
96+
Ok((act_option, complete_handshake_info.take()))
9197
},
9298
_ => { Ok((act_option, None)) }
9399
};
@@ -197,14 +203,14 @@ mod test {
197203
let mut test_ctx = TestCtx::new();
198204
let act2 = do_process_act_or_panic!(test_ctx.inbound_handshake, &test_ctx.act1);
199205

200-
let (act3, inbound_remote_pubkey) = if let (Some(act3), Some((_, remote_pubkey))) = test_ctx.outbound_handshake.process_act(&act2).unwrap() {
201-
(act3, remote_pubkey)
206+
let (act3, inbound_remote_pubkey) = if let (Some(act3), Some(completed_handshake_info)) = test_ctx.outbound_handshake.process_act(&act2).unwrap() {
207+
(act3, completed_handshake_info.their_node_id)
202208
} else {
203209
panic!();
204210
};
205211

206-
let outbound_remote_pubkey = if let (None, Some((_, remote_pubkey))) = test_ctx.inbound_handshake.process_act(&act3).unwrap() {
207-
remote_pubkey
212+
let outbound_remote_pubkey = if let (None, Some(completed_handshake_info)) = test_ctx.inbound_handshake.process_act(&act3).unwrap() {
213+
completed_handshake_info.their_node_id
208214
} else {
209215
panic!();
210216
};

lightning/src/ln/peers/handshake/states.rs

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ use bitcoin::secp256k1::{SecretKey, PublicKey};
1616
use ln::peers::{chacha, hkdf5869rfc};
1717
use ln::peers::conduit::{Conduit, SymmetricKey};
1818
use ln::peers::handshake::acts::{Act, ActBuilder, ACT_ONE_LENGTH, ACT_TWO_LENGTH, ACT_THREE_LENGTH, EMPTY_ACT_ONE, EMPTY_ACT_TWO, EMPTY_ACT_THREE};
19-
use ln::peers::handshake::IHandshakeState;
19+
use ln::peers::handshake::{CompletedHandshakeInfo, IHandshakeState};
2020

2121
// Alias type to help differentiate between temporary key and chaining key when passing bytes around
2222
type ChainingKey = [u8; 32];
@@ -37,7 +37,7 @@ pub(super) enum HandshakeState {
3737
ResponderAwaitingActOne(ResponderAwaitingActOneState),
3838
InitiatorAwaitingActTwo(InitiatorAwaitingActTwoState),
3939
ResponderAwaitingActThree(ResponderAwaitingActThreeState),
40-
Complete(Option<(Conduit, PublicKey)>),
40+
Complete(Option<CompletedHandshakeInfo>),
4141
}
4242

4343
// Enum dispatch for state machine. Single public interface can statically dispatch to all states
@@ -57,7 +57,7 @@ impl IHandshakeState for HandshakeState {
5757
HandshakeState::ResponderAwaitingActOne(state) => { state.next(input) },
5858
HandshakeState::InitiatorAwaitingActTwo(state) => { state.next(input) },
5959
HandshakeState::ResponderAwaitingActThree(state) => { state.next(input) },
60-
HandshakeState::Complete(_conduit) => { panic!("no acts left to process") }
60+
HandshakeState::Complete(_completed_handshake_info) => { panic!("no acts left to process") }
6161
}
6262
}
6363
}
@@ -319,7 +319,12 @@ impl IHandshakeState for InitiatorAwaitingActTwoState {
319319
act_three[0] = 0;
320320
Ok((
321321
Some(Act::Three(act_three)),
322-
HandshakeState::Complete(Some((conduit, responder_static_public_key)))
322+
HandshakeState::Complete(Some(
323+
CompletedHandshakeInfo {
324+
conduit,
325+
their_node_id: responder_static_public_key
326+
}
327+
))
323328
))
324329
}
325330
}
@@ -399,7 +404,12 @@ impl IHandshakeState for ResponderAwaitingActThreeState {
399404

400405
Ok((
401406
None,
402-
HandshakeState::Complete(Some((conduit, initiator_pubkey)))
407+
HandshakeState::Complete(Some(
408+
CompletedHandshakeInfo {
409+
conduit,
410+
their_node_id: initiator_pubkey
411+
}
412+
))
403413
))
404414
}
405415
}
@@ -693,15 +703,15 @@ mod test {
693703
let (_act1, awaiting_act_two_state) = do_next_or_panic!(test_ctx.initiator, &[]);
694704
let (act3, complete_state) = do_next_or_panic!(awaiting_act_two_state, &test_ctx.valid_act2);
695705

696-
let (conduit, remote_pubkey) = if let Complete(Some((conduit, remote_pubkey))) = complete_state {
697-
(conduit, remote_pubkey)
706+
let completed_handshake_info = if let Complete(Some(completed_handshake_info)) = complete_state {
707+
completed_handshake_info
698708
} else {
699709
panic!();
700710
};
701711

702712
assert_eq!(act3.as_ref(), test_ctx.valid_act3.as_slice());
703-
assert_eq!(remote_pubkey, test_ctx.responder_static_public_key);
704-
assert_eq!(0, conduit.decryptor.read_buffer_length());
713+
assert_eq!(completed_handshake_info.their_node_id, test_ctx.responder_static_public_key);
714+
assert_eq!(0, completed_handshake_info.conduit.decryptor.read_buffer_length());
705715
}
706716

707717
// Initiator::AwaitingActTwo -> Complete (segmented calls)
@@ -760,14 +770,14 @@ mod test {
760770
let test_ctx = TestCtx::new();
761771
let (_act2, awaiting_act_three_state) = do_next_or_panic!(test_ctx.responder, &test_ctx.valid_act1);
762772

763-
let (conduit, remote_pubkey) = if let (None, Complete(Some((conduit, remote_pubkey)))) = awaiting_act_three_state.next(&test_ctx.valid_act3).unwrap() {
764-
(conduit, remote_pubkey)
773+
let completed_handshake_info = if let (None, Complete(Some(completed_handshake_info))) = awaiting_act_three_state.next(&test_ctx.valid_act3).unwrap() {
774+
completed_handshake_info
765775
} else {
766776
panic!();
767777
};
768778

769-
assert_eq!(remote_pubkey, test_ctx.initiator_public_key);
770-
assert_eq!(0, conduit.decryptor.read_buffer_length());
779+
assert_eq!(completed_handshake_info.their_node_id, test_ctx.initiator_public_key);
780+
assert_eq!(0, completed_handshake_info.conduit.decryptor.read_buffer_length());
771781
}
772782

773783
// Responder::AwaitingActThree -> None (with extra bytes)
@@ -780,14 +790,14 @@ mod test {
780790
let mut act3 = test_ctx.valid_act3;
781791
act3.extend_from_slice(&[2; 16]);
782792

783-
let (conduit, remote_pubkey) = if let (None, Complete(Some((conduit, remote_pubkey)))) = awaiting_act_three_state.next(&act3).unwrap() {
784-
(conduit, remote_pubkey)
793+
let completed_handshake_info = if let (None, Complete(Some(completed_handshake_info))) = awaiting_act_three_state.next(&act3).unwrap() {
794+
completed_handshake_info
785795
} else {
786796
panic!();
787797
};
788798

789-
assert_eq!(remote_pubkey, test_ctx.initiator_public_key);
790-
assert_eq!(16, conduit.decryptor.read_buffer_length());
799+
assert_eq!(completed_handshake_info.their_node_id, test_ctx.initiator_public_key);
800+
assert_eq!(16, completed_handshake_info.conduit.decryptor.read_buffer_length());
791801
}
792802

793803
// Responder::AwaitingActThree -> Error (bad version bytes)

lightning/src/ln/peers/test_util.rs

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use bitcoin::secp256k1::key::{PublicKey, SecretKey};
1414

1515
use ln::peers::conduit::Conduit;
1616
use ln::peers::handler::{SocketDescriptor, IOutboundQueue, ITransport, PeerHandleError};
17+
use ln::peers::handshake::CompletedHandshakeInfo;
1718
use ln::peers::transport::IPeerHandshake;
1819

1920
use std::rc::Rc;
@@ -51,7 +52,7 @@ impl IPeerHandshake for PeerHandshakeTestStubFail {
5152
PeerHandshakeTestStubFail { }
5253
}
5354

54-
fn process_act(&mut self, _input: &[u8]) -> Result<(Option<Vec<u8>>, Option<(Conduit, PublicKey)>), String> {
55+
fn process_act(&mut self, _input: &[u8]) -> Result<(Option<Vec<u8>>, Option<CompletedHandshakeInfo>), String> {
5556
Err("Oh no!".to_string())
5657
}
5758
}
@@ -77,7 +78,7 @@ impl IPeerHandshake for PeerHandshakeTestStubBytes {
7778
PeerHandshakeTestStubBytes { }
7879
}
7980

80-
fn process_act(&mut self, _input: &[u8]) -> Result<(Option<Vec<u8>>, Option<(Conduit, PublicKey)>), String> {
81+
fn process_act(&mut self, _input: &[u8]) -> Result<(Option<Vec<u8>>, Option<CompletedHandshakeInfo>), String> {
8182
Ok((Some(Self::RETURNED_BYTES[..].to_vec()), None))
8283
}
8384
}
@@ -98,13 +99,18 @@ impl IPeerHandshake for PeerHandshakeTestStubComplete {
9899
PeerHandshakeTestStubComplete { }
99100
}
100101

101-
fn process_act(&mut self, _input: &[u8]) -> Result<(Option<Vec<u8>>, Option<(Conduit, PublicKey)>), String> {
102+
fn process_act(&mut self, _input: &[u8]) -> Result<(Option<Vec<u8>>, Option<CompletedHandshakeInfo>), String> {
102103
let curve = secp256k1::Secp256k1::new();
103104
let private_key = SecretKey::from_slice(&[0x_21_u8; 32]).unwrap();
104105
let public_key = PublicKey::from_secret_key(&curve, &private_key);
105106
let conduit = Conduit::new([0;32], [0;32], [0;32]);
106107

107-
Ok((None, Some((conduit, public_key))))
108+
Ok((None, Some(
109+
CompletedHandshakeInfo {
110+
conduit,
111+
their_node_id: public_key
112+
}
113+
)))
108114
}
109115
}
110116

lightning/src/ln/peers/transport.rs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use bitcoin::secp256k1::{SecretKey, PublicKey};
1313

1414
use ln::peers::conduit::Conduit;
1515
use ln::peers::handler::{IOutboundQueue, ITransport, PeerHandleError};
16-
use ln::peers::handshake::PeerHandshake;
16+
use ln::peers::handshake::{CompletedHandshakeInfo, PeerHandshake};
1717
use ln::{wire, msgs};
1818
use ln::wire::{Encode, Message};
1919

@@ -34,7 +34,7 @@ pub trait IPeerHandshake {
3434

3535
/// Progress the handshake given bytes received from the peer. Returns Some(Conduit, PublicKey) when the handshake
3636
/// is complete.
37-
fn process_act(&mut self, input: &[u8]) -> Result<(Option<Vec<u8>>, Option<(Conduit, PublicKey)>), String>;
37+
fn process_act(&mut self, input: &[u8]) -> Result<(Option<Vec<u8>>, Option<CompletedHandshakeInfo>), String>;
3838
}
3939

4040
/// Used by the PeerManager to work with a BOLT8 authenticated transport layer. Abstracts the NOISE
@@ -71,17 +71,17 @@ impl<PeerHandshakeImpl: IPeerHandshake> ITransport for Transport<PeerHandshakeIm
7171
match self.conduit {
7272
// Continue handshake
7373
None => {
74-
let (response_option, conduit_and_remote_pubkey_option) = self.handshake.process_act(input)?;
74+
let (response_option, completed_handshake_info_option) = self.handshake.process_act(input)?;
7575

7676
// Any response generated by the handshake sequence is put into the response buffer
7777
if let Some(response) = response_option {
7878
outbound_queue.push_back(response.to_vec());
7979
}
8080

8181
// If handshake is complete change the state
82-
if let Some((conduit, remote_pubkey)) = conduit_and_remote_pubkey_option {
83-
self.conduit = Some(conduit);
84-
self.their_node_id = Some(remote_pubkey);
82+
if let Some(completed_handshake_info) = completed_handshake_info_option {
83+
self.conduit = Some(completed_handshake_info.conduit);
84+
self.their_node_id = Some(completed_handshake_info.their_node_id);
8585
Ok(true)
8686
} else {
8787
Ok(false)

0 commit comments

Comments
 (0)