1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
//! Implement what components may be parsed from the wire

use bytes::{Buf, BytesMut};
use enum_dispatch::enum_dispatch;

use crate::actions::{Abort, Continue, Discard, Quit, QuitNc, Reject, Replycode, Skip, Tempfail};

use crate::{
    error::STAGE_DECODING, AddHeader, AddRecipient, ChangeHeader, DeleteRecipient, InsertHeader,
    InvalidData, NotEnoughData, ProtocolError, Quarantine, ReplaceBody,
};

use super::commands::Connect;
use super::commands::Helo;
use super::commands::Macro;
use super::commands::Recipient;
use super::commands::Unknown;
use super::commands::{Body, EndOfBody};
use super::commands::{Data, Mail};
use super::commands::{EndOfHeader, Header};
use super::optneg::OptNeg;

/// Parse something 'from the wire'.
pub(crate) trait Parsable: Sized {
    /// The unique id code for this item
    const CODE: u8;

    /// Parse a `Self` from the given `BytesMut` buffer.
    ///
    /// # Errors
    /// This can fail to parse, returning a [`ProtocolError`].
    fn parse(buffer: BytesMut) -> Result<Self, ProtocolError>;
}

macro_rules! parse_command {
    ($container_name:ident, $($variant:ident),+$(,)?) => {
        /// See the contained variants for more.
        #[allow(missing_docs)]
        #[enum_dispatch]
        #[derive(Debug, Clone)]
        pub enum $container_name {
            $($variant($variant),)+
        }

        impl $container_name {
            /// Parse a bytes buffer into this structured data
            /// 
            /// # Errors
            /// This fn may return errors if the received data did not match
            /// valid data for this command.
            pub fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
                if buffer.is_empty() {
                    return Err(NotEnoughData::new(
                        STAGE_DECODING,
                        "Command",
                        "code missing to detect which command it is",
                        1,
                        0,
                        buffer
                    ).into());
                }
                let code = buffer.get_u8();
                match code {
                    $($variant::CODE => Ok($variant::parse(buffer)?.into()),)+
                    _ => {
                        Err(InvalidData{msg: "Unknown command sent with code", offending_bytes: BytesMut::from_iter(&[code])}.into())
                    }
                }
            }
        }

        $(impl From<$variant> for $container_name {
            fn from(value: $variant) -> Self {
                Self::$variant(value)
            }
        })+

    }
}

// Parse a command sent by the client.
parse_command!(
    // The name of this enum
    ClientCommand,
    // Include actions
    Abort,
    // Milter Control
    OptNeg,
    Quit,
    QuitNc,
    // Special Info
    Macro,
    Unknown,
    // SMTP opening
    Connect,
    Helo,
    // Header
    Mail,
    Recipient,
    Header,
    EndOfHeader,
    // Body
    Data,
    Body,
    EndOfBody,
);

// Parse a command sent by the server.
parse_command!(
    // The name of this enum
    ServerCommand,
    // Option negotiation
    OptNeg,
    // The actions
    Abort,
    Continue,
    Discard,
    Reject,
    Tempfail,
    Skip,
    Replycode,
    // Modifications
    AddRecipient,
    DeleteRecipient,
    ReplaceBody,
    AddHeader,
    InsertHeader,
    ChangeHeader,
    Quarantine,
);

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use bytes::BytesMut;

    use super::*;

    #[test]
    fn test_create_abort() {
        let data = vec![b'A'];

        let command =
            ClientCommand::parse(BytesMut::from_iter(data)).expect("Failed parsing abort data");

        assert_matches!(command, ClientCommand::Abort(_));
    }

    #[test]
    fn test_create_optneg() {
        let data = vec![b'O', 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 0];

        let command =
            ClientCommand::parse(BytesMut::from_iter(data)).expect("Failed parsing optneg data");

        assert_matches!(command, ClientCommand::OptNeg(o) if o.version == 6);
    }
}