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
            cmd_type: $crate::relaycell::RelayCmd,
53
            unrecognized: $crate::relaycell::msg::Unrecognized,
54
            body_trait: $crate::relaycell::msg::Body,
55
            msg_trait: $crate::relaycell::RelayMsg,
56
            omit_from: $($omit_from)?
57
            ]
58
            $(#[$meta])*
59
            $v enum $name { $($tt)*}
60
        }
61
    };
62
    {
63
        $(#[$meta:meta])*
64
        $(@omit_from $omit_from:literal)?
65
        $v:vis enum $name:ident : ChanMsg {
66
            $($tt:tt)*
67
        }
68
    } => {
69
        $crate::restrict::restricted_msg!{
70
            [
71
            any_type: $crate::chancell::msg::AnyChanMsg,
72
            msg_mod: $crate::chancell::msg,
73
            cmd_type: $crate::chancell::ChanCmd,
74
            unrecognized: $crate::chancell::msg::Unrecognized,
75
            body_trait: $crate::chancell::msg::Body,
76
            msg_trait: $crate::chancell::ChanMsg,
77
            omit_from: $($omit_from)?
78
            ]
79
            $(#[$meta])*
80
            $v enum $name { $($tt)*}
81
        }
82
    };
83
    {
84
        [
85
          any_type: $any_msg:ty,
86
          msg_mod: $msg_mod:path,
87
          cmd_type: $cmd_type:ty,
88
          unrecognized: $unrec_type:ty,
89
          body_trait: $body_type:ty,
90
          msg_trait: $msg_trait:ty,
91
          omit_from: $($omit_from:literal)?
92
        ]
93
        $(#[$meta:meta])*
94
        $v:vis enum $name:ident {
95
            $(
96
                $(#[$case_meta:meta])*
97
                $([feature=$feat:literal])?
98
                $case:ident
99
            ),*
100
            $(, _ =>
101
                $(#[$unrec_meta:meta])*
102
                $unrecognized:ident )?
103
            $(,)?
104
        }
105
    } => {
106
    $crate::restrict::paste::paste!{
107
        $(#[$meta])*
108
        $v enum $name {
109
            $(
110
                $(#[$case_meta])*
111
                $( #[cfg(feature=$feat)] )?
112
                $case($msg_mod :: $case),
113
            )*
114
            $(
115
                $(#[$unrec_meta])*
116
                $unrecognized($unrec_type)
117
            )?
118
        }
119

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

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

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

            
165
151
        #[allow(unexpected_cfgs)]
166
151
        const _: () = {
167
151
            $(
168
151
                #[cfg(feature = $omit_from)]
169
151
            )?
170
151
            impl From<$name> for $any_msg {
171
340
                fn from(msg: $name) -> $any_msg {
172
282
                    match msg {
173
2
                        $(
174
2
                            $( #[cfg(feature=$feat)] )?
175
2
                            $name::$case(b) => Self::$case(b),
176
2
                        )*
177
2
                        $(
178
2
                            $name::$unrecognized(u) => $any_msg::Unrecognized(u),
179
2
                        )?
180
2
                    }
181
142
                }
182
142
            }
183
142
        };
184
142

            
185
142
        #[allow(unexpected_cfgs)]
186
142
        const _: () = {
187
142
            $(
188
142
                #[cfg(feature = $omit_from)]
189
142
            )?
190
142
            impl TryFrom<$any_msg> for $name {
191
142
                type Error = $any_msg;
192
142
                fn try_from(msg: $any_msg) -> std::result::Result<$name, $any_msg> {
193
                    Ok(match msg {
194
                        $(
195
                            $( #[cfg(feature=$feat)] )?
196
                            $any_msg::$case(b) => $name::$case(b),
197
                        )*
198
                        $(
199
                            $any_msg::Unrecognized(u) => Self::$unrecognized(u),
200
                        )?
201
                        #[allow(unreachable_patterns)]
202
                        other => return Err(other),
203
                    })
204
                }
205
            }
206
        };
207

            
208
        $(
209
            $( #[cfg(feature=$feat)] )?
210
            impl From<$msg_mod :: $case> for $name {
211
145954
                fn from(m: $msg_mod::$case) -> $name {
212
145954
                    $name :: $case(m)
213
145954
                }
214
            }
215
        )*
216
        $(
217
            impl From<$unrec_type> for $name {
218
196
                fn from (u: $unrec_type) -> $name {
219
196
                    $name::$unrecognized(u)
220
196
                }
221
            }
222
        )?
223
    }
224
    }
225
}
226

            
227
pub use restricted_msg;
228

            
229
#[cfg(test)]
230
mod test {
231
    use super::*;
232
    // Here we do a couple of other variations of the example in the doctest, to
233
    // make sure they work.
234

            
235
    // As in the doctest, but no "unrecognized" variant.
236
    restricted_msg! {
237
        enum StrictOpenStreamMsg : RelayMsg {
238
            Data,
239
            Sendme,
240
            End,
241
       }
242
    }
243

            
244
    // Try it with chanmsg.
245
    restricted_msg! {
246
        enum CircuitBuildReply : ChanMsg {
247
            Created,
248
            Created2,
249
            CreatedFast,
250
            Destroy,
251
            _ => Unrecognized,
252
       }
253
    }
254

            
255
    // As above, but no "unrecognized" variant.
256
    restricted_msg! {
257
        enum StrictCircuitBuildReply : ChanMsg {
258
            Created,
259
            Created2,
260
            CreatedFast,
261
            Destroy,
262
       }
263
    }
264
}