Merge pull request #1809 from valentinewallace/2022-10-custom-om-self

Give us a self when reading a custom onion message

Author: TheBlueMatt

fuzz/src/onion_message.rs

@@ -10,7 +10,7 @@ use lightning::ln::msgs::{self, DecodeError, OnionMessageHandler};
 use lightning::ln::script::ShutdownScript;
 use lightning::util::enforcing_trait_impls::EnforcingSigner;
 use lightning::util::logger::Logger;
-use lightning::util::ser::{MaybeReadableArgs, Readable, Writeable, Writer};
+use lightning::util::ser::{Readable, Writeable, Writer};
 use lightning::onion_message::{CustomOnionMessageContents, CustomOnionMessageHandler, OnionMessenger};
 
 use crate::utils::test_logger;
@@ -67,19 +67,16 @@ impl Writeable for TestCustomMessage {
 	}
 }
 
-impl MaybeReadableArgs<u64> for TestCustomMessage {
-	fn read<R: io::Read>(buffer: &mut R, _message_type: u64,) -> Result<Option<Self>, DecodeError> where Self: Sized {
-		let mut buf = Vec::new();
-		buffer.read_to_end(&mut buf)?;
-		return Ok(Some(TestCustomMessage {}))
-	}
-}
-
 struct TestCustomMessageHandler {}
 
 impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
 	fn handle_custom_message(&self, _msg: Self::CustomMessage) {}
+	fn read_custom_message<R: io::Read>(&self, _message_type: u64, buffer: &mut R) -> Result<Option<Self::CustomMessage>, msgs::DecodeError> {
+		let mut buf = Vec::new();
+		buffer.read_to_end(&mut buf)?;
+		return Ok(Some(TestCustomMessage {}))
+	}
 }
 
 pub struct VecWriter(pub Vec<u8>);

lightning/src/ln/onion_utils.rs

@@ -589,31 +589,6 @@ pub(super) fn process_onion_failure<T: secp256k1::Signing, L: Deref>(secp_ctx: &
 	} else { unreachable!(); }
 }
 
-/// An input used when decoding an onion packet.
-pub(crate) trait DecodeInput {
-	type Arg;
-	/// If Some, this is the input when checking the hmac of the onion packet.
-	fn payment_hash(&self) -> Option<&PaymentHash>;
-	/// Read argument when decrypting our hop payload.
-	fn read_arg(self) -> Self::Arg;
-}
-
-impl DecodeInput for PaymentHash {
-	type Arg = ();
-	fn payment_hash(&self) -> Option<&PaymentHash> {
-		Some(self)
-	}
-	fn read_arg(self) -> Self::Arg { () }
-}
-
-impl DecodeInput for SharedSecret {
-	type Arg = SharedSecret;
-	fn payment_hash(&self) -> Option<&PaymentHash> {
-		None
-	}
-	fn read_arg(self) -> Self::Arg { self }
-}
-
 /// Allows `decode_next_hop` to return the next hop packet bytes for either payments or onion
 /// message forwards.
 pub(crate) trait NextPacketBytes: AsMut<[u8]> {
@@ -664,7 +639,7 @@ pub(crate) enum OnionDecodeErr {
 }
 
 pub(crate) fn decode_next_payment_hop(shared_secret: [u8; 32], hop_data: &[u8], hmac_bytes: [u8; 32], payment_hash: PaymentHash) -> Result<Hop, OnionDecodeErr> {
-	match decode_next_hop(shared_secret, hop_data, hmac_bytes, payment_hash) {
+	match decode_next_hop(shared_secret, hop_data, hmac_bytes, Some(payment_hash), ()) {
 		Ok((next_hop_data, None)) => Ok(Hop::Receive(next_hop_data)),
 		Ok((next_hop_data, Some((next_hop_hmac, FixedSizeOnionPacket(new_packet_bytes))))) => {
 			Ok(Hop::Forward {
@@ -677,12 +652,16 @@ pub(crate) fn decode_next_payment_hop(shared_secret: [u8; 32], hop_data: &[u8],
 	}
 }
 
-pub(crate) fn decode_next_hop<D: DecodeInput, R: ReadableArgs<D::Arg>, N: NextPacketBytes>(shared_secret: [u8; 32], hop_data: &[u8], hmac_bytes: [u8; 32], decode_input: D) -> Result<(R, Option<([u8; 32], N)>), OnionDecodeErr> {
+pub(crate) fn decode_next_untagged_hop<T, R: ReadableArgs<T>, N: NextPacketBytes>(shared_secret: [u8; 32], hop_data: &[u8], hmac_bytes: [u8; 32], read_args: T) -> Result<(R, Option<([u8; 32], N)>), OnionDecodeErr> {
+	decode_next_hop(shared_secret, hop_data, hmac_bytes, None, read_args)
+}
+
+fn decode_next_hop<T, R: ReadableArgs<T>, N: NextPacketBytes>(shared_secret: [u8; 32], hop_data: &[u8], hmac_bytes: [u8; 32], payment_hash: Option<PaymentHash>, read_args: T) -> Result<(R, Option<([u8; 32], N)>), OnionDecodeErr> {
 	let (rho, mu) = gen_rho_mu_from_shared_secret(&shared_secret);
 	let mut hmac = HmacEngine::<Sha256>::new(&mu);
 	hmac.input(hop_data);
-	if let Some(payment_hash) = decode_input.payment_hash() {
-		hmac.input(&payment_hash.0[..]);
+	if let Some(tag) = payment_hash {
+		hmac.input(&tag.0[..]);
 	}
 	if !fixed_time_eq(&Hmac::from_engine(hmac).into_inner(), &hmac_bytes) {
 		return Err(OnionDecodeErr::Malformed {
@@ -693,7 +672,7 @@ pub(crate) fn decode_next_hop<D: DecodeInput, R: ReadableArgs<D::Arg>, N: NextPa
 
 	let mut chacha = ChaCha20::new(&rho, &[0u8; 8]);
 	let mut chacha_stream = ChaChaReader { chacha: &mut chacha, read: Cursor::new(&hop_data[..]) };
-	match R::read(&mut chacha_stream, decode_input.read_arg()) {
+	match R::read(&mut chacha_stream, read_args) {
 		Err(err) => {
 			let error_code = match err {
 				msgs::DecodeError::UnknownVersion => 0x4000 | 1, // unknown realm byte

lightning/src/ln/peer_handler.rs

@@ -21,7 +21,7 @@ use crate::ln::features::{InitFeatures, NodeFeatures};
 use crate::ln::msgs;
 use crate::ln::msgs::{ChannelMessageHandler, LightningError, NetAddress, OnionMessageHandler, RoutingMessageHandler};
 use crate::ln::channelmanager::{SimpleArcChannelManager, SimpleRefChannelManager};
-use crate::util::ser::{MaybeReadableArgs, VecWriter, Writeable, Writer};
+use crate::util::ser::{VecWriter, Writeable, Writer};
 use crate::ln::peer_channel_encryptor::{PeerChannelEncryptor,NextNoiseStep};
 use crate::ln::wire;
 use crate::ln::wire::Encode;
@@ -97,13 +97,11 @@ impl OnionMessageHandler for IgnoringMessageHandler {
 }
 impl CustomOnionMessageHandler for IgnoringMessageHandler {
 	type CustomMessage = Infallible;
-	fn handle_custom_message(&self, _msg: Self::CustomMessage) {
+	fn handle_custom_message(&self, _msg: Infallible) {
 		// Since we always return `None` in the read the handle method should never be called.
 		unreachable!();
 	}
-}
-impl MaybeReadableArgs<u64> for Infallible {
-	fn read<R: io::Read>(_buffer: &mut R, _msg_type: u64) -> Result<Option<Self>, msgs::DecodeError> where Self: Sized {
+	fn read_custom_message<R: io::Read>(&self, _msg_type: u64, _buffer: &mut R) -> Result<Option<Infallible>, msgs::DecodeError> where Self: Sized {
 		Ok(None)
 	}
 }

lightning/src/onion_message/functional_tests.rs

@@ -14,7 +14,7 @@ use crate::ln::features::InitFeatures;
 use crate::ln::msgs::{self, DecodeError, OnionMessageHandler};
 use super::{BlindedRoute, CustomOnionMessageContents, CustomOnionMessageHandler, Destination, OnionMessageContents, OnionMessenger, SendError};
 use crate::util::enforcing_trait_impls::EnforcingSigner;
-use crate::util::ser::{MaybeReadableArgs, Writeable, Writer};
+use crate::util::ser::{ Writeable, Writer};
 use crate::util::test_utils;
 
 use bitcoin::network::constants::Network;
@@ -54,8 +54,12 @@ impl Writeable for TestCustomMessage {
 	}
 }
 
-impl MaybeReadableArgs<u64> for TestCustomMessage {
-	fn read<R: io::Read>(buffer: &mut R, message_type: u64) -> Result<Option<Self>, DecodeError> where Self: Sized {
+struct TestCustomMessageHandler {}
+
+impl CustomOnionMessageHandler for TestCustomMessageHandler {
+	type CustomMessage = TestCustomMessage;
+	fn handle_custom_message(&self, _msg: Self::CustomMessage) {}
+	fn read_custom_message<R: io::Read>(&self, message_type: u64, buffer: &mut R) -> Result<Option<Self::CustomMessage>, DecodeError> where Self: Sized {
 		if message_type == CUSTOM_MESSAGE_TYPE {
 			let mut buf = Vec::new();
 			buffer.read_to_end(&mut buf)?;
@@ -66,13 +70,6 @@ impl MaybeReadableArgs<u64> for TestCustomMessage {
 	}
 }
 
-struct TestCustomMessageHandler {}
-
-impl CustomOnionMessageHandler for TestCustomMessageHandler {
-	type CustomMessage = TestCustomMessage;
-	fn handle_custom_message(&self, _msg: Self::CustomMessage) {}
-}
-
 fn create_nodes(num_messengers: u8) -> Vec<MessengerNode> {
 	let mut nodes = Vec::new();
 	for i in 0..num_messengers {
@@ -233,12 +230,6 @@ fn invalid_custom_message_type() {
 		fn write<W: Writer>(&self, _w: &mut W) -> Result<(), io::Error> { unreachable!() }
 	}
 
-	impl MaybeReadableArgs<u64> for InvalidCustomMessage {
-		fn read<R: io::Read>(_buffer: &mut R, _message_type: u64) -> Result<Option<Self>, DecodeError> where Self: Sized {
-			unreachable!()
-		}
-	}
-
 	let test_msg = OnionMessageContents::Custom(InvalidCustomMessage {});
 	let err = nodes[0].messenger.send_onion_message(&[], Destination::Node(nodes[1].get_node_pk()), test_msg, None).unwrap_err();
 	assert_eq!(err, SendError::InvalidMessage);

lightning/src/onion_message/messenger.rs

@@ -29,6 +29,7 @@ use crate::util::logger::Logger;
 use crate::util::ser::Writeable;
 
 use core::ops::Deref;
+use crate::io;
 use crate::sync::{Arc, Mutex};
 use crate::prelude::*;
 
@@ -47,7 +48,7 @@ use crate::prelude::*;
 /// # use lightning::ln::peer_handler::IgnoringMessageHandler;
 /// # use lightning::onion_message::{BlindedRoute, CustomOnionMessageContents, Destination, OnionMessageContents, OnionMessenger};
 /// # use lightning::util::logger::{Logger, Record};
-/// # use lightning::util::ser::{MaybeReadableArgs, Writeable, Writer};
+/// # use lightning::util::ser::{Writeable, Writer};
 /// # use lightning::io;
 /// # use std::sync::Arc;
 /// # struct FakeLogger {};
@@ -81,13 +82,6 @@ use crate::prelude::*;
 /// 		your_custom_message_type
 /// 	}
 /// }
-/// impl MaybeReadableArgs<u64> for YourCustomMessage {
-/// 	fn read<R: io::Read>(r: &mut R, message_type: u64) -> Result<Option<Self>, DecodeError> {
-/// 		# unreachable!()
-/// 		// Read your custom onion message of type `message_type` from `r`, or return `None`
-/// 		// if the message type is unknown
-/// 	}
-/// }
 /// // Send a custom onion message to a node id.
 /// let intermediate_hops = [hop_node_id1, hop_node_id2];
 /// let reply_path = None;
@@ -178,6 +172,9 @@ pub trait CustomOnionMessageHandler {
 	type CustomMessage: CustomOnionMessageContents;
 	/// Called with the custom message that was received.
 	fn handle_custom_message(&self, msg: Self::CustomMessage);
+	/// Read a custom message of type `message_type` from `buffer`, returning `Ok(None)` if the
+	/// message type is unknown.
+	fn read_custom_message<R: io::Read>(&self, message_type: u64, buffer: &mut R) -> Result<Option<Self::CustomMessage>, msgs::DecodeError>;
 }
 
 impl<Signer: Sign, K: Deref, L: Deref, CMH: Deref> OnionMessenger<Signer, K, L, CMH>
@@ -279,7 +276,7 @@ fn outbound_buffer_full(peer_node_id: &PublicKey, buffer: &HashMap<PublicKey, Ve
 impl<Signer: Sign, K: Deref, L: Deref, CMH: Deref> OnionMessageHandler for OnionMessenger<Signer, K, L, CMH>
 	where K::Target: KeysInterface<Signer = Signer>,
 	      L::Target: Logger,
-	      CMH::Target: CustomOnionMessageHandler,
+	      CMH::Target: CustomOnionMessageHandler + Sized,
 {
 	/// Handle an incoming onion message. Currently, if a message was destined for us we will log, but
 	/// soon we'll delegate the onion message to a handler that can generate invoices or send
@@ -308,8 +305,8 @@ impl<Signer: Sign, K: Deref, L: Deref, CMH: Deref> OnionMessageHandler for Onion
 				}
 			}
 		};
-		match onion_utils::decode_next_hop(onion_decode_ss, &msg.onion_routing_packet.hop_data[..],
-			msg.onion_routing_packet.hmac, control_tlvs_ss)
+		match onion_utils::decode_next_untagged_hop(onion_decode_ss, &msg.onion_routing_packet.hop_data[..],
+			msg.onion_routing_packet.hmac, (control_tlvs_ss, &*self.custom_handler))
 		{
 			Ok((Payload::Receive::<<<CMH as Deref>::Target as CustomOnionMessageHandler>::CustomMessage> {
 				message, control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id }), reply_path,

lightning/src/onion_message/packet.rs

@@ -15,8 +15,9 @@ use bitcoin::secp256k1::ecdh::SharedSecret;
 use crate::ln::msgs::DecodeError;
 use crate::ln::onion_utils;
 use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
+use super::messenger::CustomOnionMessageHandler;
 use crate::util::chacha20poly1305rfc::{ChaChaPolyReadAdapter, ChaChaPolyWriteAdapter};
-use crate::util::ser::{BigSize, FixedLengthReader, LengthRead, LengthReadable, LengthReadableArgs, MaybeReadableArgs, Readable, ReadableArgs, Writeable, Writer};
+use crate::util::ser::{BigSize, FixedLengthReader, LengthRead, LengthReadable, LengthReadableArgs, Readable, ReadableArgs, Writeable, Writer};
 
 use core::cmp;
 use crate::io::{self, Read};
@@ -106,7 +107,7 @@ pub(super) enum Payload<T: CustomOnionMessageContents> {
 #[derive(Debug)]
 /// The contents of an onion message. In the context of offers, this would be the invoice, invoice
 /// request, or invoice error.
-pub enum OnionMessageContents<T> where T: CustomOnionMessageContents {
+pub enum OnionMessageContents<T: CustomOnionMessageContents> {
 	// Coming soon:
 	// Invoice,
 	// InvoiceRequest,
@@ -115,7 +116,7 @@ pub enum OnionMessageContents<T> where T: CustomOnionMessageContents {
 	Custom(T),
 }
 
-impl<T> OnionMessageContents<T> where T: CustomOnionMessageContents {
+impl<T: CustomOnionMessageContents> OnionMessageContents<T> {
 	/// Returns the type that was used to decode the message payload.
 	pub fn tlv_type(&self) -> u64 {
 		match self {
@@ -132,9 +133,8 @@ impl<T: CustomOnionMessageContents> Writeable for OnionMessageContents<T> {
 	}
 }
 
-/// The contents of a custom onion message. Must implement `MaybeReadableArgs<u64>` where the `u64`
-/// is the custom TLV type attempting to be read, and return `Ok(None)` if the TLV type is unknown.
-pub trait CustomOnionMessageContents: Writeable + MaybeReadableArgs<u64> {
+/// The contents of a custom onion message.
+pub trait CustomOnionMessageContents: Writeable {
 	/// Returns the TLV type identifying the message contents. MUST be >= 64.
 	fn tlv_type(&self) -> u64;
 }
@@ -198,8 +198,10 @@ impl<T: CustomOnionMessageContents> Writeable for (Payload<T>, [u8; 32]) {
 }
 
 // Uses the provided secret to simultaneously decode and decrypt the control TLVs and data TLV.
-impl<T: CustomOnionMessageContents> ReadableArgs<SharedSecret> for Payload<T> {
-	fn read<R: Read>(r: &mut R, encrypted_tlvs_ss: SharedSecret) -> Result<Self, DecodeError> {
+impl<H: CustomOnionMessageHandler> ReadableArgs<(SharedSecret, &H)> for Payload<<H as CustomOnionMessageHandler>::CustomMessage> {
+	fn read<R: Read>(r: &mut R, args: (SharedSecret, &H)) -> Result<Self, DecodeError> {
+		let (encrypted_tlvs_ss, handler) = args;
+
 		let v: BigSize = Readable::read(r)?;
 		let mut rd = FixedLengthReader::new(r, v.0);
 		let mut reply_path: Option<BlindedRoute> = None;
@@ -216,7 +218,7 @@ impl<T: CustomOnionMessageContents> ReadableArgs<SharedSecret> for Payload<T> {
 			if message_type.is_some() { return Err(DecodeError::InvalidValue) }
 
 			message_type = Some(msg_type);
-			match T::read(msg_reader, msg_type) {
+			match handler.read_custom_message(msg_type, msg_reader) {
 				Ok(Some(msg)) => {
 					message = Some(msg);
 					Ok(true)

lightning/src/util/ser.rs

@@ -269,15 +269,6 @@ impl<T: Readable> MaybeReadable for T {
 	}
 }
 
-/// A trait that various rust-lightning types implement allowing them to (maybe) be read in from a
-/// Read, given some additional set of arguments which is required to deserialize.
-///
-/// (C-not exported) as we only export serialization to/from byte arrays instead
-pub trait MaybeReadableArgs<P> {
-	/// Reads a Self in from the given Read
-	fn read<R: Read>(reader: &mut R, params: P) -> Result<Option<Self>, DecodeError> where Self: Sized;
-}
-
 pub(crate) struct OptionDeserWrapper<T: Readable>(pub Option<T>);
 impl<T: Readable> Readable for OptionDeserWrapper<T> {
 	#[inline]