miltr_common/
decoding.rs

1//! Implement what components may be parsed from the wire
2
3use bytes::{Buf, BytesMut};
4use enum_dispatch::enum_dispatch;
5
6use crate::actions::{Abort, Continue, Discard, Quit, QuitNc, Reject, Replycode, Skip, Tempfail};
7
8use crate::{
9    error::STAGE_DECODING, AddHeader, AddRecipient, ChangeHeader, DeleteRecipient, InsertHeader,
10    InvalidData, NotEnoughData, ProtocolError, Quarantine, ReplaceBody,
11};
12
13use super::commands::Connect;
14use super::commands::Helo;
15use super::commands::Macro;
16use super::commands::Recipient;
17use super::commands::Unknown;
18use super::commands::{Body, EndOfBody};
19use super::commands::{Data, Mail};
20use super::commands::{EndOfHeader, Header};
21use super::optneg::OptNeg;
22
23/// Parse something 'from the wire'.
24pub(crate) trait Parsable: Sized {
25    /// The unique id code for this item
26    const CODE: u8;
27
28    /// Parse a `Self` from the given `BytesMut` buffer.
29    ///
30    /// # Errors
31    /// This can fail to parse, returning a [`ProtocolError`].
32    fn parse(buffer: BytesMut) -> Result<Self, ProtocolError>;
33}
34
35macro_rules! parse_command {
36    ($container_name:ident, $($variant:ident),+$(,)?) => {
37        /// See the contained variants for more.
38        #[allow(missing_docs)]
39        #[cfg_attr(feature = "tracing", derive(strum::Display))]
40        #[enum_dispatch]
41        #[derive(Debug, Clone)]
42        pub enum $container_name {
43            $($variant($variant),)+
44        }
45
46        impl $container_name {
47            /// Parse a bytes buffer into this structured data
48            ///
49            /// # Errors
50            /// This fn may return errors if the received data did not match
51            /// valid data for this command.
52            pub fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
53                if buffer.is_empty() {
54                    return Err(NotEnoughData::new(
55                        STAGE_DECODING,
56                        "Command",
57                        "code missing to detect which command it is",
58                        1,
59                        0,
60                        buffer
61                    ).into());
62                }
63                let code = buffer.get_u8();
64                match code {
65                    $($variant::CODE => Ok($variant::parse(buffer)?.into()),)+
66                    _ => {
67                        Err(InvalidData{msg: "Unknown command sent with code", offending_bytes: BytesMut::from_iter(&[code])}.into())
68                    }
69                }
70            }
71        }
72
73        $(impl From<$variant> for $container_name {
74            fn from(value: $variant) -> Self {
75                Self::$variant(value)
76            }
77        })+
78
79    }
80}
81
82// Parse a command sent by the client.
83parse_command!(
84    // The name of this enum
85    ClientCommand,
86    // Include actions
87    Abort,
88    // Milter Control
89    OptNeg,
90    Quit,
91    QuitNc,
92    // Special Info
93    Macro,
94    Unknown,
95    // SMTP opening
96    Connect,
97    Helo,
98    // Header
99    Mail,
100    Recipient,
101    Header,
102    EndOfHeader,
103    // Body
104    Data,
105    Body,
106    EndOfBody,
107);
108
109// Parse a command sent by the server.
110parse_command!(
111    // The name of this enum
112    ServerCommand,
113    // Option negotiation
114    OptNeg,
115    // The actions
116    Abort,
117    Continue,
118    Discard,
119    Reject,
120    Tempfail,
121    Skip,
122    Replycode,
123    // Modifications
124    AddRecipient,
125    DeleteRecipient,
126    ReplaceBody,
127    AddHeader,
128    InsertHeader,
129    ChangeHeader,
130    Quarantine,
131);
132
133#[cfg(test)]
134mod tests {
135    use assert_matches::assert_matches;
136    use bytes::BytesMut;
137
138    use super::*;
139
140    #[test]
141    fn test_create_abort() {
142        let data = vec![b'A'];
143
144        let command =
145            ClientCommand::parse(BytesMut::from_iter(data)).expect("Failed parsing abort data");
146
147        assert_matches!(command, ClientCommand::Abort(_));
148    }
149
150    #[test]
151    fn test_create_optneg() {
152        let data = vec![b'O', 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 0];
153
154        let command =
155            ClientCommand::parse(BytesMut::from_iter(data)).expect("Failed parsing optneg data");
156
157        assert_matches!(command, ClientCommand::OptNeg(o) if o.version == 6);
158    }
159}