Parameterize OnionMessenger by new CustomOnionMessageHandler trait

OnionMessenger::new will now take a custom onion message handler trait
implementation. This handler will be used in upcoming commit(s) to handle
inbound custom onion messages.

The new trait also specifies what custom messages are supported via its
associated type, CustomMessage. This associated type must implement a new
CustomOnionMessagesContents trait, which requires custom messages to support
being written, being read, and supplying their TLV type.

Author: valentinewallace

fuzz/src/onion_message.rs

@@ -10,12 +10,12 @@ 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::{Readable, Writer};
-use lightning::onion_message::OnionMessenger;
+use lightning::util::ser::{MaybeReadableArgs, Readable, Writeable, Writer};
+use lightning::onion_message::{CustomOnionMessageContents, CustomOnionMessageHandler, OnionMessenger};
 
 use utils::test_logger;
 
-use std::io::Cursor;
+use std::io::{self, Cursor};
 use std::sync::atomic::{AtomicU64, Ordering};
 
 #[inline]
@@ -29,7 +29,8 @@ pub fn do_test<L: Logger>(data: &[u8], logger: &L) {
 			node_secret: secret,
 			counter: AtomicU64::new(0),
 		};
-		let onion_messenger = OnionMessenger::new(&keys_manager, logger);
+		let custom_msg_handler = TestCustomMessageHandler {};
+		let onion_messenger = OnionMessenger::new(&keys_manager, logger, &custom_msg_handler);
 		let mut pk = [2; 33]; pk[1] = 0xff;
 		let peer_node_id_not_used = PublicKey::from_slice(&pk).unwrap();
 		onion_messenger.handle_onion_message(&peer_node_id_not_used, &msg);
@@ -49,6 +50,38 @@ pub extern "C" fn onion_message_run(data: *const u8, datalen: usize) {
 	do_test(unsafe { std::slice::from_raw_parts(data, datalen) }, &logger);
 }
 
+struct TestCustomMessage {}
+
+const CUSTOM_MESSAGE_TYPE: u64 = 4242;
+const CUSTOM_MESSAGE_CONTENTS: [u8; 32] = [42; 32];
+
+impl CustomOnionMessageContents for TestCustomMessage {
+	fn tlv_type(&self) -> u64 {
+		CUSTOM_MESSAGE_TYPE
+	}
+}
+
+impl Writeable for TestCustomMessage {
+	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
+		Ok(CUSTOM_MESSAGE_CONTENTS.write(w)?)
+	}
+}
+
+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) {}
+}
+
 pub struct VecWriter(pub Vec<u8>);
 impl Writer for VecWriter {
 	fn write_all(&mut self, buf: &[u8]) -> Result<(), ::std::io::Error> {

lightning/src/ln/peer_handler.rs

@@ -21,11 +21,11 @@ use ln::features::{InitFeatures, NodeFeatures};
 use ln::msgs;
 use ln::msgs::{ChannelMessageHandler, LightningError, NetAddress, OnionMessageHandler, RoutingMessageHandler};
 use ln::channelmanager::{SimpleArcChannelManager, SimpleRefChannelManager};
-use util::ser::{VecWriter, Writeable, Writer};
+use util::ser::{MaybeReadableArgs, VecWriter, Writeable, Writer};
 use ln::peer_channel_encryptor::{PeerChannelEncryptor,NextNoiseStep};
 use ln::wire;
 use ln::wire::Encode;
-use onion_message::{SimpleArcOnionMessenger, SimpleRefOnionMessenger};
+use onion_message::{CustomOnionMessageContents, CustomOnionMessageHandler, SimpleArcOnionMessenger, SimpleRefOnionMessenger};
 use routing::gossip::{NetworkGraph, P2PGossipSync};
 use util::atomic_counter::AtomicCounter;
 use util::crypto::sign;
@@ -95,6 +95,23 @@ impl OnionMessageHandler for IgnoringMessageHandler {
 		InitFeatures::empty()
 	}
 }
+impl CustomOnionMessageHandler for IgnoringMessageHandler {
+	type CustomMessage = Infallible;
+	fn handle_custom_message(&self, _msg: Self::CustomMessage) {
+		// 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 {
+		Ok(None)
+	}
+}
+
+impl CustomOnionMessageContents for Infallible {
+	fn tlv_type(&self) -> u64 { unreachable!(); }
+}
+
 impl Deref for IgnoringMessageHandler {
 	type Target = IgnoringMessageHandler;
 	fn deref(&self) -> &Self { self }

lightning/src/onion_message/functional_tests.rs

@@ -11,19 +11,21 @@
 
 use chain::keysinterface::{KeysInterface, Recipient};
 use ln::features::InitFeatures;
-use ln::msgs::{self, OnionMessageHandler};
-use super::{BlindedRoute, Destination, OnionMessenger, SendError};
+use ln::msgs::{self, DecodeError, OnionMessageHandler};
+use super::{BlindedRoute, CustomOnionMessageContents, CustomOnionMessageHandler, Destination, OnionMessenger, SendError};
 use util::enforcing_trait_impls::EnforcingSigner;
+use util::ser::{MaybeReadableArgs, Writeable, Writer};
 use util::test_utils;
 
 use bitcoin::network::constants::Network;
 use bitcoin::secp256k1::{PublicKey, Secp256k1};
 
+use io;
 use sync::Arc;
 
 struct MessengerNode {
 	keys_manager: Arc<test_utils::TestKeysInterface>,
-	messenger: OnionMessenger<EnforcingSigner, Arc<test_utils::TestKeysInterface>, Arc<test_utils::TestLogger>>,
+	messenger: OnionMessenger<EnforcingSigner, Arc<test_utils::TestKeysInterface>, Arc<test_utils::TestLogger>, Arc<TestCustomMessageHandler>>,
 	logger: Arc<test_utils::TestLogger>,
 }
 
@@ -34,6 +36,43 @@ impl MessengerNode {
 	}
 }
 
+#[derive(Clone)]
+struct TestCustomMessage {}
+
+const CUSTOM_MESSAGE_TYPE: u64 = 4242;
+const CUSTOM_MESSAGE_CONTENTS: [u8; 32] = [42; 32];
+
+impl CustomOnionMessageContents for TestCustomMessage {
+	fn tlv_type(&self) -> u64 {
+		CUSTOM_MESSAGE_TYPE
+	}
+}
+
+impl Writeable for TestCustomMessage {
+	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
+		Ok(CUSTOM_MESSAGE_CONTENTS.write(w)?)
+	}
+}
+
+impl MaybeReadableArgs<u64> for TestCustomMessage {
+	fn read<R: io::Read>(buffer: &mut R, message_type: u64) -> Result<Option<Self>, DecodeError> where Self: Sized {
+		if message_type == CUSTOM_MESSAGE_TYPE {
+			let mut buf = Vec::new();
+			buffer.read_to_end(&mut buf)?;
+			assert_eq!(buf, CUSTOM_MESSAGE_CONTENTS);
+			return Ok(Some(TestCustomMessage {}))
+		}
+		Ok(None)
+	}
+}
+
+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 {
@@ -42,7 +81,7 @@ fn create_nodes(num_messengers: u8) -> Vec<MessengerNode> {
 		let keys_manager = Arc::new(test_utils::TestKeysInterface::new(&seed, Network::Testnet));
 		nodes.push(MessengerNode {
 			keys_manager: keys_manager.clone(),
-			messenger: OnionMessenger::new(keys_manager, logger.clone()),
+			messenger: OnionMessenger::new(keys_manager, logger.clone(), Arc::new(TestCustomMessageHandler {})),
 			logger,
 		});
 	}

lightning/src/onion_message/messenger.rs

@@ -19,7 +19,9 @@ use chain::keysinterface::{InMemorySigner, KeysInterface, KeysManager, Recipient
 use ln::features::{InitFeatures, NodeFeatures};
 use ln::msgs::{self, OnionMessageHandler};
 use ln::onion_utils;
+use ln::peer_handler::IgnoringMessageHandler;
 use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
+pub use super::packet::CustomOnionMessageContents;
 use super::packet::{BIG_PACKET_HOP_DATA_LEN, ForwardControlTlvs, Packet, Payload, ReceiveControlTlvs, SMALL_PACKET_HOP_DATA_LEN};
 use super::utils;
 use util::events::OnionMessageProvider;
@@ -41,8 +43,12 @@ use prelude::*;
 /// # use bitcoin::hashes::_export::_core::time::Duration;
 /// # use bitcoin::secp256k1::{PublicKey, Secp256k1, SecretKey};
 /// # use lightning::chain::keysinterface::{InMemorySigner, KeysManager, KeysInterface};
-/// # use lightning::onion_message::{BlindedRoute, Destination, OnionMessenger};
+/// # use lightning::ln::msgs::DecodeError;
+/// # use lightning::ln::peer_handler::IgnoringMessageHandler;
+/// # use lightning::onion_message::{BlindedRoute, CustomOnionMessageContents, Destination, OnionMessenger};
 /// # use lightning::util::logger::{Logger, Record};
+/// # use lightning::util::ser::{MaybeReadableArgs, Writeable, Writer};
+/// # use lightning::io;
 /// # use std::sync::Arc;
 /// # struct FakeLogger {};
 /// # impl Logger for FakeLogger {
@@ -58,12 +64,31 @@ use prelude::*;
 /// # let (hop_node_id2, hop_node_id3, hop_node_id4) = (hop_node_id1, hop_node_id1,
 /// hop_node_id1);
 /// # let destination_node_id = hop_node_id1;
-/// #
+/// # let your_custom_message_handler = IgnoringMessageHandler {};
 /// // Create the onion messenger. This must use the same `keys_manager` as is passed to your
 /// // ChannelManager.
-/// let onion_messenger = OnionMessenger::new(&keys_manager, logger);
+/// let onion_messenger = OnionMessenger::new(&keys_manager, logger, your_custom_message_handler);
 ///
-/// // Send an empty onion message to a node id.
+/// # struct YourCustomMessage {}
+/// impl Writeable for YourCustomMessage {
+/// 	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
+/// 		# Ok(())
+/// 		// Write your custom onion message to `w`
+/// 	}
+/// }
+/// impl CustomOnionMessageContents for YourCustomMessage {
+/// 	fn tlv_type(&self) -> u64 {
+/// 		# let your_custom_message_type = 42;
+/// 		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
+/// 	}
+/// }
 /// let intermediate_hops = [hop_node_id1, hop_node_id2];
 /// let reply_path = None;
 /// onion_messenger.send_onion_message(&intermediate_hops, Destination::Node(destination_node_id), reply_path);
@@ -81,17 +106,18 @@ use prelude::*;
 ///
 /// [offers]: <https://github.com/lightning/bolts/pull/798>
 /// [`OnionMessenger`]: crate::onion_message::OnionMessenger
-pub struct OnionMessenger<Signer: Sign, K: Deref, L: Deref>
+pub struct OnionMessenger<Signer: Sign, K: Deref, L: Deref, CMH: Deref>
 	where K::Target: KeysInterface<Signer = Signer>,
 	      L::Target: Logger,
+	      CMH:: Target: CustomOnionMessageHandler,
 {
 	keys_manager: K,
 	logger: L,
 	pending_messages: Mutex<HashMap<PublicKey, VecDeque<msgs::OnionMessage>>>,
 	secp_ctx: Secp256k1<secp256k1::All>,
+	custom_handler: CMH,
 	// Coming soon:
 	// invoice_handler: InvoiceHandler,
-	// custom_handler: CustomHandler, // handles custom onion messages
 }
 
 /// The destination of an onion message.
@@ -130,20 +156,40 @@ pub enum SendError {
 	BufferFull,
 }
 
-impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
+/// Handler for custom onion messages. If you are using [`SimpleArcOnionMessenger`],
+/// [`SimpleRefOnionMessenger`], or prefer to ignore inbound custom onion messages,
+/// [`IgnoringMessageHandler`] must be provided to [`OnionMessenger::new`]. Otherwise, a custom
+/// implementation of this trait must be provided, with [`CustomMessage`] specifying the supported
+/// message types.
+///
+/// See [`OnionMessenger`] for example usage.
+///
+/// [`IgnoringMessageHandler`]: crate::ln::peer_handler::IgnoringMessageHandler
+/// [`CustomMessage`]: Self::CustomMessage
+pub trait CustomOnionMessageHandler {
+	/// The message known to the handler. To support multiple message types, you may want to make this
+	/// an enum with a variant for each supported message.
+	type CustomMessage: CustomOnionMessageContents;
+	/// Called with the custom message that was received.
+	fn handle_custom_message(&self, msg: Self::CustomMessage);
+}
+
+impl<Signer: Sign, K: Deref, L: Deref, CMH: Deref> OnionMessenger<Signer, K, L, CMH>
 	where K::Target: KeysInterface<Signer = Signer>,
 	      L::Target: Logger,
+	      CMH::Target: CustomOnionMessageHandler,
 {
 	/// Constructs a new `OnionMessenger` to send, forward, and delegate received onion messages to
 	/// their respective handlers.
-	pub fn new(keys_manager: K, logger: L) -> Self {
+	pub fn new(keys_manager: K, logger: L, custom_handler: CMH) -> Self {
 		let mut secp_ctx = Secp256k1::new();
 		secp_ctx.seeded_randomize(&keys_manager.get_secure_random_bytes());
 		OnionMessenger {
 			keys_manager,
 			pending_messages: Mutex::new(HashMap::new()),
 			secp_ctx,
 			logger,
+			custom_handler,
 		}
 	}
 
@@ -221,9 +267,10 @@ fn outbound_buffer_full(peer_node_id: &PublicKey, buffer: &HashMap<PublicKey, Ve
 	false
 }
 
-impl<Signer: Sign, K: Deref, L: Deref> OnionMessageHandler for OnionMessenger<Signer, K, L>
+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,
 {
 	/// 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
@@ -361,9 +408,10 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessageHandler for OnionMessenger<Si
 	}
 }
 
-impl<Signer: Sign, K: Deref, L: Deref> OnionMessageProvider for OnionMessenger<Signer, K, L>
+impl<Signer: Sign, K: Deref, L: Deref, CMH: Deref> OnionMessageProvider for OnionMessenger<Signer, K, L, CMH>
 	where K::Target: KeysInterface<Signer = Signer>,
 	      L::Target: Logger,
+	      CMH::Target: CustomOnionMessageHandler,
 {
 	fn next_onion_message_for_peer(&self, peer_node_id: PublicKey) -> Option<msgs::OnionMessage> {
 		let mut pending_msgs = self.pending_messages.lock().unwrap();
@@ -383,15 +431,15 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessageProvider for OnionMessenger<S
 ///
 /// [`SimpleArcChannelManager`]: crate::ln::channelmanager::SimpleArcChannelManager
 /// [`SimpleArcPeerManager`]: crate::ln::peer_handler::SimpleArcPeerManager
-pub type SimpleArcOnionMessenger<L> = OnionMessenger<InMemorySigner, Arc<KeysManager>, Arc<L>>;
+pub type SimpleArcOnionMessenger<L> = OnionMessenger<InMemorySigner, Arc<KeysManager>, Arc<L>, IgnoringMessageHandler>;
 /// Useful for simplifying the parameters of [`SimpleRefChannelManager`] and
 /// [`SimpleRefPeerManager`]. See their docs for more details.
 ///
 /// (C-not exported) as general type aliases don't make sense in bindings.
 ///
 /// [`SimpleRefChannelManager`]: crate::ln::channelmanager::SimpleRefChannelManager
 /// [`SimpleRefPeerManager`]: crate::ln::peer_handler::SimpleRefPeerManager
-pub type SimpleRefOnionMessenger<'a, 'b, L> = OnionMessenger<InMemorySigner, &'a KeysManager, &'b L>;
+pub type SimpleRefOnionMessenger<'a, 'b, L> = OnionMessenger<InMemorySigner, &'a KeysManager, &'b L, IgnoringMessageHandler>;
 
 /// Construct onion packet payloads and keys for sending an onion message along the given
 /// `unblinded_path` to the given `destination`.

lightning/src/onion_message/mod.rs

@@ -29,5 +29,5 @@ mod functional_tests;
 
 // Re-export structs so they can be imported with just the `onion_message::` module prefix.
 pub use self::blinded_route::{BlindedRoute, BlindedHop};
-pub use self::messenger::{Destination, OnionMessenger, SendError, SimpleArcOnionMessenger, SimpleRefOnionMessenger};
+pub use self::messenger::{CustomOnionMessageContents, CustomOnionMessageHandler, Destination, OnionMessenger, SendError, SimpleArcOnionMessenger, SimpleRefOnionMessenger};
 pub(crate) use self::packet::Packet;

lightning/src/onion_message/packet.rs

@@ -16,7 +16,7 @@ use ln::msgs::DecodeError;
 use ln::onion_utils;
 use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
 use util::chacha20poly1305rfc::{ChaChaPolyReadAdapter, ChaChaPolyWriteAdapter};
-use util::ser::{BigSize, FixedLengthReader, LengthRead, LengthReadable, LengthReadableArgs, Readable, ReadableArgs, Writeable, Writer};
+use util::ser::{BigSize, FixedLengthReader, LengthRead, LengthReadable, LengthReadableArgs, MaybeReadableArgs, Readable, ReadableArgs, Writeable, Writer};
 
 use core::cmp;
 use io::{self, Read};
@@ -112,6 +112,13 @@ pub(super) enum Payload {
 //	CustomMessage<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> {
+	/// Returns the TLV type identifying the message contents. MUST be >= 64.
+	fn tlv_type(&self) -> u64;
+}
+
 /// Forward control TLVs in their blinded and unblinded form.
 pub(super) enum ForwardControlTlvs {
 	/// If we're sending to a blinded route, the node that constructed the blinded route has provided