1
//! Declare a restricted variant of our message types.
2

            
3
/// Re-export tor_bytes and paste here, so that the macro can use it.
4
pub use {paste, tor_bytes};
5

            
6
/// Declare a restricted version of
7
/// [`AnyRelayMsg`](crate::relaycell::msg::AnyRelayMsg) or
8
/// [`AnyChanMsg`](crate::chancell::msg::AnyChanMsg).
9
///
10
/// Frequently we only want to handle a subset of the possible channel or relay
11
/// commands that we might see.  In those situations, it makes sense to define a
12
/// a message types that will only try to parse the allowable commands.  That way,
13
/// we can avoid exposing any unnecessary parsers to a possible attacker.
14
///
15
/// The restricted message type is an enum, and is declared with a syntax as follows:
16
/// ```
17
/// use tor_cell::{restrict::restricted_msg, relaycell::RelayMsgOuter};
18
///
19
/// restricted_msg! {
20
///     enum OpenStreamMsg : RelayMsg {
21
///         Data,
22
///         Sendme,
23
///         End,
24
///         _ => Unrecognized,
25
///    }
26
/// }
27
///
28
/// type OpenStreamMsgOuter = RelayMsgOuter<OpenStreamMsg>;
29
/// ```
30
///
31
/// Instead of `RelayMsg`, you can say `ChanMsg` to get a restricted channel
32
/// message.
33
///
34
/// Only message variants exposed from the `tor_cell::{chan,relay}cell::msg` are
35
/// supported.
36
///
37
/// You can omit the `_ => Unrecognized` clause at the end.  If you do, then any
38
/// unexpected command types will be treated as a parse error.
39
#[macro_export]
40
macro_rules! restricted_msg {
41
    {
42
        $(#[$meta:meta])*
43
        $(@omit_from $omit_from:literal)?
44
        $v:vis enum $name:ident : RelayMsg {
45
            $($tt:tt)*
46
        }
47
    } => {
48
        $crate::restrict::restricted_msg!{
49
            [
50
            any_type: $crate::relaycell::msg::AnyRelayMsg,
51
            msg_mod: $crate::relaycell::msg,
52
239875
            cmd_type: $crate::relaycell::RelayCmd,
53
239875
            unrecognized: $crate::relaycell::msg::Unrecognized,
54
239875
            body_trait: $crate::relaycell::msg::Body,
55
239875
            msg_trait: $crate::relaycell::RelayMsg,
56
239875
            omit_from: $($omit_from)?
57
239875
            ]
58
239875
            $(#[$meta])*
59
239875
            $v enum $name { $($tt)*}
60
239875
        }
61
239875
    };
62
239875
    {
63
239875
        $(#[$meta:meta])*
64
239875
        $(@omit_from $omit_from:literal)?
65
239875
        $v:vis enum $name:ident : ChanMsg {
66
239875
            $($tt:tt)*
67
239875
        }
68
239875
    } => {
69
239875
        $crate::restrict::restricted_msg!{
70
239875
            [
71
240396
            any_type: $crate::chancell::msg::AnyChanMsg,
72
41909
            msg_mod: $crate::chancell::msg,
73
44502
            cmd_type: $crate::chancell::ChanCmd,
74
8799
            unrecognized: $crate::chancell::msg::Unrecognized,
75
8799
            body_trait: $crate::chancell::msg::Body,
76
8799
            msg_trait: $crate::chancell::ChanMsg,
77
8799
            omit_from: $($omit_from)?
78
8799
            ]
79
8799
            $(#[$meta])*
80
8799
            $v enum $name { $($tt)*}
81
8799
        }
82
8799
    };
83
8799
    {
84
8799
        [
85
8799
          any_type: $any_msg:ty,
86
8799
          msg_mod: $msg_mod:path,
87
8799
          cmd_type: $cmd_type:ty,
88
8799
          unrecognized: $unrec_type:ty,
89
8799
          body_trait: $body_type:ty,
90
8799
          msg_trait: $msg_trait:ty,
91
8799
          omit_from: $($omit_from:literal)?
92
8799
        ]
93
8799
        $(#[$meta:meta])*
94
8799
        $v:vis enum $name:ident {
95
8799
            $(
96
8799
                $(#[$case_meta:meta])*
97
8799
                $([feature=$feat:literal])?
98
8799
                $case:ident
99
8799
            ),*
100
8799
            $(, _ =>
101
8799
                $(#[$unrec_meta:meta])*
102
8799
                $unrecognized:ident )?
103
8799
            $(,)?
104
8799
        }
105
8799
    } => {
106
8799
    $crate::restrict::paste::paste!{
107
8799
        $(#[$meta])*
108
8799
        $v enum $name {
109
8799
            $(
110
8799
                $(#[$case_meta])*
111
8799
                $( #[cfg(feature=$feat)] )?
112
8799
                $case($msg_mod :: $case),
113
8799
            )*
114
8799
            $(
115
8799
                $(#[$unrec_meta])*
116
8799
                $unrecognized($unrec_type)
117
8799
            )?
118
8799
        }
119
8799

            
120
8799
        impl $msg_trait for $name {
121
8799
            fn cmd(&self) -> $cmd_type {
122
8799
                match self {
123
1141
                    $(
124
1141
                        $( #[cfg(feature=$feat)] )?
125
1141
                        Self::$case(_) => $cmd_type:: [<$case:snake:upper>] ,
126
1141
                    )*
127
1141
                    $(
128
1371
                        Self::$unrecognized(u) => u.cmd(),
129
1141
                    )?
130
1141
                }
131
243557
            }
132
1141

            
133
6518
             fn encode_onto<W:>(self, w: &mut W) -> $crate::restrict::tor_bytes::EncodeResult<()>
134
6052
             where
135
6052
                W: $crate::restrict::tor_bytes::Writer + ?Sized
136
6052
             {
137
6052
                match self {
138
1141
                    $(
139
1141
                        $( #[cfg(feature=$feat)] )?
140
1149
                        Self::$case(m) => $body_type::encode_onto(m, w),
141
1141
                    )*
142
1141
                    $(
143
1157
                        Self::$unrecognized(u) => $body_type::encode_onto(u, w),
144
1141
                    )?
145
1141
                }
146
6518
            }
147
1141

            
148
122093
            fn decode_from_reader(cmd: $cmd_type, r: &mut $crate::restrict::tor_bytes::Reader<'_>) -> $crate::restrict::tor_bytes::Result<Self> {
149
122093
                Ok(match cmd {
150
1141
                    $(
151
1141
                        $( #[cfg(feature=$feat)] )?
152
1948
                        $cmd_type:: [<$case:snake:upper>] => Self::$case( <$msg_mod :: $case as $body_type> :: decode_from_reader(r)? ),
153
1141
                    )*
154
1141
                    $(
155
1371
                        _ => Self::$unrecognized($unrec_type::decode_with_cmd(cmd, r)?),
156
1141
                    )?
157
1141
                    #[allow(unreachable_patterns)] // This is unreachable if we had an Unrecognized variant above.
158
1142
                    _ => return Err($crate::restrict::tor_bytes::Error::InvalidMessage(
159
183
                        format!("Unexpected command {} in {}", cmd, stringify!($name)).into()
160
183
                    )),
161
1141
                })
162
122093
            }
163
1141
        }
164
1141

            
165
1141
        $(
166
1141
            #[cfg(feature = $omit_from)]
167
1141
        )?
168
1141
        impl From<$name> for $any_msg {
169
1141
            fn from(msg: $name) -> $any_msg {
170
1141
                match msg {
171
                    $(
172
                        $( #[cfg(feature=$feat)] )?
173
32
                        $name::$case(b) => Self::$case(b),
174
                    )*
175
                    $(
176
                        $name::$unrecognized(u) => $any_msg::Unrecognized(u),
177
                    )?
178
                }
179
1141
            }
180
        }
181

            
182
        $(
183
            #[cfg(feature = $omit_from)]
184
        )?
185
        impl TryFrom<$any_msg> for $name {
186
            type Error = $any_msg;
187
            fn try_from(msg: $any_msg) -> std::result::Result<$name, $any_msg> {
188
                Ok(match msg {
189
                    $(
190
                        $( #[cfg(feature=$feat)] )?
191
                        $any_msg::$case(b) => $name::$case(b),
192
                    )*
193
                    $(
194
                        $any_msg::Unrecognized(u) => Self::$unrecognized(u),
195
                    )?
196
                    #[allow(unreachable_patterns)]
197
                    other => return Err(other),
198
                })
199
            }
200
        }
201
        $(
202
            $( #[cfg(feature=$feat)] )?
203
            impl From<$msg_mod :: $case> for $name {
204
127080
                fn from(m: $msg_mod::$case) -> $name {
205
127080
                    $name :: $case(m)
206
127080
                }
207
            }
208
        )*
209
        $(
210
            impl From<$unrec_type> for $name {
211
188
                fn from (u: $unrec_type) -> $name {
212
188
                    $name::$unrecognized(u)
213
188
                }
214
            }
215
        )?
216
    }
217
    }
218
}
219

            
220
pub use restricted_msg;
221

            
222
#[cfg(test)]
223
mod test {
224
    use super::*;
225
    // Here we do a couple of other variations of the example in the doctest, to
226
    // make sure they work.
227

            
228
    // As in the doctest, but no "unrecognized" variant.
229
    restricted_msg! {
230
        enum StrictOpenStreamMsg : RelayMsg {
231
            Data,
232
            Sendme,
233
            End,
234
       }
235
    }
236

            
237
    // Try it with chanmsg.
238
    restricted_msg! {
239
        enum CircuitBuildReply : ChanMsg {
240
            Created,
241
            Created2,
242
            CreatedFast,
243
            Destroy,
244
            _ => Unrecognized,
245
       }
246
    }
247

            
248
    // As above, but no "unrecognized" variant.
249
    restricted_msg! {
250
        enum StrictCircuitBuildReply : ChanMsg {
251
            Created,
252
            Created2,
253
            CreatedFast,
254
            Destroy,
255
       }
256
    }
257
}