address some of Jeff's comments pertaining to message decryption, constants, and type definitions

Author: arik-so

lightning/src/ln/peer_handler.rs

@@ -108,16 +108,15 @@ enum InitSyncTracker{
 }
 
 enum PeerState {
-	Handshake(PeerHandshake),
+	Authenticating(PeerHandshake),
 	Connected(Conduit),
 }
 
 impl PeerState {
 	fn is_ready_for_encryption(&self) -> bool {
-		if let &PeerState::Connected(_) = self {
-			true
-		} else {
-			false
+		match self {
+			&PeerState::Connected(_) => true,
+			_ => false
 		}
 	}
 }
@@ -283,7 +282,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref> PeerManager<Descriptor, CM> where
 
 		let mut peers = self.peers.lock().unwrap();
 		if peers.peers.insert(descriptor, Peer {
-			encryptor: PeerState::Handshake(handshake),
+			encryptor: PeerState::Authenticating(handshake),
 			outbound: true,
 			their_node_id: Some(their_node_id.clone()),
 			their_features: None,
@@ -318,7 +317,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref> PeerManager<Descriptor, CM> where
 
 		let mut peers = self.peers.lock().unwrap();
 		if peers.peers.insert(descriptor, Peer {
-			encryptor: PeerState::Handshake(handshake),
+			encryptor: PeerState::Authenticating(handshake),
 			outbound: false,
 			their_node_id: None,
 			their_features: None,
@@ -469,9 +468,9 @@ impl<Descriptor: SocketDescriptor, CM: Deref> PeerManager<Descriptor, CM> where
 						let mut conduit_option = None;
 						let mut remote_pubkey_option = None;
 
-						if let &mut PeerState::Handshake(ref mut handshake) = &mut peer.encryptor {
+						if let &mut PeerState::Authenticating(ref mut handshake) = &mut peer.encryptor {
 							let (response, conduit, remote_pubkey) = handshake.process_act(&peer.pending_read_buffer, None).unwrap();
-							peer.pending_read_buffer.drain(..); // we read it all
+							peer.pending_read_buffer = Vec::new(); // empty the pending read buffer
 
 							if let Some(key) = remote_pubkey {
 								remote_pubkey_option = Some(key);
@@ -531,18 +530,12 @@ impl<Descriptor: SocketDescriptor, CM: Deref> PeerManager<Descriptor, CM> where
 								}
 							}
 
-							let mut next_message_result = conduit.decrypt(&peer.pending_read_buffer);
-
-							let offset = next_message_result.1;
-							if offset == 0 {
-								// nothing got read
+							let message_option = conduit.decrypt_single_message(Some(&peer.pending_read_buffer));
+							let msg_data = if let Some(message) = message_option {
+								message
+							} else {
 								break;
-							}else{
-								peer.pending_read_buffer.drain(0..offset);
-							}
-
-							// something got read, so we definitely have a message
-							let msg_data = next_message_result.0.unwrap();
+							};
 
 							let mut reader = ::std::io::Cursor::new(&msg_data[..]);
 							let message_result = wire::read(&mut reader);

lightning/src/ln/peers/chacha.rs

@@ -1,6 +1,8 @@
 use util::byte_utils;
 use util::chacha20poly1305rfc::ChaCha20Poly1305RFC;
 
+pub const TAG_SIZE: usize = 16;
+
 pub fn encrypt(key: &[u8], nonce: u64, associated_data: &[u8], plaintext: &[u8]) -> Vec<u8> {
 	let mut nonce_bytes = [0; 12];
 	nonce_bytes[4..].copy_from_slice(&byte_utils::le64_to_array(nonce));
@@ -15,7 +17,6 @@ pub fn encrypt(key: &[u8], nonce: u64, associated_data: &[u8], plaintext: &[u8])
 	tagged_ciphertext
 }
 
-
 pub fn decrypt(key: &[u8], nonce: u64, associated_data: &[u8], tagged_ciphertext: &[u8]) -> Result<Vec<u8>, String> {
 	let mut nonce_bytes = [0; 12];
 	nonce_bytes[4..].copy_from_slice(&byte_utils::le64_to_array(nonce));

lightning/src/ln/peers/conduit.rs

@@ -3,16 +3,20 @@
 use ln::peers::{chacha, hkdf};
 use util::byte_utils;
 
+type SymmetricKey = [u8; 32];
+const MESSAGE_LENGTH_HEADER_SIZE: usize = 2;
+const TAGGED_MESSAGE_LENGTH_HEADER_SIZE: usize = MESSAGE_LENGTH_HEADER_SIZE + chacha::TAG_SIZE;
+
 /// Returned after a successful handshake to encrypt and decrypt communication with peer nodes.
 /// It should not normally be manually instantiated.
 /// Automatically handles key rotation.
 /// For decryption, it is recommended to call `decrypt_message_stream` for automatic buffering.
 pub struct Conduit {
-	pub(crate) sending_key: [u8; 32],
-	pub(crate) receiving_key: [u8; 32],
+	pub(crate) sending_key: SymmetricKey,
+	pub(crate) receiving_key: SymmetricKey,
 
-	pub(crate) sending_chaining_key: [u8; 32],
-	pub(crate) receiving_chaining_key: [u8; 32],
+	pub(crate) sending_chaining_key: SymmetricKey,
+	pub(crate) receiving_chaining_key: SymmetricKey,
 
 	pub(crate) receiving_nonce: u32,
 	pub(crate) sending_nonce: u32,
@@ -26,12 +30,12 @@ impl Conduit {
 		let length = buffer.len() as u16;
 		let length_bytes = byte_utils::be16_to_array(length);
 
-		let mut ciphertext = vec![0u8; 18 + length as usize + 16];
+		let mut ciphertext = vec![0u8; TAGGED_MESSAGE_LENGTH_HEADER_SIZE + length as usize + chacha::TAG_SIZE];
 
-		ciphertext[0..18].copy_from_slice(&chacha::encrypt(&self.sending_key, self.sending_nonce as u64, &[0; 0], &length_bytes));
+		ciphertext[0..TAGGED_MESSAGE_LENGTH_HEADER_SIZE].copy_from_slice(&chacha::encrypt(&self.sending_key, self.sending_nonce as u64, &[0; 0], &length_bytes));
 		self.increment_sending_nonce();
 
-		ciphertext[18..].copy_from_slice(&chacha::encrypt(&self.sending_key, self.sending_nonce as u64, &[0; 0], buffer));
+		ciphertext[TAGGED_MESSAGE_LENGTH_HEADER_SIZE..].copy_from_slice(&chacha::encrypt(&self.sending_key, self.sending_nonce as u64, &[0; 0], buffer));
 		self.increment_sending_nonce();
 
 		ciphertext
@@ -42,8 +46,11 @@ impl Conduit {
 		read_buffer.extend_from_slice(data);
 	}
 
-	/// Add newly received data from the peer node to the buffer and decrypt all possible messages
-	pub fn decrypt_message_stream(&mut self, new_data: Option<&[u8]>) -> Vec<Vec<u8>> {
+	/// Decrypt a single message. If data containing more than one message has been received,
+	/// only the first message will be returned, and the rest stored in the internal buffer.
+	/// If a message pending in the buffer still hasn't been decrypted, that message will be
+	/// returned in lieu of anything new, even if new data is provided.
+	pub fn decrypt_single_message(&mut self, new_data: Option<&[u8]>) -> Option<Vec<u8>> {
 		let mut read_buffer = if let Some(buffer) = self.read_buffer.take() {
 			buffer
 		} else {
@@ -54,45 +61,32 @@ impl Conduit {
 			read_buffer.extend_from_slice(data);
 		}
 
-		let mut messages = Vec::new();
-
-		loop {
-			// todo: find way that won't require cloning the entire buffer
-			let (current_message, offset) = self.decrypt(&read_buffer[..]);
-			if offset == 0 {
-				break;
-			}
-
-			read_buffer.drain(0..offset);
-
-			if let Some(message) = current_message {
-				messages.push(message);
-			} else {
-				break;
-			}
+		let (current_message, offset) = self.decrypt(&read_buffer[..]);
+		if offset == 0 {
+			return None;
 		}
 
-		messages
+		read_buffer.drain(0..offset);
+		current_message
 	}
 
-	/// Decrypt a single message. Buffer is an undelimited amount of bytes
-	pub(crate) fn decrypt(&mut self, buffer: &[u8]) -> (Option<Vec<u8>>, usize) { // the response slice should have the same lifetime as the argument. It's the slice data is read from
-		if buffer.len() < 18 {
+	fn decrypt(&mut self, buffer: &[u8]) -> (Option<Vec<u8>>, usize) { // the response slice should have the same lifetime as the argument. It's the slice data is read from
+		if buffer.len() < TAGGED_MESSAGE_LENGTH_HEADER_SIZE {
 			return (None, 0);
 		}
 
-		let encrypted_length = &buffer[0..18]; // todo: abort if too short
+		let encrypted_length = &buffer[0..TAGGED_MESSAGE_LENGTH_HEADER_SIZE]; // todo: abort if too short
 		let length_vec = chacha::decrypt(&self.receiving_key, self.receiving_nonce as u64, &[0; 0], encrypted_length).unwrap();
-		let mut length_bytes = [0u8; 2];
+		let mut length_bytes = [0u8; MESSAGE_LENGTH_HEADER_SIZE];
 		length_bytes.copy_from_slice(length_vec.as_slice());
 		let message_length = byte_utils::slice_to_be16(&length_bytes) as usize;
 
-		let message_end_index = message_length + 18 + 16; // todo: abort if too short
+		let message_end_index = message_length + TAGGED_MESSAGE_LENGTH_HEADER_SIZE + chacha::TAG_SIZE; // todo: abort if too short
 		if buffer.len() < message_end_index {
 			return (None, 0);
 		}
 
-		let encrypted_message = &buffer[18..message_end_index];
+		let encrypted_message = &buffer[TAGGED_MESSAGE_LENGTH_HEADER_SIZE..message_end_index];
 
 		self.increment_receiving_nonce();
 
@@ -111,15 +105,15 @@ impl Conduit {
 		Self::increment_nonce(&mut self.receiving_nonce, &mut self.receiving_chaining_key, &mut self.receiving_key);
 	}
 
-	fn increment_nonce(nonce: &mut u32, chaining_key: &mut [u8; 32], key: &mut [u8; 32]) {
+	fn increment_nonce(nonce: &mut u32, chaining_key: &mut SymmetricKey, key: &mut SymmetricKey) {
 		*nonce += 1;
 		if *nonce == 1000 {
 			Self::rotate_key(chaining_key, key);
 			*nonce = 0;
 		}
 	}
 
-	fn rotate_key(chaining_key: &mut [u8; 32], key: &mut [u8; 32]) {
+	fn rotate_key(chaining_key: &mut SymmetricKey, key: &mut SymmetricKey) {
 		let (new_chaining_key, new_key) = hkdf::derive(chaining_key, key);
 		chaining_key.copy_from_slice(&new_chaining_key);
 		key.copy_from_slice(&new_key);
@@ -132,6 +126,7 @@ mod tests {
 	use ln::peers::conduit::Conduit;
 
 	#[test]
+	/// Based on RFC test vectors: https://github.com/lightningnetwork/lightning-rfc/blob/master/08-transport.md#message-encryption-tests
 	fn test_chaining() {
 		let chaining_key_vec = hex::decode("919219dbb2920afa8db80f9a51787a840bcf111ed8d588caf9ab4be716e42b01").unwrap();
 		let mut chaining_key = [0u8; 32];
@@ -182,9 +177,7 @@ mod tests {
 
 		for _ in 0..1002 {
 			let encrypted_message = encrypted_messages.remove(0);
-			let mut decrypted_messages = remote_peer.decrypt_message_stream(Some(&encrypted_message));
-			assert_eq!(decrypted_messages.len(), 1);
-			let decrypted_message = decrypted_messages.remove(0);
+			let decrypted_message = remote_peer.decrypt_single_message(Some(&encrypted_message)).unwrap();
 			assert_eq!(decrypted_message, hex::decode("68656c6c6f").unwrap());
 		}
 	}

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

@@ -32,13 +32,12 @@ pub struct PeerHandshake {
 impl PeerHandshake {
 	/// Instantiate a new handshake with a node identity secret key and an ephemeral private key
 	pub fn new(private_key: &SecretKey, ephemeral_private_key: &SecretKey) -> Self {
-		let handshake = PeerHandshake {
+		Self {
 			state: Some(HandshakeState::Blank),
 			private_key: (*private_key).clone(),
 			ephemeral_private_key: (*ephemeral_private_key).clone(),
 			read_buffer: Vec::new(),
-		};
-		handshake
+		}
 	}
 
 	/// Make the handshake object inbound in anticipation of a peer's first handshake act
@@ -52,7 +51,6 @@ impl PeerHandshake {
 	}
 
 	fn initialize_state(public_key: &PublicKey) -> (HandshakeHash, [u8; 32]) {
-		// do the proper initialization
 		let protocol_name = b"Noise_XK_secp256k1_ChaChaPoly_SHA256";
 		let prologue = b"lightning";
 
@@ -66,7 +64,7 @@ impl PeerHandshake {
 		let mut hash = HandshakeHash::new(initial_hash_preimage.as_slice());
 		hash.update(&public_key.serialize());
 
-		(hash, chaining_key) // hash, chaining_key
+		(hash, chaining_key)
 	}
 
 	/// Process act dynamically
@@ -154,7 +152,7 @@ impl PeerHandshake {
 		let (mut hash, chaining_key) = Self::initialize_state(&remote_public_key);
 
 		// serialize act one
-		let (act_one, chaining_key, temporary_key) = self.calculate_act_message(
+		let (act_one, chaining_key, temporary_key) = Self::calculate_act_message(
 			&self.ephemeral_private_key,
 			remote_public_key,
 			chaining_key,
@@ -193,7 +191,7 @@ impl PeerHandshake {
 		};
 
 		let mut hash = act_one_expectation.hash;
-		let (remote_ephemeral_public_key, chaining_key, _) = self.process_act_message(
+		let (remote_ephemeral_public_key, chaining_key, _) = Self::process_act_message(
 			act.0,
 			&self.private_key,
 			act_one_expectation.chaining_key,
@@ -202,7 +200,7 @@ impl PeerHandshake {
 
 		let ephemeral_private_key = (*&self.ephemeral_private_key).clone();
 
-		let (act_two, chaining_key, temporary_key) = self.calculate_act_message(
+		let (act_two, chaining_key, temporary_key) = Self::calculate_act_message(
 			&ephemeral_private_key,
 			&remote_ephemeral_public_key,
 			chaining_key,
@@ -232,7 +230,7 @@ impl PeerHandshake {
 		};
 
 		let mut hash = act_two_expectation.hash;
-		let (remote_ephemeral_public_key, chaining_key, temporary_key) = self.process_act_message(
+		let (remote_ephemeral_public_key, chaining_key, temporary_key) = Self::process_act_message(
 			act.0,
 			&act_two_expectation.ephemeral_private_key,
 			act_two_expectation.chaining_key,
@@ -317,7 +315,7 @@ impl PeerHandshake {
 		Ok((remote_pubkey, connected_peer))
 	}
 
-	fn calculate_act_message(&self, local_private_key: &SecretKey, remote_public_key: &PublicKey, chaining_key: [u8; 32], hash: &mut HandshakeHash) -> ([u8; 50], [u8; 32], [u8; 32]) {
+	fn calculate_act_message(local_private_key: &SecretKey, remote_public_key: &PublicKey, chaining_key: [u8; 32], hash: &mut HandshakeHash) -> ([u8; 50], [u8; 32], [u8; 32]) {
 		let local_public_key = Self::private_key_to_public_key(local_private_key);
 
 		hash.update(&local_public_key.serialize());
@@ -335,7 +333,7 @@ impl PeerHandshake {
 		(act, chaining_key, temporary_key)
 	}
 
-	fn process_act_message(&self, act_bytes: [u8; 50], local_private_key: &SecretKey, chaining_key: [u8; 32], hash: &mut HandshakeHash) -> Result<(PublicKey, [u8; 32], [u8; 32]), String> {
+	fn process_act_message(act_bytes: [u8; 50], local_private_key: &SecretKey, chaining_key: [u8; 32], hash: &mut HandshakeHash) -> Result<(PublicKey, [u8; 32], [u8; 32]), String> {
 		let version = act_bytes[0];
 		if version != 0 {
 			return Err("unexpected version".to_string());