Skip to content

Commit 243ccd4

Browse files
julianknutsenjulianknutsen
committed
Move Conduit decryption into read()
This gets rid of the complexity required to handle Iterators that return errors and makes way for fewer copies in the decryption path.
1 parent 72a5400 commit 243ccd4

File tree

3 files changed

+57
-152
lines changed

3 files changed

+57
-152
lines changed

‎lightning/src/ln/peers/conduit.rs

Lines changed: 28 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
33
use ln::peers::{chacha, hkdf5869rfc};
44
use util::byte_utils;
5+
use std::collections::VecDeque;
56

67
pub(super) type SymmetricKey = [u8; 32];
78

@@ -38,29 +39,14 @@ pub(super) struct Decryptor {
3839

3940
pending_message_length: Option<usize>,
4041
read_buffer: Option<Vec<u8>>,
41-
poisoned: bool, // signal an error has occurred so None is returned on iteration after failure
42+
decrypted_payloads: VecDeque<Vec<u8>>,
4243
}
4344

4445
impl Iterator for Decryptor {
45-
type Item = Result<Option<Vec<u8>>, String>;
46+
type Item = Vec<u8>;
4647

4748
fn next(&mut self) -> Option<Self::Item> {
48-
if self.poisoned {
49-
return None;
50-
}
51-
52-
match self.decrypt_single_message(None) {
53-
Ok(Some(result)) => {
54-
Some(Ok(Some(result)))
55-
},
56-
Ok(None) => {
57-
None
58-
}
59-
Err(e) => {
60-
self.poisoned = true;
61-
Some(Err(e))
62-
}
63-
}
49+
self.decrypted_payloads.pop_front()
6450
}
6551
}
6652

@@ -79,7 +65,7 @@ impl Conduit {
7965
receiving_nonce: 0,
8066
read_buffer: None,
8167
pending_message_length: None,
82-
poisoned: false
68+
decrypted_payloads: VecDeque::new(),
8369
}
8470
}
8571
}
@@ -89,7 +75,7 @@ impl Conduit {
8975
self.encryptor.encrypt(buffer)
9076
}
9177

92-
pub(super) fn read(&mut self, data: &[u8]) {
78+
pub(super) fn read(&mut self, data: &[u8]) -> Result<(), String>{
9379
self.decryptor.read(data)
9480
}
9581

@@ -143,9 +129,25 @@ impl Encryptor {
143129
}
144130

145131
impl Decryptor {
146-
pub(super) fn read(&mut self, data: &[u8]) {
147-
let read_buffer = self.read_buffer.get_or_insert(Vec::new());
148-
read_buffer.extend_from_slice(data);
132+
pub(super) fn read(&mut self, data: &[u8]) -> Result<(), String> {
133+
let mut input_data = Some(data);
134+
135+
loop {
136+
match self.decrypt_single_message(input_data) {
137+
Ok(Some(result)) => {
138+
self.decrypted_payloads.push_back(result);
139+
},
140+
Ok(None) => {
141+
break;
142+
}
143+
Err(e) => {
144+
return Err(e);
145+
}
146+
}
147+
input_data = None;
148+
}
149+
150+
Ok(())
149151
}
150152

151153
/// Decrypt a single message. If data containing more than one message has been received,
@@ -332,7 +334,7 @@ mod tests {
332334
let encrypted = remote_peer.encrypt(&[1]);
333335

334336
connected_peer.decryptor.receiving_key = [0; 32];
335-
assert_eq!(connected_peer.decrypt_single_message(Some(&encrypted)), Err("invalid hmac".to_string()));
337+
assert_eq!(connected_peer.read(&encrypted), Err("invalid hmac".to_string()));
336338
}
337339

338340
// Test next()::None
@@ -348,86 +350,9 @@ mod tests {
348350
fn decryptor_iterator_one_item_valid() {
349351
let (mut connected_peer, mut remote_peer) = setup_peers();
350352
let encrypted = remote_peer.encrypt(&[1]);
351-
connected_peer.read(&encrypted);
352-
353-
assert_eq!(connected_peer.decryptor.next(), Some(Ok(Some(vec![1]))));
354-
assert_eq!(connected_peer.decryptor.next(), None);
355-
}
356-
357-
// Test next()::err -> next()::None
358-
#[test]
359-
fn decryptor_iterator_error() {
360-
let (mut connected_peer, mut remote_peer) = setup_peers();
361-
let encrypted = remote_peer.encrypt(&[1]);
362-
connected_peer.read(&encrypted);
363-
364-
connected_peer.decryptor.receiving_key = [0; 32];
365-
assert_eq!(connected_peer.decryptor.next(), Some(Err("invalid hmac".to_string())));
366-
assert_eq!(connected_peer.decryptor.next(), None);
367-
}
368-
369-
// Test next()::Some -> next()::err -> next()::None
370-
#[test]
371-
fn decryptor_iterator_error_after_success() {
372-
let (mut connected_peer, mut remote_peer) = setup_peers();
373-
let encrypted = remote_peer.encrypt(&[1]);
374-
connected_peer.read(&encrypted);
375-
let encrypted = remote_peer.encrypt(&[2]);
376-
connected_peer.read(&encrypted);
377-
378-
assert_eq!(connected_peer.decryptor.next(), Some(Ok(Some(vec![1]))));
379-
connected_peer.decryptor.receiving_key = [0; 32];
380-
assert_eq!(connected_peer.decryptor.next(), Some(Err("invalid hmac".to_string())));
381-
assert_eq!(connected_peer.decryptor.next(), None);
382-
}
383-
384-
// Test that next()::Some -> next()::err -> next()::None
385-
// Error should poison decryptor
386-
#[test]
387-
fn decryptor_iterator_next_after_error_returns_none() {
388-
let (mut connected_peer, mut remote_peer) = setup_peers();
389-
let encrypted = remote_peer.encrypt(&[1]);
390-
connected_peer.read(&encrypted);
391-
let encrypted = remote_peer.encrypt(&[2]);
392-
connected_peer.read(&encrypted);
393-
let encrypted = remote_peer.encrypt(&[3]);
394-
connected_peer.read(&encrypted);
395-
396-
// Get one valid value
397-
assert_eq!(connected_peer.decryptor.next(), Some(Ok(Some(vec![1]))));
398-
let valid_receiving_key = connected_peer.decryptor.receiving_key;
399-
400-
// Corrupt the receiving key and ensure we get a failure
401-
connected_peer.decryptor.receiving_key = [0; 32];
402-
assert_eq!(connected_peer.decryptor.next(), Some(Err("invalid hmac".to_string())));
403-
404-
// Restore the receiving key, do a read and ensure None is returned (poisoned)
405-
connected_peer.decryptor.receiving_key = valid_receiving_key;
406-
assert_eq!(connected_peer.decryptor.next(), None);
407-
}
408-
409-
// Test next()::Some -> next()::err -> read() -> next()::None
410-
// Error should poison decryptor even after future reads
411-
#[test]
412-
fn decryptor_iterator_read_next_after_error_returns_none() {
413-
let (mut connected_peer, mut remote_peer) = setup_peers();
414-
let encrypted = remote_peer.encrypt(&[1]);
415-
connected_peer.read(&encrypted);
416-
let encrypted = remote_peer.encrypt(&[2]);
417-
connected_peer.read(&encrypted);
418-
419-
// Get one valid value
420-
assert_eq!(connected_peer.decryptor.next(), Some(Ok(Some(vec![1]))));
421-
let valid_receiving_key = connected_peer.decryptor.receiving_key;
422-
423-
// Corrupt the receiving key and ensure we get a failure
424-
connected_peer.decryptor.receiving_key = [0; 32];
425-
assert_eq!(connected_peer.decryptor.next(), Some(Err("invalid hmac".to_string())));
353+
connected_peer.read(&encrypted).unwrap();
426354

427-
// Restore the receiving key, do a read and ensure None is returned (poisoned)
428-
let encrypted = remote_peer.encrypt(&[3]);
429-
connected_peer.read(&encrypted);
430-
connected_peer.decryptor.receiving_key = valid_receiving_key;
355+
assert_eq!(connected_peer.decryptor.next(), Some(vec![1]));
431356
assert_eq!(connected_peer.decryptor.next(), None);
432357
}
433358

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -386,7 +386,7 @@ impl IHandshakeState for ResponderAwaitingActThreeState {
386386

387387
// Any remaining data in the read buffer would be encrypted, so transfer ownership
388388
// to the Conduit for future use.
389-
conduit.read(&input[bytes_read..]);
389+
conduit.read(&input[bytes_read..])?;
390390

391391
Ok((
392392
None,
@@ -769,7 +769,7 @@ mod test {
769769
let test_ctx = TestCtx::new();
770770
let (_act2, awaiting_act_three_state) = do_next_or_panic!(test_ctx.responder, &test_ctx.valid_act1);
771771
let mut act3 = test_ctx.valid_act3;
772-
act3.extend_from_slice(&[2; 100]);
772+
act3.extend_from_slice(&[2; 16]);
773773

774774
let (conduit, remote_pubkey) = if let (None, Complete(Some((conduit, remote_pubkey)))) = awaiting_act_three_state.next(&act3).unwrap() {
775775
(conduit, remote_pubkey)
@@ -778,7 +778,7 @@ mod test {
778778
};
779779

780780
assert_eq!(remote_pubkey, test_ctx.initiator_public_key);
781-
assert_eq!(100, conduit.decryptor.read_buffer_length());
781+
assert_eq!(16, conduit.decryptor.read_buffer_length());
782782
}
783783

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

‎lightning/src/ln/peers/transport.rs

Lines changed: 26 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ impl<PeerHandshakeImpl: IPeerHandshake> ITransport for Transport<PeerHandshakeIm
9191
}
9292
}
9393
Some(ref mut conduit) => {
94-
conduit.read(input);
94+
conduit.read(input)?;
9595
Ok(false) // newly connected
9696
}
9797
}
@@ -105,53 +105,33 @@ impl<PeerHandshakeImpl: IPeerHandshake> ITransport for Transport<PeerHandshakeIm
105105
match self.conduit {
106106
None => {}
107107
Some(ref mut conduit) => {
108-
// Using Iterators that can error requires special handling
109-
// The item returned from next() has type Option<Result<Option<Vec>, String>>
110-
// The Some wrapper is stripped for each item inside the loop
111-
// There are 3 valid match cases:
112-
// 1) Some(Ok(Some(msg_data))) => Indicates a valid decrypted msg accessed via msg_data
113-
// 2) Some(Err(_)) => Indicates an error during decryption that should be handled
114-
// 3) None -> Indicates there were no messages available to decrypt
115-
// Invalid Cases
116-
// 1) Some(Ok(None)) => Translated to None case above so users of iterators can stop correctly
117-
for msg_data_result in &mut conduit.decryptor {
118-
match msg_data_result {
119-
Ok(Some(msg_data)) => {
120-
let mut reader = ::std::io::Cursor::new(&msg_data[..]);
121-
let message_result = wire::read(&mut reader);
122-
let message = match message_result {
123-
Ok(x) => x,
124-
Err(e) => {
125-
match e {
126-
msgs::DecodeError::UnknownVersion => return Err(PeerHandleError { no_connection_possible: false }),
127-
msgs::DecodeError::UnknownRequiredFeature => {
128-
log_debug!(logger, "Got a channel/node announcement with an known required feature flag, you may want to update!");
129-
continue;
130-
}
131-
msgs::DecodeError::InvalidValue => {
132-
log_debug!(logger, "Got an invalid value while deserializing message");
133-
return Err(PeerHandleError { no_connection_possible: false });
134-
}
135-
msgs::DecodeError::ShortRead => {
136-
log_debug!(logger, "Deserialization failed due to shortness of message");
137-
return Err(PeerHandleError { no_connection_possible: false });
138-
}
139-
msgs::DecodeError::BadLengthDescriptor => return Err(PeerHandleError { no_connection_possible: false }),
140-
msgs::DecodeError::Io(_) => return Err(PeerHandleError { no_connection_possible: false }),
141-
}
142-
}
143-
};
144-
145-
received_messages.push(message);
146-
},
108+
for msg_data in &mut conduit.decryptor {
109+
let mut reader = ::std::io::Cursor::new(&msg_data[..]);
110+
let message_result = wire::read(&mut reader);
111+
let message = match message_result {
112+
Ok(x) => x,
147113
Err(e) => {
148-
log_trace!(logger, "Message decryption failed due to: {}", e);
149-
return Err(PeerHandleError { no_connection_possible: false });
150-
}
151-
Ok(None) => {
152-
panic!("Invalid behavior. Conduit iterator should never return this match.")
114+
match e {
115+
msgs::DecodeError::UnknownVersion => return Err(PeerHandleError { no_connection_possible: false }),
116+
msgs::DecodeError::UnknownRequiredFeature => {
117+
log_debug!(logger, "Got a channel/node announcement with an known required feature flag, you may want to update!");
118+
continue;
119+
}
120+
msgs::DecodeError::InvalidValue => {
121+
log_debug!(logger, "Got an invalid value while deserializing message");
122+
return Err(PeerHandleError { no_connection_possible: false });
123+
}
124+
msgs::DecodeError::ShortRead => {
125+
log_debug!(logger, "Deserialization failed due to shortness of message");
126+
return Err(PeerHandleError { no_connection_possible: false });
127+
}
128+
msgs::DecodeError::BadLengthDescriptor => return Err(PeerHandleError { no_connection_possible: false }),
129+
msgs::DecodeError::Io(_) => return Err(PeerHandleError { no_connection_possible: false }),
130+
}
153131
}
154-
}
132+
};
133+
134+
received_messages.push(message);
155135
}
156136
}
157137
}

0 commit comments

Comments
 (0)