Optimize length-prefixed writes

Author: devrandom

lightning/src/ln/msgs.rs

@@ -855,6 +855,7 @@ pub trait RoutingMessageHandler : Send + Sync + MessageSendEventsProvider {
 
 mod fuzzy_internal_msgs {
 	use ln::channelmanager::PaymentSecret;
+	use util::ser::Writeable;
 
 	// These types aren't intended to be pub, but are exposed for direct fuzzing (as we deserialize
 	// them from untrusted input):
@@ -887,6 +888,39 @@ mod fuzzy_internal_msgs {
 		// 12 bytes of 0-padding for Legacy format
 	}
 
+	pub(crate) trait OnionHopDataLength: Writeable {
+		fn encoding_length(&self) -> usize;
+	}
+
+	impl OnionHopDataLength for OnionHopData {
+		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>,
@@ -1814,7 +1848,7 @@ impl Writeable for GossipTimestampFilter {
 mod tests {
 	use hex;
 	use ln::msgs;
-	use ln::msgs::{ChannelFeatures, FinalOnionHopData, InitFeatures, NodeFeatures, OptionalField, OnionErrorPacket, OnionHopDataFormat};
+	use ln::msgs::{ChannelFeatures, FinalOnionHopData, InitFeatures, NodeFeatures, OptionalField, OnionErrorPacket, OnionHopDataFormat, OnionHopDataLength};
 	use ln::channelmanager::{PaymentPreimage, PaymentHash, PaymentSecret};
 	use util::ser::{Writeable, Readable};
 
@@ -2477,6 +2511,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 +2532,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 +2551,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 +2581,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_route_tests.rs

@@ -17,7 +17,7 @@ use ln::onion_utils;
 use routing::router::{Route, get_route};
 use ln::features::InitFeatures;
 use ln::msgs;
-use ln::msgs::{ChannelMessageHandler, HTLCFailChannelUpdate, OptionalField};
+use ln::msgs::{ChannelMessageHandler, HTLCFailChannelUpdate, OptionalField, OnionHopDataLength};
 use util::test_utils;
 use util::events::{Event, EventsProvider, MessageSendEvent, MessageSendEventsProvider};
 use util::ser::{Writeable, Writer};
@@ -237,6 +237,11 @@ impl Writeable for BogusOnionHopData {
 		writer.write_all(&self.data[..])
 	}
 }
+impl OnionHopDataLength for BogusOnionHopData {
+	fn encoding_length(&self) -> usize {
+		self.data.len()
+	}
+}
 
 #[test]
 fn test_onion_failure() {

lightning/src/ln/onion_utils.rs

@@ -13,7 +13,7 @@ use routing::router::RouteHop;
 use util::byte_utils;
 use util::chacha20::ChaCha20;
 use util::errors::{self, APIError};
-use util::ser::{Readable, Writeable, Writer};
+use util::ser::{Readable, Writeable};
 use util::logger::Logger;
 
 use bitcoin::hashes::{Hash, HashEngine};
@@ -28,6 +28,7 @@ use bitcoin::secp256k1;
 
 use std::io::Cursor;
 use std::ops::Deref;
+use ln::msgs::OnionHopDataLength;
 
 pub(super) struct OnionKeys {
 	#[cfg(test)]
@@ -183,10 +184,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;
 		}
@@ -207,7 +207,7 @@ pub(super) fn construct_onion_packet(payloads: Vec<msgs::OnionHopData>, onion_ke
 #[cfg(test)]
 // Used in testing to write bogus OnionHopDatas, which is otherwise not representable in
 // msgs::OnionHopData.
-pub(super) fn construct_onion_packet_bogus_hopdata<HD: Writeable>(payloads: Vec<HD>, onion_keys: Vec<OnionKeys>, prng_seed: [u8; 32], associated_data: &PaymentHash) -> msgs::OnionPacket {
+pub(super) fn construct_onion_packet_bogus_hopdata<HD: OnionHopDataLength>(payloads: Vec<HD>, onion_keys: Vec<OnionKeys>, prng_seed: [u8; 32], associated_data: &PaymentHash) -> msgs::OnionPacket {
 	let mut packet_data = [0; ONION_DATA_LEN];
 
 	let mut chacha = ChaCha20::new(&prng_seed, &[0; 8]);
@@ -217,7 +217,7 @@ pub(super) fn construct_onion_packet_bogus_hopdata<HD: Writeable>(payloads: Vec<
 }
 
 /// panics if route_size_insane(paylods)
-fn construct_onion_packet_with_init_noise<HD: Writeable>(mut payloads: Vec<HD>, onion_keys: Vec<OnionKeys>, mut packet_data: [u8; ONION_DATA_LEN], associated_data: &PaymentHash) -> msgs::OnionPacket {
+fn construct_onion_packet_with_init_noise<HD: OnionHopDataLength>(mut payloads: Vec<HD>, onion_keys: Vec<OnionKeys>, mut packet_data: [u8; ONION_DATA_LEN], associated_data: &PaymentHash) -> msgs::OnionPacket {
 	let filler = {
 		const ONION_HOP_DATA_LEN: usize = 65; // We may decrease this eventually after TLV is common
 		let mut res = Vec::with_capacity(ONION_HOP_DATA_LEN * (payloads.len() - 1));
@@ -232,9 +232,8 @@ fn construct_onion_packet_with_init_noise<HD: Writeable>(mut payloads: Vec<HD>,
 				chacha.process_in_place(&mut dummy); // We don't have a seek function :(
 			}
 
-			let mut payload_len = Writer::new();
-			payload.write(&mut payload_len).expect("Failed to calculate length");
-			pos += payload_len.len() + 32;
+			let payload_len = payload.encoding_length();
+			pos += payload_len + 32;
 			assert!(pos <= ONION_DATA_LEN);
 
 			res.resize(pos, 0u8);
@@ -245,9 +244,7 @@ fn construct_onion_packet_with_init_noise<HD: Writeable>(mut payloads: Vec<HD>,
 
 	let mut hmac_res = [0; 32];
 	for (i, (payload, keys)) in payloads.iter_mut().zip(onion_keys.iter()).rev().enumerate() {
-		let mut payload_len = Writer::new();
-		payload.write(&mut payload_len).expect("Failed to calculate length");
-		let len = payload_len.len();
+		let len = payload.encoding_length();
 		shift_arr_right(&mut packet_data, len + 32);
 		packet_data[0..len].copy_from_slice(&payload.encode()[..]);
 		packet_data[len..(len + 32)].copy_from_slice(&hmac_res);
@@ -490,6 +487,7 @@ mod tests {
 	use bitcoin::secp256k1::key::{PublicKey,SecretKey};
 
 	use super::OnionKeys;
+	use ln::msgs::OnionHopDataLength;
 
 	fn build_test_onion_keys() -> Vec<OnionKeys> {
 		// Keys from BOLT 4, used in both test vector tests
@@ -649,6 +647,11 @@ mod tests {
 			writer.write_all(&self.data[..])
 		}
 	}
+	impl OnionHopDataLength for RawOnionHopData {
+		fn encoding_length(&self) -> usize {
+			self.data.len()
+		}
+	}
 
 	#[test]
 	fn variable_length_onion_vectors() {

lightning/src/util/ser.rs

@@ -10,7 +10,7 @@
 //! A very simple serialization framework which is used to serialize/deserialize messages as well
 //! as ChannelsManagers and ChannelMonitors.
 
-use std::io::{Read, Write};
+use std::io::{Read, Write, Cursor};
 use std::collections::HashMap;
 use std::hash::Hash;
 use std::sync::Mutex;
@@ -62,6 +62,32 @@ impl Writer {
 	pub fn len(&self) -> usize {
 		self.0.len()
 	}
+
+	// Write the data with a varint length prefix.
+	// Use length_adjust to add to the length for special cases, such as writing an Option.
+	pub(crate) fn write_all_with_length<T: Writeable>(&mut self, data: &T, length_adjust: usize) -> Result<(), ::std::io::Error> {
+		// Most common case is a single byte length prefix, so leave a space for that
+		self.0.push(0);
+		let start = self.len();
+		data.write(self).expect("No in-memory data may fail to serialize");
+		self.insert_length(start, length_adjust)
+	}
+
+	fn insert_length(&mut self, slice_start: usize, length_adjust: usize) -> Result<(), ::std::io::Error> {
+		let slice_end = self.len();
+		let length = slice_end - slice_start + length_adjust;
+		let big_size = BigSize(length as u64);
+		let len_len = big_size.expected_length();
+		// At this point we have a placeholder for a varint of length 1 just before the start of data
+		// Shift the slice to make space for the length.  If len_len is 1, no shift is needed.
+		for _ in 1..len_len {
+			self.0.insert(slice_start, 0);
+		}
+		let mut cursor = Cursor::new(&mut self.0);
+		cursor.set_position((slice_start - 1) as u64);
+		// Write the length in the space we created before the start of data
+		BigSize(length as u64).write_write(&mut cursor)
+	}
 }
 
 impl Write for Writer {
@@ -225,27 +251,42 @@ impl Readable for U48 {
 /// you're looking for an example of a variable-length integer to use for your own project, move
 /// along, this is a rather poor design.
 pub(crate) struct BigSize(pub u64);
-impl Writeable for BigSize {
-	#[inline]
-	fn write(&self, writer: &mut Writer) -> Result<(), ::std::io::Error> {
+impl BigSize {
+	pub fn write_write<W: Write>(&self, write: &mut W) -> Result<(), ::std::io::Error> {
 		match self.0 {
 			0...0xFC => {
-				(self.0 as u8).write(writer)
+				write.write_all(&[self.0 as u8])
 			},
 			0xFD...0xFFFF => {
-				0xFDu8.write(writer)?;
-				(self.0 as u16).write(writer)
+				write.write_all(&[0xFDu8])?;
+				write.write_all(&be16_to_array(self.0 as u16))
 			},
 			0x10000...0xFFFFFFFF => {
-				0xFEu8.write(writer)?;
-				(self.0 as u32).write(writer)
+				write.write_all(&[0xFEu8])?;
+				write.write_all(&be32_to_array(self.0 as u32))
 			},
 			_ => {
-				0xFFu8.write(writer)?;
-				(self.0 as u64).write(writer)
+				write.write_all(&[0xFFu8])?;
+				write.write_all(&be64_to_array(self.0 as u64))
 			},
 		}
 	}
+
+	pub fn expected_length(&self) -> usize {
+		match self.0 {
+			0...0xFC => 1,
+			0xFD...0xFFFF => 3,
+			0x10000...0xFFFFFFFF => 5,
+			_ => 9,
+		}
+	}
+}
+
+impl Writeable for BigSize {
+	#[inline]
+	fn write(&self, writer: &mut Writer) -> Result<(), ::std::io::Error> {
+		self.write_write(writer)
+	}
 }
 impl Readable for BigSize {
 	#[inline]
@@ -609,10 +650,7 @@ impl<T: Writeable> Writeable for Option<T> {
 		match *self {
 			None => 0u8.write(w)?,
 			Some(ref data) => {
-				let mut len_calc = Writer::new();
-				data.write(&mut len_calc).expect("No in-memory data may fail to serialize");
-				BigSize(len_calc.len() as u64 + 1).write(w)?;
-				data.write(w)?;
+				w.write_all_with_length(data, 1)?;
 			}
 		}
 		Ok(())
@@ -732,3 +770,56 @@ impl<A: Writeable, B: Writeable> Writeable for (A, B) {
 		self.1.write(w)
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use std::io::Cursor;
+	use util::ser::*;
+
+	#[test]
+	fn big_size() {
+		for i in vec![0, 0xfc, 0xfd, 0xffff, 0x10000, 0xFFFFFFFF, 0xFFFFFFFFF] {
+			let mut writer = Writer::new();
+			let size = BigSize(i);
+			size.write(&mut writer).unwrap();
+			assert_eq!(writer.len(), size.expected_length());
+			let read_size = BigSize::read(&mut Cursor::new(writer.0)).unwrap();
+			assert_eq!(read_size.0, size.0);
+		}
+	}
+
+	#[test]
+	fn length_prefix_zero() {
+		impl_array!(0);
+		let mut writer = Writer::new();
+		writer.write_all_with_length(&[], 0).unwrap();
+		// Writing a zero-length length-prefixed item should result in just one zero byte
+		assert_eq!(writer.0, vec![0]);
+	}
+
+	#[test]
+	fn length_prefix_one() {
+		impl_array!(1);
+		let mut writer = Writer::new();
+		writer.write_all_with_length(&[0x11u8], 0).unwrap();
+		assert_eq!(writer.0, vec![1, 0x11]);
+	}
+
+	#[test]
+	fn length_prefix() {
+		for size in vec![3, 0xfc, 0xfd, 0xffff, 0x10000] {
+			let mut writer = Writer::new();
+			let mut data = Vec::new();
+			// A Vec serializes its own 2 bytes length, so adjust this
+			data.resize(size-3, 0x11u8);
+			// Writing an Option is convenient for round-trip checking of length-prefixed writes
+			Some(data.clone()).write(&mut writer).unwrap();
+			let mut reader = Cursor::new(&writer.0);
+			let read_data: Option<Vec<u8>> = Readable::read(&mut reader).unwrap();
+            assert_eq!(data, read_data.unwrap());
+			// Ensure everything was consumed
+            assert_eq!(reader.position() as usize, writer.len());
+			println!("{}", writer.len())
+		}
+	}
+}

lightning/src/util/ser_macros.rs

@@ -9,13 +9,10 @@
 
 macro_rules! encode_tlv {
 	($stream: expr, {$(($type: expr, $field: expr)),*}) => { {
-		use util::ser::{BigSize, Writer};
+		use util::ser::BigSize;
 		$(
 			BigSize($type).write($stream)?;
-			let mut len_calc = Writer::new();
-			$field.write(&mut len_calc)?;
-			BigSize(len_calc.len() as u64).write($stream)?;
-			$field.write($stream)?;
+			$stream.write_all_with_length(&$field, 0)?;
 		)*
 	} }
 }