milter 0.2.4

//! The example milter included in the libmilter distribution. See:
//! https://salsa.debian.org/debian/sendmail/blob/master/libmilter/example.c

use milter::{
    on_abort, on_body, on_close, on_data, on_eoh, on_eom, on_header, on_mail, on_negotiate,
    on_unknown, Actions, Context, Error, Milter, ProtocolOpts, Status,
};
use once_cell::sync::Lazy;
use std::{
    env,
    fs::{self, File, OpenOptions},
    io::Write,
    path::PathBuf,
    process,
    sync::Mutex,
    time::SystemTime,
};

struct Message {
    path: PathBuf,
    file: File,
}

static MTA_CAPS: Lazy<Mutex<ProtocolOpts>> = Lazy::new(|| Mutex::new(ProtocolOpts::empty()));

fn set_mta_caps(protocol_opts: ProtocolOpts) {
    *MTA_CAPS.lock().unwrap() = protocol_opts;
}

fn mta_caps() -> ProtocolOpts {
    *MTA_CAPS.lock().unwrap()
}

#[on_negotiate(negotiate_callback)]
fn handle_negotiate(
    _: Context<Message>,
    _: Actions,
    protocol_opts: ProtocolOpts,
) -> (Status, Actions, ProtocolOpts) {
    set_mta_caps(protocol_opts);

    let actions = Actions::ADD_HEADER;

    let mut protocol_opts =
        ProtocolOpts::NO_CONNECT | ProtocolOpts::NO_HELO | ProtocolOpts::NO_RCPT;
    if mta_caps().contains(ProtocolOpts::NOREPLY_HEADER) {
        protocol_opts |= ProtocolOpts::NOREPLY_HEADER;
    }

    (Status::Continue, actions, protocol_opts)
}

#[on_mail(mail_callback)]
fn handle_mail(mut ctx: Context<Message>, _: Vec<&str>) -> milter::Result<Status> {
    let epoch_secs = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .expect("system time out of range")
        .as_secs();
    let path = PathBuf::from(format!("/tmp/msg.{}", epoch_secs));

    let file = match OpenOptions::new()
        .write(true)
        .truncate(true)
        .create(true)
        .open(&path)
    {
        Ok(file) => file,
        Err(_) => return Ok(Status::Tempfail),
    };

    ctx.data.replace(Message { path, file })?;

    Ok(Status::Continue)
}

#[on_data(data_callback)]
fn handle_data(_: Context<Message>) -> Status {
    Status::Continue
}

#[on_header(header_callback)]
fn handle_header(mut ctx: Context<Message>, name: &str, value: &str) -> milter::Result<Status> {
    let msg = ctx.data.borrow_mut().unwrap();

    write!(&mut msg.file, "{}: {}\r\n", name, value).map_err(|e| Error::Custom(e.into()))?;

    Ok(if mta_caps().contains(ProtocolOpts::NOREPLY_HEADER) {
        Status::Noreply
    } else {
        Status::Continue
    })
}

#[on_eoh(eoh_callback)]
fn handle_eoh(mut ctx: Context<Message>) -> milter::Result<Status> {
    let msg = ctx.data.borrow_mut().unwrap();

    write!(&mut msg.file, "\r\n").map_err(|e| Error::Custom(e.into()))?;

    Ok(Status::Continue)
}

#[on_body(body_callback)]
fn handle_body(mut ctx: Context<Message>, content: &[u8]) -> milter::Result<Status> {
    let msg = ctx.data.borrow_mut().unwrap();

    if msg.file.write_all(content).is_err() {
        cleanup(ctx)?;
        return Ok(Status::Tempfail);
    }

    Ok(Status::Continue)
}

#[on_eom(eom_callback)]
fn handle_eom(mut ctx: Context<Message>) -> milter::Result<Status> {
    let msg = ctx.data.take()?.unwrap();

    let path = msg.path.to_str().expect("invalid characters in path");

    ctx.api.add_header("X-Archived", path)?;

    Ok(Status::Continue)
}

#[on_abort(abort_callback)]
fn handle_abort(ctx: Context<Message>) -> milter::Result<Status> {
    cleanup(ctx)?;

    Ok(Status::Continue)
}

fn cleanup(mut ctx: Context<Message>) -> milter::Result<()> {
    if let Some(msg) = ctx.data.take()? {
        fs::remove_file(msg.path).map_err(|e| Error::Custom(e.into()))?;
    }
    Ok(())
}

#[on_close(close_callback)]
fn handle_close(_: Context<Message>) -> Status {
    Status::Accept
}

#[on_unknown(unknown_callback)]
fn handle_unknown(_: Context<Message>, _: &str) -> Status {
    Status::Continue
}

fn main() {
    let args = env::args().collect::<Vec<_>>();

    if args.len() != 3 || args[1] != "-p" {
        eprintln!("usage: {} -p <socket>", args[0]);
        process::exit(1);
    }

    if let Err(e) = Milter::new(&args[2])
        .name("SampleFilter")
        .on_negotiate(negotiate_callback)
        .on_mail(mail_callback)
        .on_data(data_callback)
        .on_header(header_callback)
        .on_eoh(eoh_callback)
        .on_body(body_callback)
        .on_eom(eom_callback)
        .on_abort(abort_callback)
        .on_close(close_callback)
        .on_unknown(unknown_callback)
        .actions(Actions::ADD_HEADER)
        .run()
    {
        eprintln!("milter execution failed: {}", e);
        process::exit(1);
    }
}