1 files changed, 72 insertions, 31 deletions

diff --git a/crates/tor-proto/src/client/circuit/celltypes.rs b/crates/tor-proto/src/client/circuit/celltypes.rs
index 5cd553741..aab47f215 100644
--- a/crates/tor-proto/src/client/circuit/celltypes.rs
+++ b/crates/tor-proto/src/client/circuit/celltypes.rs
@@ -5,7 +5,7 @@
 //! so that Rust's typesafety can help enforce protocol properties.
 
 use crate::{Error, Result};
-use derive_deftly::Deftly;
+use derive_deftly::{Deftly, define_derive_deftly};
 use std::fmt::{self, Display};
 use tor_cell::chancell::{
     ChanMsg,
@@ -13,12 +13,39 @@ use tor_cell::chancell::{
 };
 use tor_memquota::derive_deftly_template_HasMemoryCost;
 
+define_derive_deftly! {
+    /// Derives a `TryFrom<AnyChanMsg>` implementation for enums
+    /// that represent restricted subsets of ChanMsgs
+    ///
+    /// # Limitations
+    ///
+    /// The variants of the enum this is derived for *must* be a
+    /// subset of the variants of [`AnyChanMsg`].
+    RestrictedChanMsgSet:
+
+    impl TryFrom<AnyChanMsg> for $ttype {
+        type Error = crate::Error;
+
+        fn try_from(m: AnyChanMsg) -> Result<$ttype> {
+            match m {
+                $( AnyChanMsg::$vname(m) => Ok($ttype::$vname(m)), )
+                _ => Err(Error::ChanProto(format!(
+                    "Got a {} {}",
+                    m.cmd(), ${tmeta(usage) as str},
+                ))),
+            }
+        }
+    }
+}
+
 /// A subclass of ChanMsg that can arrive in response to a CREATE* cell
 /// that we send.
 #[cfg_attr(docsrs, doc(cfg(feature = "testing")))]
-#[derive(Debug)]
+#[derive(Debug, Deftly)]
 #[allow(unreachable_pub)] // Only `pub` with feature `testing`; otherwise, visible in crate
 #[allow(clippy::exhaustive_enums)]
+#[derive_deftly(RestrictedChanMsgSet)]
+#[deftly(usage = "in response to circuit creation")]
 pub enum CreateResponse {
     /// Destroy cell: the CREATE failed.
     Destroy(chanmsg::Destroy),
@@ -39,27 +66,13 @@ impl Display for CreateResponse {
     }
 }
 
-impl TryFrom<AnyChanMsg> for CreateResponse {
-    type Error = crate::Error;
-
-    fn try_from(m: AnyChanMsg) -> Result<CreateResponse> {
-        match m {
-            AnyChanMsg::Destroy(m) => Ok(CreateResponse::Destroy(m)),
-            AnyChanMsg::CreatedFast(m) => Ok(CreateResponse::CreatedFast(m)),
-            AnyChanMsg::Created2(m) => Ok(CreateResponse::Created2(m)),
-            _ => Err(Error::ChanProto(format!(
-                "Got a {} in response to circuit creation",
-                m.cmd()
-            ))),
-        }
-    }
-}
-
 /// A subclass of ChanMsg that can correctly arrive on a live client
 /// circuit (one where a CREATED* has been received).
 #[derive(Debug, Deftly)]
 #[allow(unreachable_pub)] // Only `pub` with feature `testing`; otherwise, visible in crate
 #[derive_deftly(HasMemoryCost)]
+#[derive_deftly(RestrictedChanMsgSet)]
+#[deftly(usage = "on an open client circuit")]
 pub enum ClientCircChanMsg {
     /// A relay cell telling us some kind of remote command from some
     /// party on the circuit.
@@ -69,19 +82,24 @@ pub enum ClientCircChanMsg {
     // Note: RelayEarly is not valid for clients!
 }
 
-impl TryFrom<AnyChanMsg> for ClientCircChanMsg {
-    type Error = crate::Error;
-
-    fn try_from(m: AnyChanMsg) -> Result<ClientCircChanMsg> {
-        match m {
-            AnyChanMsg::Destroy(m) => Ok(ClientCircChanMsg::Destroy(m)),
-            AnyChanMsg::Relay(m) => Ok(ClientCircChanMsg::Relay(m)),
-            _ => Err(Error::ChanProto(format!(
-                "Got a {} cell on an open circuit",
-                m.cmd()
-            ))),
-        }
-    }
+/// A subclass of ChanMsg that can correctly arrive on a live relay
+/// circuit (one where a CREATE* has been received).
+#[derive(Debug, Deftly)]
+#[derive_deftly(HasMemoryCost)]
+#[derive_deftly(RestrictedChanMsgSet)]
+#[deftly(usage = "on an open relay circuit")]
+#[cfg(feature = "relay")]
+#[cfg_attr(not(test), expect(unused))] // TODO(relay)
+pub(crate) enum RelayCircChanMsg {
+    /// A relay cell telling us some kind of remote command from some
+    /// party on the circuit.
+    Relay(chanmsg::Relay),
+    /// A relay early cell that is allowed to contain a CREATE message.
+    RelayEarly(chanmsg::RelayEarly),
+    /// A cell telling us to destroy the circuit.
+    Destroy(chanmsg::Destroy),
+    /// A cell telling us to enable/disable channel padding.
+    PaddingNegotiate(chanmsg::PaddingNegotiate),
 }
 
 #[cfg(test)]
@@ -137,4 +155,27 @@ mod test {
         ));
         bad(msg::Versions::new([1, 2, 3]).unwrap().into());
     }
+
+    #[test]
+    #[cfg(feature = "relay")]
+    fn relay_circ_chan_msg() {
+        use tor_cell::chancell::msg::{self, AnyChanMsg};
+        fn good(m: AnyChanMsg) {
+            assert!(RelayCircChanMsg::try_from(m).is_ok());
+        }
+        fn bad(m: AnyChanMsg) {
+            assert!(RelayCircChanMsg::try_from(m).is_err());
+        }
+
+        good(msg::Destroy::new(2.into()).into());
+        bad(msg::CreatedFast::new(&b"The great globular mass"[..]).into());
+        bad(msg::Created2::new(&b"of protoplasmic slush"[..]).into());
+        good(msg::Relay::new(&b"undulated slightly,"[..]).into());
+        good(msg::AnyChanMsg::RelayEarly(
+            msg::Relay::new(&b"as if aware of him"[..]).into(),
+        ));
+        bad(msg::Versions::new([1, 2, 3]).unwrap().into());
+        good(msg::PaddingNegotiate::start_default().into());
+        good(msg::RelayEarly::from(msg::Relay::new(b"snail-like unipedular organism")).into());
+    }
 }