Give us a self when reading a custom onion message

Co-authored-by: Matt Corallo <[email protected]>
Co-authored-by: Valentine Wallace <[email protected]>

Author: TheBlueMatt

lightning/src/ln/onion_utils.rs

@@ -606,14 +606,6 @@ impl DecodeInput for PaymentHash {
 	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]> {

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<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

@@ -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<R: io::Read>(&self, message_type: u64, buffer: &mut R) -> Result<Option<TestCustomMessage>, 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 {

lightning/src/onion_message/messenger.rs

@@ -14,8 +14,10 @@ use bitcoin::hashes::{Hash, HashEngine};
 use bitcoin::hashes::hmac::{Hmac, HmacEngine};
 use bitcoin::hashes::sha256::Hash as Sha256;
 use bitcoin::secp256k1::{self, PublicKey, Scalar, Secp256k1, SecretKey};
+use bitcoin::secp256k1::ecdh::SharedSecret;
 
 use crate::chain::keysinterface::{InMemorySigner, KeysInterface, KeysManager, Recipient, Sign};
+use crate::ln::PaymentHash;
 use crate::ln::features::{InitFeatures, NodeFeatures};
 use crate::ln::msgs::{self, OnionMessageHandler};
 use crate::ln::onion_utils;
@@ -29,6 +31,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::*;
 
@@ -178,6 +181,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<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>
@@ -276,10 +282,19 @@ fn outbound_buffer_full(peer_node_id: &PublicKey, buffer: &HashMap<PublicKey, Ve
 	false
 }
 
+impl<'a, T: CustomOnionMessageHandler> onion_utils::DecodeInput for (SharedSecret, &'a T) {
+	type Arg = (SharedSecret, &'a T);
+	fn payment_hash(&self) -> Option<&PaymentHash> {
+		None
+	}
+	fn read_arg(self) -> Self::Arg { self }
+}
+
 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, // XXX: No idea why rustc desperately
+	                                                      // wants Sized here, it's kinda dumb
 {
 	/// 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
@@ -309,7 +324,7 @@ 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)
+			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 {
@@ -134,7 +135,7 @@ 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> {
+pub trait CustomOnionMessageContents: Writeable {
 	/// Returns the TLV type identifying the message contents. MUST be >= 64.
 	fn tlv_type(&self) -> u64;
 }
@@ -198,8 +199,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 +219,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(msg_type, msg_reader) {
 				Ok(Some(msg)) => {
 					message = Some(msg);
 					Ok(true)