Address custom HTLC TLV fixups

[rust-lightning] / lightning / src / offers / merkle.rs

diff --git a/lightning/src/offers/merkle.rs b/lightning/src/offers/merkle.rs
index 9782dc7d1e84131f1d6e659dce2d28ac7e9de6f0..f7c33902c51441cd3e2a358637c9ddb4d2cc31e6 100644 (file)
--- a/lightning/src/offers/merkle.rs
+++ b/lightning/src/offers/merkle.rs
@@ -66,7 +66,7 @@ pub(super) fn verify_signature(
        secp_ctx.verify_schnorr(signature, &digest, &pubkey)
 }
 
-fn message_digest(tag: &str, bytes: &[u8]) -> Message {
+pub(super) fn message_digest(tag: &str, bytes: &[u8]) -> Message {
        let tag = sha256::Hash::hash(tag.as_bytes());
        let merkle_root = root_hash(bytes);
        Message::from_slice(&tagged_hash(tag, merkle_root)).unwrap()
@@ -143,28 +143,38 @@ fn tagged_branch_hash_from_engine(
 
 /// [`Iterator`] over a sequence of bytes yielding [`TlvRecord`]s. The input is assumed to be a
 /// well-formed TLV stream.
-struct TlvStream<'a> {
+#[derive(Clone)]
+pub(super) struct TlvStream<'a> {
        data: io::Cursor<&'a [u8]>,
 }
 
 impl<'a> TlvStream<'a> {
-       fn new(data: &'a [u8]) -> Self {
+       pub fn new(data: &'a [u8]) -> Self {
                Self {
                        data: io::Cursor::new(data),
                }
        }
 
+       pub fn range<T>(self, types: T) -> impl core::iter::Iterator<Item = TlvRecord<'a>>
+       where
+               T: core::ops::RangeBounds<u64> + Clone,
+       {
+               let take_range = types.clone();
+               self.skip_while(move |record| !types.contains(&record.r#type))
+                       .take_while(move |record| take_range.contains(&record.r#type))
+       }
+
        fn skip_signatures(self) -> core::iter::Filter<TlvStream<'a>, fn(&TlvRecord) -> bool> {
                self.filter(|record| !SIGNATURE_TYPES.contains(&record.r#type))
        }
 }
 
 /// A slice into a [`TlvStream`] for a record.
-struct TlvRecord<'a> {
-       r#type: u64,
+pub(super) struct TlvRecord<'a> {
+       pub(super) r#type: u64,
        type_bytes: &'a [u8],
        // The entire TLV record.
-       record_bytes: &'a [u8],
+       pub(super) record_bytes: &'a [u8],
 }
 
 impl<'a> Iterator for TlvStream<'a> {
@@ -212,10 +222,11 @@ impl<'a> Writeable for WithoutSignatures<'a> {
 
 #[cfg(test)]
 mod tests {
-       use super::{TlvStream, WithoutSignatures};
+       use super::{SIGNATURE_TYPES, TlvStream, WithoutSignatures};
 
        use bitcoin::hashes::{Hash, sha256};
        use bitcoin::secp256k1::{KeyPair, Secp256k1, SecretKey};
+       use bitcoin::secp256k1::schnorr::Signature;
        use core::convert::Infallible;
        use crate::offers::offer::{Amount, OfferBuilder};
        use crate::offers::invoice_request::InvoiceRequest;
@@ -270,6 +281,10 @@ mod tests {
                        super::root_hash(&invoice_request.bytes[..]),
                        sha256::Hash::from_slice(&hex::decode("608407c18ad9a94d9ea2bcdbe170b6c20c462a7833a197621c916f78cf18e624").unwrap()).unwrap(),
                );
+               assert_eq!(
+                       invoice_request.signature(),
+                       Signature::from_slice(&hex::decode("b8f83ea3288cfd6ea510cdb481472575141e8d8744157f98562d162cc1c472526fdb24befefbdebab4dbb726bbd1b7d8aec057f8fa805187e5950d2bbe0e5642").unwrap()).unwrap(),
+               );
        }
 
        #[test]
@@ -302,6 +317,38 @@ mod tests {
                );
        }
 
+       #[test]
+       fn iterates_over_tlv_stream_range() {
+               let secp_ctx = Secp256k1::new();
+               let recipient_pubkey = {
+                       let secret_key = SecretKey::from_slice(&[41; 32]).unwrap();
+                       KeyPair::from_secret_key(&secp_ctx, &secret_key).public_key()
+               };
+               let payer_keys = {
+                       let secret_key = SecretKey::from_slice(&[42; 32]).unwrap();
+                       KeyPair::from_secret_key(&secp_ctx, &secret_key)
+               };
+
+               let invoice_request = OfferBuilder::new("foo".into(), recipient_pubkey)
+                       .amount_msats(100)
+                       .build_unchecked()
+                       .request_invoice(vec![0; 8], payer_keys.public_key()).unwrap()
+                       .build_unchecked()
+                       .sign::<_, Infallible>(|digest| Ok(secp_ctx.sign_schnorr_no_aux_rand(digest, &payer_keys)))
+                       .unwrap();
+
+               let tlv_stream = TlvStream::new(&invoice_request.bytes).range(0..1)
+                       .chain(TlvStream::new(&invoice_request.bytes).range(1..80))
+                       .chain(TlvStream::new(&invoice_request.bytes).range(80..160))
+                       .chain(TlvStream::new(&invoice_request.bytes).range(160..240))
+                       .chain(TlvStream::new(&invoice_request.bytes).range(SIGNATURE_TYPES))
+                       .map(|r| r.record_bytes.to_vec())
+                       .flatten()
+                       .collect::<Vec<u8>>();
+
+               assert_eq!(tlv_stream, invoice_request.bytes);
+       }
+
        impl AsRef<[u8]> for InvoiceRequest {
                fn as_ref(&self) -> &[u8] {
                        &self.bytes