Skip to content

Commit

Permalink
f - Refactor Encode and TypedMessage
Browse files Browse the repository at this point in the history
  • Loading branch information
jkczyz authored and Tibo-lg committed Aug 19, 2021
1 parent 3b48b4c commit b6f743f
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 34 deletions.
8 changes: 4 additions & 4 deletions lightning/src/ln/peer_handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use ln::channelmanager::{SimpleArcChannelManager, SimpleRefChannelManager};
use util::ser::{VecWriter, Writeable, Writer};
use ln::peer_channel_encryptor::{PeerChannelEncryptor,NextNoiseStep};
use ln::wire;
use ln::wire::TypedMessage;
use ln::wire::MessageType;
use util::byte_utils;
use util::events::{MessageSendEvent, MessageSendEventsProvider};
use util::logger::Logger;
Expand Down Expand Up @@ -85,8 +85,8 @@ pub struct IgnoringCustomMessageHandler{}

type DummyCustomType = ();

impl TypedMessage for DummyCustomType {
fn msg_type(&self) -> u16 {
impl wire::Type for DummyCustomType {
fn type_id(&self) -> MessageType {
// We should never call this for `DummyCustomType`
unreachable!();
}
Expand Down Expand Up @@ -728,7 +728,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
}

/// Append a message to a peer's pending outbound/write buffer, and update the map of peers needing sends accordingly.
fn enqueue_message<M: TypedMessage + Writeable + Debug>(&self, peer: &mut Peer, message: &M) {
fn enqueue_message<M: wire::Type + Writeable + Debug>(&self, peer: &mut Peer, message: &M) {
let mut buffer = VecWriter(Vec::new());
wire::write(message, &mut buffer).unwrap(); // crash if the write failed
let encoded_message = buffer.0;
Expand Down
53 changes: 23 additions & 30 deletions lightning/src/ln/wire.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
//! 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
//! Messages implementing the [`Type`] trait define a message type and can be sent over the wire
//! using [`write()`].
//!
//! [BOLT #1]: https://github.com/lightningnetwork/lightning-rfc/blob/master/01-messaging.md
Expand All @@ -26,7 +26,7 @@ use util::ser::{Readable, Writeable, Writer};
/// decoders.
pub trait CustomMessageReader {
/// The type of the message decoded by the implementation.
type CustomMessage : core::fmt::Debug + TypedMessage + Writeable;
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
Expand All @@ -38,7 +38,7 @@ pub trait CustomMessageReader {
/// variant contains a message from [`msgs`] or otherwise the message type if unknown.
#[allow(missing_docs)]
#[derive(Debug)]
pub(crate) enum Message<T> where T: core::fmt::Debug + TypedMessage {
pub(crate) enum Message<T> where T: core::fmt::Debug + Type {
Init(msgs::Init),
Error(msgs::ErrorMessage),
Ping(msgs::Ping),
Expand Down Expand Up @@ -76,7 +76,7 @@ pub(crate) enum Message<T> where T: core::fmt::Debug + TypedMessage {
#[derive(Clone, Copy, Debug)]
pub struct MessageType(u16);

impl<T> Message<T> where T: core::fmt::Debug + TypedMessage {
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 {
Expand Down Expand Up @@ -110,7 +110,7 @@ impl<T> Message<T> where T: core::fmt::Debug + TypedMessage {
&Message::ReplyChannelRange(ref msg) => msg.type_id(),
&Message::GossipTimestampFilter(ref msg) => msg.type_id(),
&Message::Unknown(type_id) => type_id,
&Message::Custom(ref msg) => MessageType(msg.msg_type()),
&Message::Custom(ref msg) => msg.type_id(),
}
}
}
Expand Down Expand Up @@ -139,7 +139,7 @@ pub(crate) fn read<R: io::Read, T, H: core::ops::Deref>(
custom_reader: &H
) -> Result<Message<T>, msgs::DecodeError>
where
T: core::fmt::Debug + TypedMessage + Writeable,
T: core::fmt::Debug + Type + Writeable,
H::Target: CustomMessageReader<CustomMessage = T>
{
let message_type = <u16 as Readable>::read(buffer)?;
Expand Down Expand Up @@ -244,39 +244,32 @@ where
/// # Errors
///
/// Returns an I/O error if the write could not be completed.
pub fn write<M: TypedMessage + Writeable, W: Writer>(message: &M, buffer: &mut W) -> Result<(), io::Error> {
message.msg_type().write(buffer)?;
pub fn write<M: Type + Writeable, W: Writer>(message: &M, buffer: &mut W) -> Result<(), io::Error> {
message.type_id().0.write(buffer)?;
message.write(buffer)
}

pub(crate) mod encode {
use super::*;
/// Defines a type-identified encoding for sending messages over the wire.
///
/// Messages implementing this trait specify a type and must be [`Writeable`] to use with [`write()`].
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;

/// Returns the type identifying the message payload. Convenience method for accessing
/// [`Self::TYPE`].
fn type_id(&self) -> MessageType {
MessageType(Self::TYPE)
}
}
}

pub(crate) use self::encode::Encode;

/// A message that has an associated type id.
pub trait TypedMessage {
/// The type id for the implementing message.
fn msg_type(&self) -> u16;
/// 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 Type {
/// Returns the type identifying the message payload.
fn type_id(&self) -> MessageType;
}

impl<T> TypedMessage for T where T: Encode {
fn msg_type(&self) -> u16 {
T::TYPE
impl<T> Type for T where T: Encode {
fn type_id(&self) -> MessageType {
MessageType(T::TYPE)
}
}

Expand Down Expand Up @@ -559,9 +552,9 @@ mod tests {

const CUSTOM_MESSAGE_TYPE : u16 = 9000;

impl TypedMessage for TestCustomMessage {
fn msg_type(&self) -> u16 {
CUSTOM_MESSAGE_TYPE
impl Type for TestCustomMessage {
fn type_id(&self) -> MessageType {
MessageType(CUSTOM_MESSAGE_TYPE)
}
}

Expand Down Expand Up @@ -598,7 +591,7 @@ mod tests {
let decoded_msg = read(&mut reader, &TestCustomMessageReader{}).unwrap();
match decoded_msg {
Message::Custom(custom) => {
assert_eq!(custom.msg_type(), CUSTOM_MESSAGE_TYPE);
assert_eq!(custom.type_id().0, CUSTOM_MESSAGE_TYPE);
assert_eq!(custom, TestCustomMessage {});
},
_ => panic!("Expected custom message, found message type: {}", decoded_msg.type_id()),
Expand Down

0 comments on commit b6f743f

Please sign in to comment.