Optimize checking of onion route size

Author: devrandom

lightning/src/ln/msgs.rs

@@ -887,6 +887,34 @@ mod fuzzy_internal_msgs {
 		// 12 bytes of 0-padding for Legacy format
 	}
 
+	impl OnionHopData {
+		pub fn encoding_length(&self) -> usize {
+			fn high_zero_bytes_dropped_length(data: u64) -> usize {
+				8 - ((data.leading_zeros()/8) as usize)
+			}
+			match self.format {
+				OnionHopDataFormat::Legacy { .. } => 1 + 8 + 8 + 4 + 12, // 33
+				OnionHopDataFormat::NonFinalNode { .. } => {
+					1 +
+						1 + 1 + high_zero_bytes_dropped_length(self.amt_to_forward) +
+						1 + 1 + high_zero_bytes_dropped_length(self.outgoing_cltv_value as u64) +
+						1 + 1 + 8 // short_channel_id
+				}
+				OnionHopDataFormat::FinalNode { ref payment_data } => {
+					1 +
+						1 + 1 + high_zero_bytes_dropped_length(self.amt_to_forward) +
+						1 + 1 + high_zero_bytes_dropped_length(self.outgoing_cltv_value as u64) +
+						match payment_data {
+							None => 0,
+							Some(payment_data) =>
+								1 + 1 + high_zero_bytes_dropped_length(payment_data.total_msat) +
+									32 // payment_secret
+						}
+				}
+			}
+		}
+	}
+
 	pub struct DecodedOnionErrorPacket {
 		pub(crate) hmac: [u8; 32],
 		pub(crate) failuremsg: Vec<u8>,
@@ -2477,6 +2505,7 @@ mod tests {
 		let encoded_value = msg.encode();
 		let target_value = hex::decode("00deadbeef1bad1dea0badf00d01020304ffffffff000000000000000000000000").unwrap();
 		assert_eq!(encoded_value, target_value);
+		assert_eq!(encoded_value.len(), msg.encoding_length())
 	}
 
 	#[test]
@@ -2497,6 +2526,7 @@ mod tests {
 		} else { panic!(); }
 		assert_eq!(msg.amt_to_forward, 0x0badf00d01020304);
 		assert_eq!(msg.outgoing_cltv_value, 0xffffffff);
+		assert_eq!(encoded_value.len(), msg.encoding_length())
 	}
 
 	#[test]
@@ -2515,6 +2545,7 @@ mod tests {
 		if let OnionHopDataFormat::FinalNode { payment_data: None } = msg.format { } else { panic!(); }
 		assert_eq!(msg.amt_to_forward, 0x0badf00d01020304);
 		assert_eq!(msg.outgoing_cltv_value, 0xffffffff);
+		assert_eq!(encoded_value.len(), msg.encoding_length())
 	}
 
 	#[test]
@@ -2544,6 +2575,7 @@ mod tests {
 		} else { panic!(); }
 		assert_eq!(msg.amt_to_forward, 0x0badf00d01020304);
 		assert_eq!(msg.outgoing_cltv_value, 0xffffffff);
+		assert_eq!(encoded_value.len(), msg.encoding_length())
 	}
 
 	#[test]

lightning/src/ln/onion_utils.rs

@@ -183,10 +183,9 @@ fn shift_arr_right(arr: &mut [u8; ONION_DATA_LEN], amt: usize) {
 pub(super) fn route_size_insane(payloads: &Vec<msgs::OnionHopData>) -> bool {
 	let mut len = 0;
 	for payload in payloads.iter() {
-		let mut payload_len = Writer::new();
-		payload.write(&mut payload_len).expect("Failed to calculate length");
-		assert!(payload_len.len() + 32 < ONION_DATA_LEN);
-		len += payload_len.len() + 32;
+		let encoding_length = payload.encoding_length();
+		assert!(encoding_length + 32 < ONION_DATA_LEN);
+		len += encoding_length + 32;
 		if len > ONION_DATA_LEN {
 			return true;
 		}