indymilter 0.3.0

// indymilter – asynchronous milter library
// Copyright © 2021–2024 David Bürgin <dbuergin@gluet.ch>
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or (at your option) any later
// version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
// details.
//
// You should have received a copy of the GNU General Public License along with
// this program. If not, see <https://www.gnu.org/licenses/>.

//! Milter replies.

use crate::{
    macros::MacroStage,
    message::{Byte, Message, TryFromByteError, Version},
    proto_util::{Actions, ProtoOpts},
};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::{
    collections::HashMap,
    error::Error,
    ffi::CString,
    fmt::{self, Display, Formatter},
};

/// The kind of a reply.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum ReplyKind {
    /// The `+` reply.
    AddRcpt,
    /// The `-` reply.
    DeleteRcpt,
    /// The `2` reply.
    AddRcptExt,
    /// The `O` reply.
    OptNeg,
    /// The `a` reply.
    Accept,
    /// The `b` reply.
    ReplaceBody,
    /// The `c` reply.
    Continue,
    /// The `d` reply.
    Discard,
    /// The `e` reply.
    ChangeSender,
    /// The `h` reply.
    AddHeader,
    /// The `i` reply.
    InsertHeader,
    /// The `m` reply.
    ChangeHeader,
    /// The `p` reply.
    Progress,
    /// The `q` reply.
    Quarantine,
    /// The `r` reply.
    Reject,
    /// The `s` reply.
    Skip,
    /// The `t` reply.
    Tempfail,
    /// The `y` reply.
    ReplyCode,
}

impl From<ReplyKind> for u8 {
    fn from(kind: ReplyKind) -> Self {
        match kind {
            ReplyKind::AddRcpt => b'+',
            ReplyKind::DeleteRcpt => b'-',
            ReplyKind::AddRcptExt => b'2',
            ReplyKind::OptNeg => b'O',
            ReplyKind::Accept => b'a',
            ReplyKind::ReplaceBody => b'b',
            ReplyKind::Continue => b'c',
            ReplyKind::Discard => b'd',
            ReplyKind::ChangeSender => b'e',
            ReplyKind::AddHeader => b'h',
            ReplyKind::InsertHeader => b'i',
            ReplyKind::ChangeHeader => b'm',
            ReplyKind::Progress => b'p',
            ReplyKind::Quarantine => b'q',
            ReplyKind::Reject => b'r',
            ReplyKind::Skip => b's',
            ReplyKind::Tempfail => b't',
            ReplyKind::ReplyCode => b'y',
        }
    }
}

impl TryFrom<u8> for ReplyKind {
    type Error = TryFromByteError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            b'+' => Ok(Self::AddRcpt),
            b'-' => Ok(Self::DeleteRcpt),
            b'2' => Ok(Self::AddRcptExt),
            b'O' => Ok(Self::OptNeg),
            b'a' => Ok(Self::Accept),
            b'b' => Ok(Self::ReplaceBody),
            b'c' => Ok(Self::Continue),
            b'd' => Ok(Self::Discard),
            b'e' => Ok(Self::ChangeSender),
            b'h' => Ok(Self::AddHeader),
            b'i' => Ok(Self::InsertHeader),
            b'm' => Ok(Self::ChangeHeader),
            b'p' => Ok(Self::Progress),
            b'q' => Ok(Self::Quarantine),
            b'r' => Ok(Self::Reject),
            b's' => Ok(Self::Skip),
            b't' => Ok(Self::Tempfail),
            b'y' => Ok(Self::ReplyCode),
            value => Err(TryFromByteError(value)),
        }
    }
}

/// An error that occurs during parsing of a reply.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum ParseReplyError {
    UnknownReply(u8),
    NoOptNegReply,
    UnknownStage(i32),
    NoI32Found,
    NoCStringFound,
}

impl Display for ParseReplyError {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match *self {
            Self::UnknownReply(byte) => write!(f, "unknown reply: {:?}", Byte(byte)),
            Self::NoOptNegReply => write!(f, "no option negotiation reply found"),
            Self::UnknownStage(index) => write!(f, "unknown macro stage: {index}"),
            Self::NoI32Found => write!(f, "no i32 found"),
            Self::NoCStringFound => write!(f, "no C string found"),
        }
    }
}

impl Error for ParseReplyError {}

/// A reply.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Reply {
    // Note: The trailing comments describe restrictions on data as produced by
    // `ContextActions` implementations. However, unlike in `Command` these
    // restrictions are not checked when parsing into this data model.

    /// The `+` reply.
    AddRcpt {
        rcpt: CString,  // non-empty
    },
    /// The `-` reply.
    DeleteRcpt {
        rcpt: CString,  // non-empty
    },
    /// The `2` reply.
    AddRcptExt {
        rcpt: CString,  // non-empty
        args: Option<CString>,
    },
    /// The `O` reply.
    OptNeg {
        version: Version,
        actions: Actions,
        opts: ProtoOpts,
        macros: HashMap<MacroStage, CString>,
    },
    /// The `a` reply.
    Accept,
    /// The `b` reply.
    ReplaceBody { chunk: Bytes },
    /// The `c` reply.
    Continue,
    /// The `d` reply.
    Discard,
    /// The `e` reply.
    ChangeSender {
        mail: CString,  // non-empty
        args: Option<CString>,
    },
    /// The `h` reply.
    AddHeader {
        name: CString,  // non-empty
        value: CString,
    },
    /// The `i` reply.
    InsertHeader {
        index: i32,  // non-negative
        name: CString,  // non-empty
        value: CString,
    },
    /// The `m` reply.
    ChangeHeader {
        name: CString,  // non-empty
        index: i32,  // non-negative
        value: CString,
    },
    /// The `p` reply.
    Progress,
    /// The `q` reply.
    Quarantine {
        reason: CString,  // non-empty
    },
    /// The `r` reply.
    Reject,
    /// The `s` reply.
    Skip,
    /// The `t` reply.
    Tempfail,
    /// The `y` reply.
    ReplyCode {
        reply: CString,  // conforms to reply as produced in Context, eg "550 ..."
    },
}

impl Reply {
    /// Parses a reply from a milter protocol message.
    pub fn parse_reply(msg: Message) -> Result<Self, ParseReplyError> {
        let kind = ReplyKind::try_from(msg.kind)
            .map_err(|e| ParseReplyError::UnknownReply(e.byte()))?;

        let mut buf = msg.buffer;

        match kind {
            ReplyKind::AddRcpt => {
                let rcpt = get_c_string(&mut buf)?;

                Ok(Self::AddRcpt { rcpt })
            }
            ReplyKind::DeleteRcpt => {
                let rcpt = get_c_string(&mut buf)?;

                Ok(Self::DeleteRcpt { rcpt })
            }
            ReplyKind::AddRcptExt => {
                let rcpt = get_c_string(&mut buf)?;
                let args = if buf.has_remaining() {
                    Some(get_c_string(&mut buf)?)
                } else {
                    None
                };

                Ok(Self::AddRcptExt { rcpt, args })
            }
            ReplyKind::OptNeg => {
                if buf.remaining() < 12 {
                    return Err(ParseReplyError::NoOptNegReply);
                }

                let version = buf.get_u32();
                let actions = Actions::from_bits_truncate(buf.get_u32());
                let opts = ProtoOpts::from_bits_truncate(buf.get_u32());

                let mut macros = HashMap::new();

                // If there is more than an `i32` at this point, require a
                // `MacroStage` and a C string following, else ignore the rest.
                while buf.remaining() > 4 {
                    let stage = MacroStage::try_from(buf.get_i32())
                        .map_err(|e| ParseReplyError::UnknownStage(e.index()))?;
                    let m = get_c_string(&mut buf)?;

                    macros.insert(stage, m);
                }

                Ok(Self::OptNeg {
                    version,
                    actions,
                    opts,
                    macros,
                })
            }
            ReplyKind::Accept => Ok(Self::Accept),
            ReplyKind::ReplaceBody => Ok(Self::ReplaceBody { chunk: buf }),
            ReplyKind::Continue => Ok(Self::Continue),
            ReplyKind::Discard => Ok(Self::Discard),
            ReplyKind::ChangeSender => {
                let mail = get_c_string(&mut buf)?;
                let args = if buf.has_remaining() {
                    Some(get_c_string(&mut buf)?)
                } else {
                    None
                };

                Ok(Self::ChangeSender { mail, args })
            }
            ReplyKind::AddHeader => {
                let name = get_c_string(&mut buf)?;
                let value = get_c_string(&mut buf)?;

                Ok(Self::AddHeader { name, value })
            }
            ReplyKind::InsertHeader => {
                let index = get_i32(&mut buf)?;
                let name = get_c_string(&mut buf)?;
                let value = get_c_string(&mut buf)?;

                Ok(Self::InsertHeader { index, name, value })
            }
            ReplyKind::ChangeHeader => {
                let index = get_i32(&mut buf)?;
                let name = get_c_string(&mut buf)?;
                let value = get_c_string(&mut buf)?;

                Ok(Self::ChangeHeader { name, index, value })
            }
            ReplyKind::Progress => Ok(Self::Progress),
            ReplyKind::Quarantine => {
                let reason = get_c_string(&mut buf)?;

                Ok(Self::Quarantine { reason })
            }
            ReplyKind::Reject => Ok(Self::Reject),
            ReplyKind::Skip => Ok(Self::Skip),
            ReplyKind::Tempfail => Ok(Self::Tempfail),
            ReplyKind::ReplyCode => {
                let reply = get_c_string(&mut buf)?;

                Ok(Self::ReplyCode { reply })
            }
        }
    }

    /// Converts this reply to a milter protocol message.
    pub fn into_message(self) -> Message {
        match self {
            Self::AddRcpt { rcpt } => {
                let rcpt = rcpt.to_bytes_with_nul();

                Message::new(ReplyKind::AddRcpt, Bytes::copy_from_slice(rcpt))
            }
            Self::DeleteRcpt { rcpt } => {
                let rcpt = rcpt.to_bytes_with_nul();

                Message::new(ReplyKind::DeleteRcpt, Bytes::copy_from_slice(rcpt))
            }
            Self::AddRcptExt { rcpt, args } => {
                let rcpt = rcpt.to_bytes_with_nul();

                let mut buf = BytesMut::with_capacity(rcpt.len());

                buf.put(rcpt);
                if let Some(args) = args {
                    buf.put(args.to_bytes_with_nul());
                }

                Message::new(ReplyKind::AddRcptExt, buf)
            }
            Self::OptNeg { version, actions, opts, macros } => {
                let mut buf = BytesMut::with_capacity(12);

                buf.put_u32(version);
                buf.put_u32(actions.bits());
                buf.put_u32(opts.bits());

                for stage in MacroStage::all_sorted_by_index() {
                    if let Some(macros) = macros.get(&stage) {
                        buf.put_i32(stage.into());
                        buf.put(macros.to_bytes_with_nul());
                    }
                }

                Message::new(ReplyKind::OptNeg, buf)
            }
            Self::Accept => Message::new(ReplyKind::Accept, Bytes::new()),
            Self::ReplaceBody { chunk } => Message::new(ReplyKind::ReplaceBody, chunk),
            Self::Continue => Message::new(ReplyKind::Continue, Bytes::new()),
            Self::Discard => Message::new(ReplyKind::Discard, Bytes::new()),
            Self::ChangeSender { mail, args } => {
                let mail = mail.to_bytes_with_nul();

                let mut buf = BytesMut::with_capacity(mail.len());

                buf.put(mail);
                if let Some(args) = args {
                    buf.put(args.to_bytes_with_nul());
                }

                Message::new(ReplyKind::ChangeSender, buf)
            }
            Self::AddHeader { name, value } => {
                let name = name.to_bytes_with_nul();
                let value = value.to_bytes_with_nul();

                let mut buf = BytesMut::with_capacity(name.len() + value.len());

                buf.put(name);
                buf.put(value);

                Message::new(ReplyKind::AddHeader, buf)
            }
            Self::InsertHeader { index, name, value } => {
                let name = name.to_bytes_with_nul();
                let value = value.to_bytes_with_nul();

                let mut buf = BytesMut::with_capacity(name.len() + value.len() + 4);

                buf.put_i32(index);
                buf.put(name);
                buf.put(value);

                Message::new(ReplyKind::InsertHeader, buf)
            }
            Self::ChangeHeader { name, index, value } => {
                let name = name.to_bytes_with_nul();
                let value = value.to_bytes_with_nul();

                let mut buf = BytesMut::with_capacity(name.len() + value.len() + 4);

                buf.put_i32(index);
                buf.put(name);
                buf.put(value);

                Message::new(ReplyKind::ChangeHeader, buf)
            }
            Self::Progress => Message::new(ReplyKind::Progress, Bytes::new()),
            Self::Quarantine { reason } => {
                let reason = reason.to_bytes_with_nul();

                Message::new(ReplyKind::Quarantine, Bytes::copy_from_slice(reason))
            }
            Self::Reject => Message::new(ReplyKind::Reject, Bytes::new()),
            Self::Skip => Message::new(ReplyKind::Skip, Bytes::new()),
            Self::Tempfail => Message::new(ReplyKind::Tempfail, Bytes::new()),
            Self::ReplyCode { reply } => {
                let reply = reply.to_bytes_with_nul();

                Message::new(ReplyKind::ReplyCode, Bytes::copy_from_slice(reply))
            }
        }
    }
}

fn get_i32(buf: &mut Bytes) -> Result<i32, ParseReplyError> {
    if buf.remaining() < 4 {
        return Err(ParseReplyError::NoI32Found);
    }
    Ok(buf.get_i32())
}

fn get_c_string(buf: &mut Bytes) -> Result<CString, ParseReplyError> {
    super::get_c_string(buf).map_err(|_| ParseReplyError::NoCStringFound)
}

#[cfg(test)]
mod tests {
    use super::*;
    use byte_strings::c_str;

    #[test]
    fn add_rcpt_into_message() {
        let rcpt = c_str!("<rcpt@example.com>");

        let add_rcpt = Reply::AddRcpt { rcpt: rcpt.into() };
        let add_rcpt_ext = Reply::AddRcptExt {
            rcpt: rcpt.into(),
            args: Some(c_str!("ARGS").into()),
        };

        assert_eq!(
            add_rcpt.into_message(),
            Message::new(ReplyKind::AddRcpt, "<rcpt@example.com>\0")
        );
        assert_eq!(
            add_rcpt_ext.into_message(),
            Message::new(ReplyKind::AddRcptExt, "<rcpt@example.com>\0ARGS\0")
        );
    }
}