From 3f1457ea046b724a1df33dee6c91c1e51a773dd5 Mon Sep 17 00:00:00 2001
From: Nick Mathewson <nickm@torproject.org>
Date: Mon, 13 Feb 2023 15:53:59 -0500
Subject: [PATCH] tor-cell: Implement {Relay,Chan}Msg for every body type

This will make it ergonomic to decode a single body type without
having to declare a variant that accepts only a single message.
---
 crates/tor-cell/src/chancell/msg.rs  | 48 ++++++++++++++++++++++++++++
 crates/tor-cell/src/relaycell/msg.rs | 46 ++++++++++++++++++++++++++
 crates/tor-cell/src/restrict.rs      |  8 ++---
 3 files changed, 97 insertions(+), 5 deletions(-)

diff --git a/crates/tor-cell/src/chancell/msg.rs b/crates/tor-cell/src/chancell/msg.rs
index d428a8e62..49571e99e 100644
--- a/crates/tor-cell/src/chancell/msg.rs
+++ b/crates/tor-cell/src/chancell/msg.rs
@@ -1170,6 +1170,54 @@ msg_into_cell!(AuthChallenge);
 msg_into_cell!(Authenticate);
 msg_into_cell!(Authorize);
 
+/// Helper: declare a ChanMsg implementation for a message type that has a
+/// fixed command.
+//
+// TODO: It might be better to merge Body with ChanMsg, but that is complex,
+// since their needs are _slightly_ different.
+macro_rules! msg_impl_chanmsg {
+    ($($body:ident,)*) =>
+    {paste::paste!{
+       $(impl crate::chancell::ChanMsg for $body {
+            fn cmd(&self) -> crate::chancell::ChanCmd { crate::chancell::ChanCmd::[< $body:snake:upper >] }
+            fn encode_onto<W: tor_bytes::Writer + ?Sized>(self, w: &mut W) -> tor_bytes::EncodeResult<()> {
+                crate::chancell::msg::Body::encode_onto(self, w)
+            }
+            fn decode_from_reader(cmd: ChanCmd, r: &mut tor_bytes::Reader<'_>) -> tor_bytes::Result<Self> {
+                if cmd != crate::chancell::ChanCmd::[< $body:snake:upper >] {
+                    return Err(tor_bytes::Error::InvalidMessage(
+                        format!("Expected {} command; got {cmd}", stringify!([< $body:snake:upper >])).into()
+                    ));
+                }
+                crate::chancell::msg::Body::decode_from_reader(r)
+            }
+        })*
+    }}
+}
+
+// We implement ChanMsg for every body type, so that you can write code that does
+// e.g. ChanCell<Relay>.
+msg_impl_chanmsg!(
+    Padding,
+    Vpadding,
+    Create,
+    CreateFast,
+    Create2,
+    Created,
+    CreatedFast,
+    Created2,
+    Relay,
+    RelayEarly,
+    Destroy,
+    Netinfo,
+    Versions,
+    PaddingNegotiate,
+    Certs,
+    AuthChallenge,
+    Authenticate,
+    Authorize,
+);
+
 #[cfg(test)]
 mod test {
     // @@ begin test lint list maintained by maint/add_warning @@
diff --git a/crates/tor-cell/src/relaycell/msg.rs b/crates/tor-cell/src/relaycell/msg.rs
index a33afbae4..32beaec34 100644
--- a/crates/tor-cell/src/relaycell/msg.rs
+++ b/crates/tor-cell/src/relaycell/msg.rs
@@ -1175,3 +1175,49 @@ empty_body! {
     /// Opens a new stream on a directory cache.
     pub struct BeginDir {}
 }
+
+/// Helper: declare a RelayMsg implementation for a message type that has a
+/// fixed command.
+//
+// TODO: It might be better to merge Body with RelayMsg, but that is complex,
+// since their needs are _slightly_ different.
+macro_rules! msg_impl_relaymsg {
+    ($($body:ident),* $(,)?) =>
+    {paste::paste!{
+       $(impl crate::relaycell::RelayMsg for $body {
+            fn cmd(&self) -> crate::relaycell::RelayCmd { crate::relaycell::RelayCmd::[< $body:snake:upper >] }
+            fn encode_onto<W: tor_bytes::Writer + ?Sized>(self, w: &mut W) -> tor_bytes::EncodeResult<()> {
+                crate::relaycell::msg::Body::encode_onto(self, w)
+            }
+            fn decode_from_reader(cmd: RelayCmd, r: &mut tor_bytes::Reader<'_>) -> tor_bytes::Result<Self> {
+                if cmd != crate::relaycell::RelayCmd::[< $body:snake:upper >] {
+                    return Err(tor_bytes::Error::InvalidMessage(
+                        format!("Expected {} command; got {cmd}", stringify!([< $body:snake:upper >])).into()
+                    ));
+                }
+                crate::relaycell::msg::Body::decode_from_reader(r)
+            }
+        })*
+    }}
+}
+
+msg_impl_relaymsg!(
+    Begin, Data, End, Connected, Sendme, Extend, Extended, Extend2, Extended2, Truncate, Truncated,
+    Drop, Resolve, Resolved, BeginDir,
+);
+
+#[cfg(feature = "experimental-udp")]
+msg_impl_relaymsg!(ConnectUdp, ConnectedUdp, Datagram);
+
+#[cfg(feature = "onion-service")]
+msg_impl_relaymsg!(
+    EstablishIntro,
+    EstablishRendezvous,
+    Introduce1,
+    Introduce2,
+    Rendezvous1,
+    Rendezvous2,
+    IntroEstablished,
+    RendezvousEstablished,
+    IntroduceAck,
+);
diff --git a/crates/tor-cell/src/restrict.rs b/crates/tor-cell/src/restrict.rs
index c975080e4..f39f9b05c 100644
--- a/crates/tor-cell/src/restrict.rs
+++ b/crates/tor-cell/src/restrict.rs
@@ -134,24 +134,22 @@ macro_rules! restricted_msg {
              where
                 W: $crate::restrict::tor_bytes::Writer + ?Sized
              {
-                use $body_type;
                 match self {
                     $(
                         $( #[cfg(feature=$feat)] )?
-                        Self::$case(m) => m.encode_onto(w),
+                        Self::$case(m) => $body_type::encode_onto(m, w),
                     )*
                     $(
-                        Self::$unrecognized(u) => u.encode_onto(w),
+                        Self::$unrecognized(u) => $body_type::encode_onto(u, w),
                     )?
                 }
             }
 
             fn decode_from_reader(cmd: $cmd_type, r: &mut $crate::restrict::tor_bytes::Reader<'_>) -> $crate::restrict::tor_bytes::Result<Self> {
-                use $body_type;
                 Ok(match cmd {
                     $(
                         $( #[cfg(feature=$feat)] )?
-                        $cmd_type:: [<$case:snake:upper>] => Self::$case( $msg_mod :: $case :: decode_from_reader(r)? ),
+                        $cmd_type:: [<$case:snake:upper>] => Self::$case( <$msg_mod :: $case as $body_type> :: decode_from_reader(r)? ),
                     )*
                     $(
                         _ => Self::$unrecognized($unrec_type::decode_with_cmd(cmd, r)?),