Add custom message handler to peer manager
[rust-lightning] / lightning / src / ln / wire.rs
index e3ff3095219e7abf1efe8ff829ad1d2b7ac65bd0..05bdb5694d8b4c4202348e4623917af41fe0e3e8 100644 (file)
@@ -7,26 +7,32 @@
 // You may not use this file except in accordance with one or both of these
 // licenses.
 
-//! Wire encoding/decoding for Lightning messages according to [BOLT #1].
-//!
-//! Messages known by this module can be read from the wire using [`read()`].
-//! The [`Message`] enum returned by [`read()`] wraps the decoded message or the message type (if
-//! unknown) to use with pattern matching.
-//!
-//! Messages implementing the [`Encode`] trait define a message type and can be sent over the wire
-//! using [`write()`].
-//!
+//! Wire encoding/decoding for Lightning messages according to [BOLT #1], and for
+//! custom message through the [`CustomMessageReader`] trait.
+//! 
 //! [BOLT #1]: https://github.com/lightningnetwork/lightning-rfc/blob/master/01-messaging.md
 
 use io;
 use ln::msgs;
 use util::ser::{Readable, Writeable, Writer};
 
+/// Trait to be implemented by custom message (unrelated to the channel/gossip LN layers)
+/// decoders.
+pub trait CustomMessageReader {
+       /// The type of the message decoded by the implementation.
+       type CustomMessage: core::fmt::Debug + Type + Writeable;
+       /// Decodes a custom message to `CustomMessageType`. If the given message type is known to the
+       /// implementation and the message could be decoded, must return `Ok(Some(message))`. If the
+       /// message type is unknown to the implementation, must return `Ok(None)`. If a decoding error
+       /// occur, must return `Err(DecodeError::X)` where `X` details the encountered error.
+       fn read<R: io::Read>(&self, message_type: u16, buffer: &mut R) -> Result<Option<Self::CustomMessage>, msgs::DecodeError>;
+}
+
 /// A Lightning message returned by [`read()`] when decoding bytes received over the wire. Each
 /// variant contains a message from [`msgs`] or otherwise the message type if unknown.
 #[allow(missing_docs)]
 #[derive(Debug)]
-pub enum Message {
+pub(crate) enum Message<T> where T: core::fmt::Debug + Type {
        Init(msgs::Init),
        Error(msgs::ErrorMessage),
        Ping(msgs::Ping),
@@ -57,13 +63,16 @@ pub enum Message {
        GossipTimestampFilter(msgs::GossipTimestampFilter),
        /// A message that could not be decoded because its type is unknown.
        Unknown(MessageType),
+       /// A message that was produced by a [`CustomMessageReader`] and is to be handled by a
+       /// [`::ln::peer_handler::CustomMessageHandler`].
+       Custom(T),
 }
 
 /// A number identifying a message to determine how it is encoded on the wire.
 #[derive(Clone, Copy, Debug)]
 pub struct MessageType(u16);
 
-impl Message {
+impl<T> Message<T> where T: core::fmt::Debug + Type {
        #[allow(dead_code)] // This method is only used in tests
        /// Returns the type that was used to decode the message payload.
        pub fn type_id(&self) -> MessageType {
@@ -97,6 +106,7 @@ impl Message {
                        &Message::ReplyChannelRange(ref msg) => msg.type_id(),
                        &Message::GossipTimestampFilter(ref msg) => msg.type_id(),
                        &Message::Unknown(type_id) => type_id,
+                       &Message::Custom(ref msg) => msg.type_id(),
                }
        }
 }
@@ -120,7 +130,14 @@ impl ::core::fmt::Display for MessageType {
 /// # Errors
 ///
 /// Returns an error if the message payload code not be decoded as the specified type.
-pub fn read<R: io::Read>(buffer: &mut R) -> Result<Message, msgs::DecodeError> {
+pub(crate) fn read<R: io::Read, T, H: core::ops::Deref>(
+       buffer: &mut R,
+       custom_reader: H,
+) -> Result<Message<T>, msgs::DecodeError>
+where
+       T: core::fmt::Debug + Type + Writeable,
+       H::Target: CustomMessageReader<CustomMessage = T>,
+{
        let message_type = <u16 as Readable>::read(buffer)?;
        match message_type {
                msgs::Init::TYPE => {
@@ -208,7 +225,11 @@ pub fn read<R: io::Read>(buffer: &mut R) -> Result<Message, msgs::DecodeError> {
                        Ok(Message::GossipTimestampFilter(Readable::read(buffer)?))
                },
                _ => {
-                       Ok(Message::Unknown(MessageType(message_type)))
+                       if let Some(custom) = custom_reader.read(message_type, buffer)? {
+                               Ok(Message::Custom(custom))
+                       } else {
+                               Ok(Message::Unknown(MessageType(message_type)))
+                       }
                },
        }
 }
@@ -219,22 +240,32 @@ pub fn read<R: io::Read>(buffer: &mut R) -> Result<Message, msgs::DecodeError> {
 /// # Errors
 ///
 /// Returns an I/O error if the write could not be completed.
-pub fn write<M: Encode + Writeable, W: Writer>(message: &M, buffer: &mut W) -> Result<(), io::Error> {
-       M::TYPE.write(buffer)?;
+pub(crate) fn write<M: Type + Writeable, W: Writer>(message: &M, buffer: &mut W) -> Result<(), io::Error> {
+       message.type_id().0.write(buffer)?;
        message.write(buffer)
 }
 
-/// Defines a type-identified encoding for sending messages over the wire.
+mod encode {
+       /// Defines a constant type identifier for reading messages from the wire.
+       pub trait Encode {
+               /// The type identifying the message payload.
+               const TYPE: u16;
+       }
+}
+
+pub(crate) use self::encode::Encode;
+
+/// Defines a type identifier for sending messages over the wire.
 ///
-/// Messages implementing this trait specify a type and must be [`Writeable`] to use with [`write()`].
-pub trait Encode {
-       /// The type identifying the message payload.
-       const TYPE: u16;
+/// Messages implementing this trait specify a type and must be [`Writeable`].
+pub trait Type {
+       /// Returns the type identifying the message payload.
+       fn type_id(&self) -> MessageType;
+}
 
-       /// Returns the type identifying the message payload. Convenience method for accessing
-       /// [`Self::TYPE`].
+impl<T> Type for T where T: Encode {
        fn type_id(&self) -> MessageType {
-               MessageType(Self::TYPE)
+               MessageType(T::TYPE)
        }
 }
 
@@ -355,6 +386,7 @@ mod tests {
        use super::*;
        use prelude::*;
        use core::convert::TryInto;
+       use ::ln::peer_handler::IgnoringMessageHandler;
 
        // Big-endian wire encoding of Pong message (type = 19, byteslen = 2).
        const ENCODED_PONG: [u8; 6] = [0u8, 19u8, 0u8, 2u8, 0u8, 0u8];
@@ -363,35 +395,35 @@ mod tests {
        fn read_empty_buffer() {
                let buffer = [];
                let mut reader = io::Cursor::new(buffer);
-               assert!(read(&mut reader).is_err());
+               assert!(read(&mut reader, &IgnoringMessageHandler{}).is_err());
        }
 
        #[test]
        fn read_incomplete_type() {
                let buffer = &ENCODED_PONG[..1];
                let mut reader = io::Cursor::new(buffer);
-               assert!(read(&mut reader).is_err());
+               assert!(read(&mut reader, &IgnoringMessageHandler{}).is_err());
        }
 
        #[test]
        fn read_empty_payload() {
                let buffer = &ENCODED_PONG[..2];
                let mut reader = io::Cursor::new(buffer);
-               assert!(read(&mut reader).is_err());
+               assert!(read(&mut reader, &IgnoringMessageHandler{}).is_err());
        }
 
        #[test]
        fn read_invalid_message() {
                let buffer = &ENCODED_PONG[..4];
                let mut reader = io::Cursor::new(buffer);
-               assert!(read(&mut reader).is_err());
+               assert!(read(&mut reader, &IgnoringMessageHandler{}).is_err());
        }
 
        #[test]
        fn read_known_message() {
                let buffer = &ENCODED_PONG[..];
                let mut reader = io::Cursor::new(buffer);
-               let message = read(&mut reader).unwrap();
+               let message = read(&mut reader, &IgnoringMessageHandler{}).unwrap();
                match message {
                        Message::Pong(_) => (),
                        _ => panic!("Expected pong message; found message type: {}", message.type_id()),
@@ -402,7 +434,7 @@ mod tests {
        fn read_unknown_message() {
                let buffer = &::core::u16::MAX.to_be_bytes();
                let mut reader = io::Cursor::new(buffer);
-               let message = read(&mut reader).unwrap();
+               let message = read(&mut reader, &IgnoringMessageHandler{}).unwrap();
                match message {
                        Message::Unknown(MessageType(::core::u16::MAX)) => (),
                        _ => panic!("Expected message type {}; found: {}", ::core::u16::MAX, message.type_id()),
@@ -428,7 +460,7 @@ mod tests {
                assert!(write(&message, &mut buffer).is_ok());
 
                let mut reader = io::Cursor::new(buffer);
-               let decoded_message = read(&mut reader).unwrap();
+               let decoded_message = read(&mut reader, &IgnoringMessageHandler{}).unwrap();
                match decoded_message {
                        Message::Pong(msgs::Pong { byteslen: 2u16 }) => (),
                        Message::Pong(msgs::Pong { byteslen }) => {
@@ -440,13 +472,13 @@ mod tests {
 
        #[test]
        fn is_even_message_type() {
-               let message = Message::Unknown(MessageType(42));
+               let message = Message::<()>::Unknown(MessageType(42));
                assert!(message.type_id().is_even());
        }
 
        #[test]
        fn is_odd_message_type() {
-               let message = Message::Unknown(MessageType(43));
+               let message = Message::<()>::Unknown(MessageType(43));
                assert!(!message.type_id().is_even());
        }
 
@@ -466,7 +498,7 @@ mod tests {
 
        fn check_init_msg(buffer: Vec<u8>, expect_unknown: bool) {
                let mut reader = io::Cursor::new(buffer);
-               let decoded_msg = read(&mut reader).unwrap();
+               let decoded_msg = read(&mut reader, &IgnoringMessageHandler{}).unwrap();
                match decoded_msg {
                        Message::Init(msgs::Init { features }) => {
                                assert!(features.supports_variable_length_onion());
@@ -485,7 +517,7 @@ mod tests {
                // Taken from lnd v0.9.0-beta.
                let buffer = vec![1, 1, 91, 164, 146, 213, 213, 165, 21, 227, 102, 33, 105, 179, 214, 21, 221, 175, 228, 93, 57, 177, 191, 127, 107, 229, 31, 50, 21, 81, 179, 71, 39, 18, 35, 2, 89, 224, 110, 123, 66, 39, 148, 246, 177, 85, 12, 19, 70, 226, 173, 132, 156, 26, 122, 146, 71, 213, 247, 48, 93, 190, 185, 177, 12, 172, 0, 3, 2, 162, 161, 94, 103, 195, 37, 2, 37, 242, 97, 140, 2, 111, 69, 85, 39, 118, 30, 221, 99, 254, 120, 49, 103, 22, 170, 227, 111, 172, 164, 160, 49, 68, 138, 116, 16, 22, 206, 107, 51, 153, 255, 97, 108, 105, 99, 101, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 1, 172, 21, 0, 2, 38, 7];
                let mut reader = io::Cursor::new(buffer);
-               let decoded_msg = read(&mut reader).unwrap();
+               let decoded_msg = read(&mut reader, &IgnoringMessageHandler{}).unwrap();
                match decoded_msg {
                        Message::NodeAnnouncement(msgs::NodeAnnouncement { contents: msgs::UnsignedNodeAnnouncement { features, ..}, ..}) => {
                                assert!(features.supports_variable_length_onion());
@@ -502,7 +534,7 @@ mod tests {
                // Taken from lnd v0.9.0-beta.
                let buffer = vec![1, 0, 82, 238, 153, 33, 128, 87, 215, 2, 28, 241, 140, 250, 98, 255, 56, 5, 79, 240, 214, 231, 172, 35, 240, 171, 44, 9, 78, 91, 8, 193, 102, 5, 17, 178, 142, 106, 180, 183, 46, 38, 217, 212, 25, 236, 69, 47, 92, 217, 181, 221, 161, 205, 121, 201, 99, 38, 158, 216, 186, 193, 230, 86, 222, 6, 206, 67, 22, 255, 137, 212, 141, 161, 62, 134, 76, 48, 241, 54, 50, 167, 187, 247, 73, 27, 74, 1, 129, 185, 197, 153, 38, 90, 255, 138, 39, 161, 102, 172, 213, 74, 107, 88, 150, 90, 0, 49, 104, 7, 182, 184, 194, 219, 181, 172, 8, 245, 65, 226, 19, 228, 101, 145, 25, 159, 52, 31, 58, 93, 53, 59, 218, 91, 37, 84, 103, 17, 74, 133, 33, 35, 2, 203, 101, 73, 19, 94, 175, 122, 46, 224, 47, 168, 128, 128, 25, 26, 25, 214, 52, 247, 43, 241, 117, 52, 206, 94, 135, 156, 52, 164, 143, 234, 58, 185, 50, 185, 140, 198, 174, 71, 65, 18, 105, 70, 131, 172, 137, 0, 164, 51, 215, 143, 117, 119, 217, 241, 197, 177, 227, 227, 170, 199, 114, 7, 218, 12, 107, 30, 191, 236, 203, 21, 61, 242, 48, 192, 90, 233, 200, 199, 111, 162, 68, 234, 54, 219, 1, 233, 66, 5, 82, 74, 84, 211, 95, 199, 245, 202, 89, 223, 102, 124, 62, 166, 253, 253, 90, 180, 118, 21, 61, 110, 37, 5, 96, 167, 0, 0, 6, 34, 110, 70, 17, 26, 11, 89, 202, 175, 18, 96, 67, 235, 91, 191, 40, 195, 79, 58, 94, 51, 42, 31, 199, 178, 183, 60, 241, 136, 145, 15, 0, 2, 65, 0, 0, 1, 0, 0, 2, 37, 242, 97, 140, 2, 111, 69, 85, 39, 118, 30, 221, 99, 254, 120, 49, 103, 22, 170, 227, 111, 172, 164, 160, 49, 68, 138, 116, 16, 22, 206, 107, 3, 54, 61, 144, 88, 171, 247, 136, 208, 99, 9, 135, 37, 201, 178, 253, 136, 0, 185, 235, 68, 160, 106, 110, 12, 46, 21, 125, 204, 18, 75, 234, 16, 3, 42, 171, 28, 52, 224, 11, 30, 30, 253, 156, 148, 175, 203, 121, 250, 111, 122, 195, 84, 122, 77, 183, 56, 135, 101, 88, 41, 60, 191, 99, 232, 85, 2, 36, 17, 156, 11, 8, 12, 189, 177, 68, 88, 28, 15, 207, 21, 179, 151, 56, 226, 158, 148, 3, 120, 113, 177, 243, 184, 17, 173, 37, 46, 222, 16];
                let mut reader = io::Cursor::new(buffer);
-               let decoded_msg = read(&mut reader).unwrap();
+               let decoded_msg = read(&mut reader, &IgnoringMessageHandler{}).unwrap();
                match decoded_msg {
                        Message::ChannelAnnouncement(msgs::ChannelAnnouncement { contents: msgs::UnsignedChannelAnnouncement { features, ..}, ..}) => {
                                assert!(!features.requires_unknown_bits());
@@ -510,4 +542,71 @@ mod tests {
                        _ => panic!("Expected node announcement, found message type: {}", decoded_msg.type_id())
                }
        }
+
+       #[derive(Eq, PartialEq, Debug)]
+       struct TestCustomMessage {}
+
+       const CUSTOM_MESSAGE_TYPE : u16 = 9000;
+
+       impl Type for TestCustomMessage {
+               fn type_id(&self) -> MessageType {
+                       MessageType(CUSTOM_MESSAGE_TYPE)
+               }
+       }
+
+       impl Writeable for TestCustomMessage {
+               fn write<W: Writer>(&self, _: &mut W) -> Result<(), io::Error> {
+                       Ok(())
+               }
+       }
+
+       struct TestCustomMessageReader {}
+
+       impl CustomMessageReader for TestCustomMessageReader {
+               type CustomMessage = TestCustomMessage;
+               fn read<R: io::Read>(
+                       &self,
+                       message_type: u16,
+                       _: &mut R
+               ) -> Result<Option<Self::CustomMessage>, msgs::DecodeError> {
+                       if message_type == CUSTOM_MESSAGE_TYPE {
+                               return Ok(Some(TestCustomMessage{}));
+                       }
+
+                       Ok(None)
+               }
+       }
+
+       #[test]
+       fn read_custom_message() {
+               let buffer = vec![35, 40];
+               let mut reader = io::Cursor::new(buffer);
+               let decoded_msg = read(&mut reader, &TestCustomMessageReader{}).unwrap();
+               match decoded_msg {
+                       Message::Custom(custom) => {
+                               assert_eq!(custom.type_id().0, CUSTOM_MESSAGE_TYPE);
+                               assert_eq!(custom, TestCustomMessage {});
+                       },
+                       _ => panic!("Expected custom message, found message type: {}", decoded_msg.type_id()),
+               }
+       }
+
+       #[test]
+       fn read_with_custom_reader_unknown_message_type() {
+               let buffer = vec![35, 42];
+               let mut reader = io::Cursor::new(buffer);
+               let decoded_msg = read(&mut reader, &TestCustomMessageReader{}).unwrap();
+               match decoded_msg {
+                       Message::Unknown(_) => {},
+                       _ => panic!("Expected unknown message, found message type: {}", decoded_msg.type_id()),
+               }
+       }
+
+       #[test]
+       fn custom_reader_unknown_message_type() {
+               let buffer = Vec::new();
+               let mut reader = io::Cursor::new(buffer);
+               let res = TestCustomMessageReader{}.read(CUSTOM_MESSAGE_TYPE + 1, &mut reader).unwrap();
+               assert!(res.is_none());
+       }
 }