indymilter 0.3.0

// indymilter – asynchronous milter library
// Copyright © 2021–2023 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/>.

use crate::message::{TryFromByteError, TryFromIndexError};
use std::{
    collections::HashMap,
    ffi::{CStr, CString},
};

/// The stage for which macros are sent.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum MacroStage {
    /// The `connect` stage.
    Connect,
    /// The `helo` stage.
    Helo,
    /// The `mail` stage.
    Mail,
    /// The `rcpt` stage.
    Rcpt,
    /// The `data` stage.
    Data,
    /// The `eoh` stage.
    Eoh,
    /// The `eom` stage.
    Eom,
}

impl MacroStage {
    /// Returns an iterator of all macro stages.
    pub fn all() -> impl DoubleEndedIterator<Item = Self> {
        use MacroStage::*;
        [Connect, Helo, Mail, Rcpt, Data, Eoh, Eom].into_iter()
    }

    /// Returns an iterator of all macro stages, sorted by their index value.
    pub fn all_sorted_by_index() -> impl DoubleEndedIterator<Item = Self> {
        // In libmilter, Eoh and Eom are out of order. This function is provided
        // to be able to iterate in the same order as libmilter.
        use MacroStage::*;
        [Connect, Helo, Mail, Rcpt, Data, Eom, Eoh].into_iter()
    }
}

impl From<MacroStage> for i32 {
    fn from(stage: MacroStage) -> Self {
        match stage {
            MacroStage::Connect => 0,
            MacroStage::Helo => 1,
            MacroStage::Mail => 2,
            MacroStage::Rcpt => 3,
            MacroStage::Data => 4,
            MacroStage::Eoh => 6,
            MacroStage::Eom => 5,
        }
    }
}

impl TryFrom<i32> for MacroStage {
    type Error = TryFromIndexError;

    fn try_from(value: i32) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(Self::Connect),
            1 => Ok(Self::Helo),
            2 => Ok(Self::Mail),
            3 => Ok(Self::Rcpt),
            4 => Ok(Self::Data),
            6 => Ok(Self::Eoh),
            5 => Ok(Self::Eom),
            value => Err(TryFromIndexError::new(value)),
        }
    }
}

impl From<MacroStage> for u8 {
    fn from(stage: MacroStage) -> Self {
        match stage {
            MacroStage::Connect => b'C',
            MacroStage::Helo => b'H',
            MacroStage::Mail => b'M',
            MacroStage::Rcpt => b'R',
            MacroStage::Data => b'T',
            MacroStage::Eoh => b'N',
            MacroStage::Eom => b'E',
        }
    }
}

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

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            b'C' => Ok(Self::Connect),
            b'H' => Ok(Self::Helo),
            b'M' => Ok(Self::Mail),
            b'R' => Ok(Self::Rcpt),
            b'T' => Ok(Self::Data),
            b'N' => Ok(Self::Eoh),
            b'E' => Ok(Self::Eom),
            value => Err(TryFromByteError::new(value)),
        }
    }
}

// Note: This struct does *not* implement `Default` or `Clone`. Callbacks
// receive a mutable reference to this struct but must not be able to replace
// it.
/// Currently defined macros.
#[derive(Debug, Eq, PartialEq)]
pub struct Macros(HashMap<MacroStage, HashMap<CString, CString>>);

impl Macros {
    pub(crate) fn new() -> Self {
        Self(
            MacroStage::all()
                .map(|stage| (stage, Default::default()))
                .collect(),
        )
    }

    pub(crate) fn clone_internal(&self) -> Self {
        Self(self.0.clone())
    }

    /// Returns the value associated with the given macro name, if any.
    ///
    /// Single-character macros like `i` or `{i}` can be looked up in either
    /// form (both names `i` and `{i}` return the same result).
    pub fn get(&self, name: &CStr) -> Option<&CStr> {
        let alt_name;
        let alt_name = match *name.to_bytes() {
            [x] => {
                alt_name = [b'{', x, b'}', 0];
                Some(CStr::from_bytes_with_nul(&alt_name[..4]).unwrap())
            }
            [b'{', x, b'}'] => {
                alt_name = [x, 0, 0, 0];
                Some(CStr::from_bytes_with_nul(&alt_name[..2]).unwrap())
            }
            _ => None,
        };

        MacroStage::all().rev().find_map(|stage| {
            let m = &self.0[&stage];
            m.get(name)
                .or_else(|| alt_name.and_then(|name| m.get(name)))
                .map(|v| v.as_ref())
        })
    }

    /// Returns a new hash map with all defined macros.
    pub fn to_hash_map(&self) -> HashMap<CString, CString> {
        MacroStage::all()
            .flat_map(|stage| self.0[&stage].clone())
            .collect()
    }

    pub(crate) fn insert(&mut self, stage: MacroStage, entries: HashMap<CString, CString>) {
        self.0.insert(stage, entries);
    }

    pub(crate) fn clear(&mut self) {
        for s in MacroStage::all() {
            self.0.get_mut(&s).unwrap().clear();
        }
    }

    pub(crate) fn clear_after(&mut self, stage: MacroStage) {
        for s in MacroStage::all().rev().take_while(|&s| s != stage) {
            self.0.get_mut(&s).unwrap().clear();
        }
    }
}

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

    #[test]
    fn macro_lookup() {
        let mut macros = Macros::new();

        macros.insert(
            MacroStage::Connect,
            HashMap::from([
                (c_str!("a").into(), c_str!("b").into()),
                (c_str!("{x}").into(), c_str!("y").into()),
                (c_str!("{name}").into(), c_str!("value").into()),
            ]),
        );

        assert_eq!(macros.get(c_str!("n")), None);
        assert_eq!(macros.get(c_str!("{none}")), None);
        assert_eq!(macros.get(c_str!("{name}")), Some(c_str!("value")));
        assert_eq!(macros.get(c_str!("a")), Some(c_str!("b")));
        assert_eq!(macros.get(c_str!("{a}")), Some(c_str!("b")));
        assert_eq!(macros.get(c_str!("x")), Some(c_str!("y")));
        assert_eq!(macros.get(c_str!("{x}")), Some(c_str!("y")));
    }

    #[test]
    fn to_hash_map_ok() {
        let mut macros = Macros::new();

        macros.insert(
            MacroStage::Connect,
            HashMap::from([
                (c_str!("{name1}").into(), c_str!("connect1").into()),
                (c_str!("{name2}").into(), c_str!("connect2").into()),
            ]),
        );
        macros.insert(
            MacroStage::Mail,
            HashMap::from([
                (c_str!("{name2}").into(), c_str!("mail2").into()),
                (c_str!("{name3}").into(), c_str!("mail3").into()),
            ]),
        );

        assert_eq!(
            macros.to_hash_map(),
            HashMap::from([
                (c_str!("{name1}").into(), c_str!("connect1").into()),
                (c_str!("{name2}").into(), c_str!("mail2").into()),
                (c_str!("{name3}").into(), c_str!("mail3").into()),
            ]),
        );

        macros.clear_after(MacroStage::Helo);

        assert_eq!(
            macros.to_hash_map(),
            HashMap::from([
                (c_str!("{name1}").into(), c_str!("connect1").into()),
                (c_str!("{name2}").into(), c_str!("connect2").into()),
            ]),
        );
    }
}