Merge pull request lightningdevkit#3494 from jkczyz/2024-12-invoice-b…

…yte-allocation Don't over-allocate invoice bytes
TheBlueMatt · Feb 7, 2025 · f045c0e · f045c0e
2 parents 8e90841 + 0912b51
commit f045c0e
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 151 deletions.
diff --git a/lightning/src/offers/invoice.rs b/lightning/src/offers/invoice.rs
@@ -122,7 +122,7 @@ use crate::offers::invoice_macros::{invoice_accessors_common, invoice_builder_me
 #[cfg(test)]
 use crate::offers::invoice_macros::invoice_builder_methods_test_common;
 use crate::offers::invoice_request::{EXPERIMENTAL_INVOICE_REQUEST_TYPES, ExperimentalInvoiceRequestTlvStream, ExperimentalInvoiceRequestTlvStreamRef, INVOICE_REQUEST_PAYER_ID_TYPE, INVOICE_REQUEST_TYPES, IV_BYTES as INVOICE_REQUEST_IV_BYTES, InvoiceRequest, InvoiceRequestContents, InvoiceRequestTlvStream, InvoiceRequestTlvStreamRef};
-use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE};
+use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self};
 use crate::offers::nonce::Nonce;
 use crate::offers::offer::{Amount, EXPERIMENTAL_OFFER_TYPES, ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OFFER_TYPES, OfferTlvStream, OfferTlvStreamRef, Quantity};
 use crate::offers::parse::{Bolt12ParseError, Bolt12SemanticError, ParsedMessage};
@@ -520,19 +520,8 @@ impl UnsignedBolt12Invoice {
 		let (_, _, _, invoice_tlv_stream, _, _, experimental_invoice_tlv_stream) =
 			contents.as_tlv_stream();
 
-		// Allocate enough space for the invoice, which will include:
-		// - all TLV records from `invreq_bytes` except signatures,
-		// - all invoice-specific TLV records, and
-		// - a signature TLV record once the invoice is signed.
-		//
-		// This assumes both the invoice request and the invoice will each only have one signature
-		// using SIGNATURE_TYPES.start as the TLV record. Thus, it is accounted for by invreq_bytes.
-		let mut bytes = Vec::with_capacity(
-			invreq_bytes.len()
-				+ invoice_tlv_stream.serialized_length()
-				+ if contents.is_for_offer() { 0 } else { SIGNATURE_TLV_RECORD_SIZE }
-				+ experimental_invoice_tlv_stream.serialized_length(),
-		);
+		const INVOICE_ALLOCATION_SIZE: usize = 1024;
+		let mut bytes = Vec::with_capacity(INVOICE_ALLOCATION_SIZE);
 
 		// Use the invoice_request bytes instead of the invoice_request TLV stream as the latter may
 		// have contained unknown TLV records, which are not stored in `InvoiceRequestContents` or
@@ -545,23 +534,16 @@ impl UnsignedBolt12Invoice {
 
 		invoice_tlv_stream.write(&mut bytes).unwrap();
 
-		let mut experimental_tlv_stream = TlvStream::new(remaining_bytes)
-			.range(EXPERIMENTAL_TYPES)
-			.peekable();
-		let mut experimental_bytes = Vec::with_capacity(
-			remaining_bytes.len()
-				- experimental_tlv_stream
-					.peek()
-					.map_or(remaining_bytes.len(), |first_record| first_record.start)
-				+ experimental_invoice_tlv_stream.serialized_length(),
-		);
+		const EXPERIMENTAL_TLV_ALLOCATION_SIZE: usize = 0;
+		let mut experimental_bytes = Vec::with_capacity(EXPERIMENTAL_TLV_ALLOCATION_SIZE);
 
+		let experimental_tlv_stream = TlvStream::new(remaining_bytes)
+			.range(EXPERIMENTAL_TYPES);
 		for record in experimental_tlv_stream {
 			record.write(&mut experimental_bytes).unwrap();
 		}
 
 		experimental_invoice_tlv_stream.write(&mut experimental_bytes).unwrap();
-		debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());
 
 		let tlv_stream = TlvStream::new(&bytes).chain(TlvStream::new(&experimental_bytes));
 		let tagged_hash = TaggedHash::from_tlv_stream(SIGNATURE_TAG, tlv_stream);
@@ -592,14 +574,6 @@ macro_rules! unsigned_invoice_sign_method { ($self: ident, $self_type: ty $(, $s
 		signature_tlv_stream.write(&mut $self.bytes).unwrap();
 
 		// Append the experimental bytes after the signature.
-		debug_assert_eq!(
-			// The two-byte overallocation results from SIGNATURE_TLV_RECORD_SIZE accommodating TLV
-			// records with types >= 253.
-			$self.bytes.len()
-				+ $self.experimental_bytes.len()
-				+ if $self.contents.is_for_offer() { 0 } else { 2 },
-			$self.bytes.capacity(),
-		);
 		$self.bytes.extend_from_slice(&$self.experimental_bytes);
 
 		Ok(Bolt12Invoice {
@@ -965,13 +939,6 @@ impl Hash for Bolt12Invoice {
 }
 
 impl InvoiceContents {
-	fn is_for_offer(&self) -> bool {
-		match self {
-			InvoiceContents::ForOffer { .. } => true,
-			InvoiceContents::ForRefund { .. } => false,
-		}
-	}
-
 	/// Whether the original offer or refund has expired.
 	#[cfg(feature = "std")]
 	fn is_offer_or_refund_expired(&self) -> bool {
@@ -1362,7 +1329,11 @@ pub(super) const EXPERIMENTAL_INVOICE_TYPES: core::ops::RangeFrom<u64> = 3_000_0
 
 #[cfg(not(test))]
 tlv_stream!(
-	ExperimentalInvoiceTlvStream, ExperimentalInvoiceTlvStreamRef, EXPERIMENTAL_INVOICE_TYPES, {}
+	ExperimentalInvoiceTlvStream, ExperimentalInvoiceTlvStreamRef, EXPERIMENTAL_INVOICE_TYPES, {
+		// When adding experimental TLVs, update EXPERIMENTAL_TLV_ALLOCATION_SIZE accordingly in
+		// both UnsignedBolt12Invoice:new and UnsignedStaticInvoice::new to avoid unnecessary
+		// allocations.
+	}
 );
 
 #[cfg(test)]
@@ -2880,9 +2851,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice.bytes.reserve_exact(
-			unsigned_invoice.bytes.capacity() - unsigned_invoice.bytes.len() + unknown_bytes.len(),
-		);
 		unsigned_invoice.bytes.extend_from_slice(&unknown_bytes);
 		unsigned_invoice.tagged_hash =
 			TaggedHash::from_valid_tlv_stream_bytes(SIGNATURE_TAG, &unsigned_invoice.bytes);
@@ -2917,9 +2885,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice.bytes.reserve_exact(
-			unsigned_invoice.bytes.capacity() - unsigned_invoice.bytes.len() + unknown_bytes.len(),
-		);
 		unsigned_invoice.bytes.extend_from_slice(&unknown_bytes);
 		unsigned_invoice.tagged_hash =
 			TaggedHash::from_valid_tlv_stream_bytes(SIGNATURE_TAG, &unsigned_invoice.bytes);
@@ -2982,9 +2947,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice.bytes.reserve_exact(
-			unsigned_invoice.bytes.capacity() - unsigned_invoice.bytes.len() + unknown_bytes.len(),
-		);
 		unsigned_invoice.experimental_bytes.extend_from_slice(&unknown_bytes);
 
 		let tlv_stream = TlvStream::new(&unsigned_invoice.bytes)
@@ -3021,9 +2983,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice.bytes.reserve_exact(
-			unsigned_invoice.bytes.capacity() - unsigned_invoice.bytes.len() + unknown_bytes.len(),
-		);
 		unsigned_invoice.experimental_bytes.extend_from_slice(&unknown_bytes);
 
 		let tlv_stream = TlvStream::new(&unsigned_invoice.bytes)

diff --git a/lightning/src/offers/invoice_request.rs b/lightning/src/offers/invoice_request.rs
@@ -77,7 +77,7 @@ use crate::ln::channelmanager::PaymentId;
 use crate::types::features::InvoiceRequestFeatures;
 use crate::ln::inbound_payment::{ExpandedKey, IV_LEN};
 use crate::ln::msgs::DecodeError;
-use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE};
+use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self};
 use crate::offers::nonce::Nonce;
 use crate::offers::offer::{Amount, EXPERIMENTAL_OFFER_TYPES, ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OFFER_TYPES, Offer, OfferContents, OfferId, OfferTlvStream, OfferTlvStreamRef};
 use crate::offers::parse::{Bolt12ParseError, ParsedMessage, Bolt12SemanticError};
@@ -473,17 +473,8 @@ impl UnsignedInvoiceRequest {
 			_experimental_offer_tlv_stream, experimental_invoice_request_tlv_stream,
 		) = contents.as_tlv_stream();
 
-		// Allocate enough space for the invoice_request, which will include:
-		// - all TLV records from `offer.bytes`,
-		// - all invoice_request-specific TLV records, and
-		// - a signature TLV record once the invoice_request is signed.
-		let mut bytes = Vec::with_capacity(
-			offer.bytes.len()
-				+ payer_tlv_stream.serialized_length()
-				+ invoice_request_tlv_stream.serialized_length()
-				+ SIGNATURE_TLV_RECORD_SIZE
-				+ experimental_invoice_request_tlv_stream.serialized_length(),
-		);
+		const INVOICE_REQUEST_ALLOCATION_SIZE: usize = 512;
+		let mut bytes = Vec::with_capacity(INVOICE_REQUEST_ALLOCATION_SIZE);
 
 		payer_tlv_stream.write(&mut bytes).unwrap();
 
@@ -495,23 +486,16 @@ impl UnsignedInvoiceRequest {
 
 		invoice_request_tlv_stream.write(&mut bytes).unwrap();
 
-		let mut experimental_tlv_stream = TlvStream::new(remaining_bytes)
-			.range(EXPERIMENTAL_OFFER_TYPES)
-			.peekable();
-		let mut experimental_bytes = Vec::with_capacity(
-			remaining_bytes.len()
-				- experimental_tlv_stream
-					.peek()
-					.map_or(remaining_bytes.len(), |first_record| first_record.start)
-				+ experimental_invoice_request_tlv_stream.serialized_length(),
-		);
+		const EXPERIMENTAL_TLV_ALLOCATION_SIZE: usize = 0;
+		let mut experimental_bytes = Vec::with_capacity(EXPERIMENTAL_TLV_ALLOCATION_SIZE);
 
+		let experimental_tlv_stream = TlvStream::new(remaining_bytes)
+			.range(EXPERIMENTAL_OFFER_TYPES);
 		for record in experimental_tlv_stream {
 			record.write(&mut experimental_bytes).unwrap();
 		}
 
 		experimental_invoice_request_tlv_stream.write(&mut experimental_bytes).unwrap();
-		debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());
 
 		let tlv_stream = TlvStream::new(&bytes).chain(TlvStream::new(&experimental_bytes));
 		let tagged_hash = TaggedHash::from_tlv_stream(SIGNATURE_TAG, tlv_stream);
@@ -544,12 +528,6 @@ macro_rules! unsigned_invoice_request_sign_method { (
 		signature_tlv_stream.write(&mut $self.bytes).unwrap();
 
 		// Append the experimental bytes after the signature.
-		debug_assert_eq!(
-			// The two-byte overallocation results from SIGNATURE_TLV_RECORD_SIZE accommodating TLV
-			// records with types >= 253.
-			$self.bytes.len() + $self.experimental_bytes.len() + 2,
-			$self.bytes.capacity(),
-		);
 		$self.bytes.extend_from_slice(&$self.experimental_bytes);
 
 		Ok(InvoiceRequest {
@@ -1127,7 +1105,10 @@ pub(super) const EXPERIMENTAL_INVOICE_REQUEST_TYPES: core::ops::Range<u64> =
 #[cfg(not(test))]
 tlv_stream!(
 	ExperimentalInvoiceRequestTlvStream, ExperimentalInvoiceRequestTlvStreamRef,
-	EXPERIMENTAL_INVOICE_REQUEST_TYPES, {}
+	EXPERIMENTAL_INVOICE_REQUEST_TYPES, {
+		// When adding experimental TLVs, update EXPERIMENTAL_TLV_ALLOCATION_SIZE accordingly in
+		// UnsignedInvoiceRequest::new to avoid unnecessary allocations.
+	}
 );
 
 #[cfg(test)]
@@ -2422,11 +2403,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice_request.bytes.reserve_exact(
-			unsigned_invoice_request.bytes.capacity()
-				- unsigned_invoice_request.bytes.len()
-				+ unknown_bytes.len(),
-		);
 		unsigned_invoice_request.bytes.extend_from_slice(&unknown_bytes);
 		unsigned_invoice_request.tagged_hash =
 			TaggedHash::from_valid_tlv_stream_bytes(SIGNATURE_TAG, &unsigned_invoice_request.bytes);
@@ -2460,11 +2436,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice_request.bytes.reserve_exact(
-			unsigned_invoice_request.bytes.capacity()
-				- unsigned_invoice_request.bytes.len()
-				+ unknown_bytes.len(),
-		);
 		unsigned_invoice_request.bytes.extend_from_slice(&unknown_bytes);
 		unsigned_invoice_request.tagged_hash =
 			TaggedHash::from_valid_tlv_stream_bytes(SIGNATURE_TAG, &unsigned_invoice_request.bytes);
@@ -2508,11 +2479,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice_request.bytes.reserve_exact(
-			unsigned_invoice_request.bytes.capacity()
-				- unsigned_invoice_request.bytes.len()
-				+ unknown_bytes.len(),
-		);
 		unsigned_invoice_request.experimental_bytes.extend_from_slice(&unknown_bytes);
 
 		let tlv_stream = TlvStream::new(&unsigned_invoice_request.bytes)
@@ -2549,11 +2515,6 @@ mod tests {
 		BigSize(32).write(&mut unknown_bytes).unwrap();
 		[42u8; 32].write(&mut unknown_bytes).unwrap();
 
-		unsigned_invoice_request.bytes.reserve_exact(
-			unsigned_invoice_request.bytes.capacity()
-				- unsigned_invoice_request.bytes.len()
-				+ unknown_bytes.len(),
-		);
 		unsigned_invoice_request.experimental_bytes.extend_from_slice(&unknown_bytes);
 
 		let tlv_stream = TlvStream::new(&unsigned_invoice_request.bytes)

diff --git a/lightning/src/offers/merkle.rs b/lightning/src/offers/merkle.rs
@@ -11,7 +11,6 @@
 
 use bitcoin::hashes::{Hash, HashEngine, sha256};
 use bitcoin::secp256k1::{Message, PublicKey, Secp256k1, self};
-use bitcoin::secp256k1::constants::SCHNORR_SIGNATURE_SIZE;
 use bitcoin::secp256k1::schnorr::Signature;
 use crate::io;
 use crate::util::ser::{BigSize, Readable, Writeable, Writer};
@@ -26,10 +25,6 @@ tlv_stream!(SignatureTlvStream, SignatureTlvStreamRef<'a>, SIGNATURE_TYPES, {
 	(240, signature: Signature),
 });
 
-/// Size of a TLV record in `SIGNATURE_TYPES` when the type is 1000. TLV types are encoded using
-/// BigSize, so a TLV record with type 240 will use two less bytes.
-pub(super) const SIGNATURE_TLV_RECORD_SIZE: usize = 3 + 1 + SCHNORR_SIGNATURE_SIZE;
-
 /// A hash for use in a specific context by tweaking with a context-dependent tag as per [BIP 340]
 /// and computed over the merkle root of a TLV stream to sign as defined in [BOLT 12].
 ///
@@ -253,7 +248,6 @@ pub(super) struct TlvRecord<'a> {
 	type_bytes: &'a [u8],
 	// The entire TLV record.
 	pub(super) record_bytes: &'a [u8],
-	pub(super) start: usize,
 	pub(super) end: usize,
 }
 
@@ -278,7 +272,7 @@ impl<'a> Iterator for TlvStream<'a> {
 			self.data.set_position(end);
 
 			Some(TlvRecord {
-				r#type, type_bytes, record_bytes, start: start as usize, end: end as usize,
+				r#type, type_bytes, record_bytes, end: end as usize,
 			})
 		} else {
 			None

diff --git a/lightning/src/offers/offer.rs b/lightning/src/offers/offer.rs
@@ -438,7 +438,8 @@ macro_rules! offer_builder_methods { (
 			}
 		}
 
-		let mut bytes = Vec::new();
+		const OFFER_ALLOCATION_SIZE: usize = 512;
+		let mut bytes = Vec::with_capacity(OFFER_ALLOCATION_SIZE);
 		$self.offer.write(&mut bytes).unwrap();
 
 		let id = OfferId::from_valid_offer_tlv_stream(&bytes);

diff --git a/lightning/src/offers/refund.rs b/lightning/src/offers/refund.rs
@@ -338,7 +338,8 @@ macro_rules! refund_builder_methods { (
 			$self.refund.payer.0 = metadata;
 		}
 
-		let mut bytes = Vec::new();
+		const REFUND_ALLOCATION_SIZE: usize = 512;
+		let mut bytes = Vec::with_capacity(REFUND_ALLOCATION_SIZE);
 		$self.refund.write(&mut bytes).unwrap();
 
 		Ok(Refund {