DRY shared hkdf_extract_expand code to new module

Author: valentinewallace

lightning/src/ln/channelmanager.rs

@@ -81,6 +81,7 @@ mod inbound_payment {
 	use ln::msgs;
 	use ln::msgs::MAX_VALUE_MSAT;
 	use util::chacha20::ChaCha20;
+	use util::crypto::hkdf_extract_expand;
 	use util::logger::Logger;
 
 	use core::convert::TryInto;
@@ -112,7 +113,12 @@ mod inbound_payment {
 
 	impl ExpandedKey {
 		pub(super) fn new(key_material: &KeyMaterial) -> ExpandedKey {
-			hkdf_extract_expand(b"LDK Inbound Payment Key Expansion", &key_material)
+			let keys = hkdf_extract_expand(b"LDK Inbound Payment Key Expansion", &key_material.0, 3);
+			Self {
+				metadata_key: keys[0],
+				ldk_pmt_hash_key: keys[1],
+				user_pmt_hash_key: keys[2],
+			}
 		}
 	}
 
@@ -330,31 +336,6 @@ mod inbound_payment {
 		}
 		return Ok(PaymentPreimage(decoded_payment_preimage))
 	}
-
-	fn hkdf_extract_expand(salt: &[u8], ikm: &KeyMaterial) -> ExpandedKey {
-		let mut hmac = HmacEngine::<Sha256>::new(salt);
-		hmac.input(&ikm.0);
-		let prk = Hmac::from_engine(hmac).into_inner();
-		let mut hmac = HmacEngine::<Sha256>::new(&prk[..]);
-		hmac.input(&[1; 1]);
-		let metadata_key = Hmac::from_engine(hmac).into_inner();
-
-		let mut hmac = HmacEngine::<Sha256>::new(&prk[..]);
-		hmac.input(&metadata_key);
-		hmac.input(&[2; 1]);
-		let ldk_pmt_hash_key = Hmac::from_engine(hmac).into_inner();
-
-		let mut hmac = HmacEngine::<Sha256>::new(&prk[..]);
-		hmac.input(&ldk_pmt_hash_key);
-		hmac.input(&[3; 1]);
-		let user_pmt_hash_key = Hmac::from_engine(hmac).into_inner();
-
-		ExpandedKey {
-			metadata_key,
-			ldk_pmt_hash_key,
-			user_pmt_hash_key,
-		}
-	}
 }
 
 // We hold various information about HTLC relay in the HTLC objects in Channel itself:

lightning/src/ln/peer_channel_encryptor.rs

@@ -12,7 +12,7 @@ use prelude::*;
 use ln::msgs::LightningError;
 use ln::msgs;
 
-use bitcoin::hashes::{Hash, HashEngine, Hmac, HmacEngine};
+use bitcoin::hashes::{Hash, HashEngine};
 use bitcoin::hashes::sha256::Hash as Sha256;
 
 use bitcoin::secp256k1::Secp256k1;
@@ -21,6 +21,7 @@ use bitcoin::secp256k1::ecdh::SharedSecret;
 use bitcoin::secp256k1;
 
 use util::chacha20poly1305rfc::ChaCha20Poly1305RFC;
+use util::crypto::hkdf_extract_expand;
 use bitcoin::hashes::hex::ToHex;
 
 /// Maximum Lightning message data length according to
@@ -160,24 +161,11 @@ impl PeerChannelEncryptor {
 		Ok(())
 	}
 
-	fn hkdf_extract_expand(salt: &[u8], ikm: &[u8]) -> ([u8; 32], [u8; 32]) {
-		let mut hmac = HmacEngine::<Sha256>::new(salt);
-		hmac.input(ikm);
-		let prk = Hmac::from_engine(hmac).into_inner();
-		let mut hmac = HmacEngine::<Sha256>::new(&prk[..]);
-		hmac.input(&[1; 1]);
-		let t1 = Hmac::from_engine(hmac).into_inner();
-		let mut hmac = HmacEngine::<Sha256>::new(&prk[..]);
-		hmac.input(&t1);
-		hmac.input(&[2; 1]);
-		(t1, Hmac::from_engine(hmac).into_inner())
-	}
-
 	#[inline]
 	fn hkdf(state: &mut BidirectionalNoiseState, ss: SharedSecret) -> [u8; 32] {
-		let (t1, t2) = Self::hkdf_extract_expand(&state.ck, &ss[..]);
-		state.ck = t1;
-		t2
+		let hkdf = hkdf_extract_expand(&state.ck, &ss[..], 2);
+		state.ck = hkdf[0];
+		hkdf[1]
 	}
 
 	#[inline]
@@ -311,7 +299,7 @@ impl PeerChannelEncryptor {
 						let temp_k = PeerChannelEncryptor::hkdf(bidirectional_state, ss);
 
 						PeerChannelEncryptor::encrypt_with_ad(&mut res[50..], 0, &temp_k, &bidirectional_state.h, &[0; 0]);
-						final_hkdf = Self::hkdf_extract_expand(&bidirectional_state.ck, &[0; 0]);
+						final_hkdf = hkdf_extract_expand(&bidirectional_state.ck, &[0; 0], 2);
 						ck = bidirectional_state.ck.clone();
 						res
 					},
@@ -320,7 +308,7 @@ impl PeerChannelEncryptor {
 			_ => panic!("Cannot get act one after noise handshake completes"),
 		};
 
-		let (sk, rk) = final_hkdf;
+		let (sk, rk) = (final_hkdf[0], final_hkdf[1]);
 		self.noise_state = NoiseState::Finished {
 			sk,
 			sn: 0,
@@ -365,15 +353,15 @@ impl PeerChannelEncryptor {
 						let temp_k = PeerChannelEncryptor::hkdf(bidirectional_state, ss);
 
 						PeerChannelEncryptor::decrypt_with_ad(&mut [0; 0], 0, &temp_k, &bidirectional_state.h, &act_three[50..])?;
-						final_hkdf = Self::hkdf_extract_expand(&bidirectional_state.ck, &[0; 0]);
+						final_hkdf = hkdf_extract_expand(&bidirectional_state.ck, &[0; 0], 2);
 						ck = bidirectional_state.ck.clone();
 					},
 					_ => panic!("Wrong direction for act"),
 				},
 			_ => panic!("Cannot get act one after noise handshake completes"),
 		}
 
-		let (rk, sk) = final_hkdf;
+		let (rk, sk) = (final_hkdf[0], final_hkdf[1]);
 		self.noise_state = NoiseState::Finished {
 			sk,
 			sn: 0,
@@ -399,9 +387,9 @@ impl PeerChannelEncryptor {
 		match self.noise_state {
 			NoiseState::Finished { ref mut sk, ref mut sn, ref mut sck, rk: _, rn: _, rck: _ } => {
 				if *sn >= 1000 {
-					let (new_sck, new_sk) = Self::hkdf_extract_expand(sck, sk);
-					*sck = new_sck;
-					*sk = new_sk;
+					let hkdf = hkdf_extract_expand(sck, sk, 2);
+					*sck = hkdf[0];
+					*sk = hkdf[1];
 					*sn = 0;
 				}
 
@@ -425,9 +413,9 @@ impl PeerChannelEncryptor {
 		match self.noise_state {
 			NoiseState::Finished { sk: _, sn: _, sck: _, ref mut rk, ref mut rn, ref mut rck } => {
 				if *rn >= 1000 {
-					let (new_rck, new_rk) = Self::hkdf_extract_expand(rck, rk);
-					*rck = new_rck;
-					*rk = new_rk;
+					let hkdf = hkdf_extract_expand(rck, rk, 2);
+					*rck = hkdf[0];
+					*rk = hkdf[1];
 					*rn = 0;
 				}
 

lightning/src/util/mod.rs

@@ -48,3 +48,25 @@ pub mod test_utils;
 /// machine errors and used in fuzz targets and tests.
 #[cfg(any(test, feature = "fuzztarget", feature = "_test_utils"))]
 pub mod enforcing_trait_impls;
+
+pub(crate) mod crypto {
+	use bitcoin::hashes::{Hash, HashEngine};
+	use bitcoin::hashes::hmac::{Hmac, HmacEngine};
+	use bitcoin::hashes::sha256::Hash as Sha256;
+
+	pub fn hkdf_extract_expand(salt: &[u8], ikm: &[u8], num_keys: u8) -> Vec<[u8; 32]> {
+		let mut keys_res: Vec<[u8; 32]> = Vec::new();
+		let mut hmac = HmacEngine::<Sha256>::new(salt);
+		hmac.input(ikm);
+		let prk = Hmac::from_engine(hmac).into_inner();
+		for i in 0..num_keys {
+			let mut hmac = HmacEngine::<Sha256>::new(&prk[..]);
+			if i != 0 {
+				hmac.input(&keys_res[keys_res.len() - 1]);
+			}
+			hmac.input(&[i + 1; 1]);
+			keys_res.push(Hmac::from_engine(hmac).into_inner());
+		}
+		keys_res
+	}
+}