From 4597d5abaff4c9a527b5139b476f55dc8eef9c46 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Tue, 25 Jun 2024 22:04:46 +0000
Subject: [PATCH] demo

---
 lightning/src/blinded_path/message.rs | 80 +++++++++++++++------------
 lightning/src/onion_message/packet.rs |  2 +-
 2 files changed, 47 insertions(+), 35 deletions(-)

diff --git a/lightning/src/blinded_path/message.rs b/lightning/src/blinded_path/message.rs
index 9a30d91bd..ca8215224 100644
--- a/lightning/src/blinded_path/message.rs
+++ b/lightning/src/blinded_path/message.rs
@@ -52,12 +52,21 @@ pub(crate) struct ForwardTlvs {
 	pub(crate) next_blinding_override: Option<PublicKey>,
 }
 
-/// Similar to [`ForwardTlvs`], but these TLVs are for the final node.
-pub(crate) struct ReceiveTlvs {
-	/// If `path_id` is `Some`, it is used to identify the blinded path that this onion message is
-	/// sending to. This is useful for receivers to check that said blinded path is being used in
-	/// the right context.
-	pub path_id: Option<MessageContext>
+
+/// Represents additional data included by the recipient in a [`BlindedPath`].
+///
+/// This data is encrypted and will be included when sending through
+/// [`BlindedPath`]. The recipient can authenticate the message and
+/// utilize it for further processing if needed.
+#[derive(Clone, Debug)]
+pub enum ReceiveTlvs {
+	/// Represents the data specific to [`OffersMessage`]
+	///
+	/// [`OffersMessage`]: crate::onion_message::offers::OffersMessage
+	Offers(OffersContext),
+	/// Represents the data appended with Custom Onion Message.
+	Custom(Vec<u8>),
+	Empty,
 }
 
 impl Writeable for ForwardTlvs {
@@ -78,38 +87,46 @@ impl Writeable for ForwardTlvs {
 
 impl Writeable for ReceiveTlvs {
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
-		// TODO: write padding
-		encode_tlv_stream!(writer, {
-			(6, self.path_id, option),
-		});
+		match self {
+			Self::Offers(o) => {
+				encode_tlv_stream!(writer, {
+					(6, (0u8, &o), required),
+				});
+			},
+			Self::Custom(v) => {
+				encode_tlv_stream!(writer, {
+					(6, (1u8, &v), required),
+				});
+			},
+			Self::Empty => {},
+		}
 		Ok(())
 	}
 }
 
 impl Readable for ReceiveTlvs {
 	fn read<R: io::Read>(r: &mut R) -> Result<Self, DecodeError> {
-		_init_and_read_tlv_stream!(r, {
-			(6, path_id, option),
-		});
-		Ok(ReceiveTlvs { path_id })
+		let mut res = Self::Empty;
+		decode_tlv_stream_with_custom_tlv_decode!(r, {},
+			|ty, reader| {
+				if ty == 6 {
+					let ty: u8 = Readable::read(reader)?;
+					match ty {
+						0 => { res = Self::Offers(Readable::read(reader)?); }
+						1 => { res = Self::Custom(Readable::read(reader)?); }
+						_ if ty % 2 == 0 => { return Err(DecodeError::UnknownRequiredFeature); }
+						_ => {},
+					}
+					Ok(true)
+				} else {
+					Ok(false)
+				}
+			}
+		);
+		Ok(res)
 	}
 }
 
-/// Represents additional data included by the recipient in a [`BlindedPath`].
-///
-/// This data is encrypted and will be included when sending through
-/// [`BlindedPath`]. The recipient can authenticate the message and
-/// utilize it for further processing if needed.
-#[derive(Clone, Debug)]
-pub enum MessageContext {
-	/// Represents the data specific to [`OffersMessage`]
-	///
-	/// [`OffersMessage`]: crate::onion_message::offers::OffersMessage
-	Offers(OffersContext),
-	/// Represents the data appended with Custom Onion Message.
-	Custom(Vec<u8>),
-}
-
 /// Contains the data specific to [`OffersMessage`]
 ///
 /// [`OffersMessage`]: crate::onion_message::offers::OffersMessage
@@ -125,11 +142,6 @@ pub enum OffersContext {
 
 }
 
-impl_writeable_tlv_based_enum!(MessageContext, ;
-	(0, Offers),
-	(1, Custom),
-);
-
 impl_writeable_tlv_based_enum!(OffersContext,
 	(0, Unknown) => {},
 	(1, OutboundPayment) => {
diff --git a/lightning/src/onion_message/packet.rs b/lightning/src/onion_message/packet.rs
index fca7dd6a9..3bd2f56fd 100644
--- a/lightning/src/onion_message/packet.rs
+++ b/lightning/src/onion_message/packet.rs
@@ -341,7 +341,7 @@ impl Readable for ControlTlvs {
 			})
 		} else if valid_recv_fmt {
 			ControlTlvs::Receive(ReceiveTlvs {
-				path_id,
+				path_id, // this would be unwrap_or(::Empty)
 			})
 		} else {
 			return Err(DecodeError::InvalidValue)
-- 
2.39.5