indymilter-test 0.2.0

use indymilter::Callbacks;
use std::{
    net::Ipv4Addr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};
use tokio::{net::TcpListener, sync::oneshot};

use indymilter_test::*;

const LOCALHOST: (Ipv4Addr, u16) = (Ipv4Addr::LOCALHOST, 0);

#[tokio::test]
async fn body_skip() {
    fn milter_callbacks() -> Callbacks<()> {
        use indymilter::Status;

        let called = Arc::new(AtomicBool::new(false));

        Callbacks::new()
            .on_negotiate(|cx, _, _| {
                Box::pin(async {
                    cx.requested_opts |= ProtoOpts::SKIP;

                    Status::Continue
                })
            })
            .on_body(move |_cx, _| {
                let called = called.clone();

                Box::pin(async move {
                    if called.swap(true, Ordering::SeqCst) {
                        Status::Skip
                    } else {
                        Status::Continue
                    }
                })
            })
    }

    let _ = tracing_subscriber::fmt::try_init();

    let listener = TcpListener::bind(LOCALHOST).await.unwrap();
    let callbacks = milter_callbacks();
    let config = Default::default();
    let (shutdown_milter, shutdown) = oneshot::channel();

    let milter_addr = listener.local_addr().unwrap();

    let milter = tokio::spawn(indymilter::run(listener, callbacks, config, shutdown));

    let mut conn = TestConnection::open(milter_addr).await.unwrap();

    let status = conn.body("first").await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.body("second").await.unwrap();
    assert_eq!(status, Status::Skip);

    let error = conn.body("third").await.unwrap_err();
    assert!(matches!(error, TestError::MilterUsage));

    conn.close().await.unwrap();

    shutdown_milter.send(()).unwrap();

    milter.await.unwrap().unwrap();
}

#[tokio::test]
async fn noreply_stages() {
    fn milter_callbacks() -> Callbacks<()> {
        use indymilter::Status;

        Callbacks::new()
            .on_negotiate(|cx, _, _| {
                Box::pin(async {
                    cx.requested_opts |= ProtoOpts::NOREPLY_DATA
                        | ProtoOpts::NOREPLY_HEADER
                        | ProtoOpts::NOREPLY_EOH;

                    Status::Continue
                })
            })
            .on_data(|_cx| Box::pin(async { Status::Noreply }))
            .on_header(|_cx, _, _| Box::pin(async { Status::Noreply }))
            // An indymilter implementation detail … This `Continue` is
            // corrected to `Noreply`, and a warning is logged:
            .on_eoh(|_cx| Box::pin(async { Status::Continue }))
    }

    let _ = tracing_subscriber::fmt::try_init();

    let listener = TcpListener::bind(LOCALHOST).await.unwrap();
    let callbacks = milter_callbacks();
    let config = Default::default();
    let (shutdown_milter, shutdown) = oneshot::channel();

    let milter_addr = listener.local_addr().unwrap();

    let milter = tokio::spawn(indymilter::run(listener, callbacks, config, shutdown));

    let mut conn = TestConnection::open(milter_addr).await.unwrap();

    let status = conn.data().await.unwrap();
    assert_eq!(status, Status::Noreply);

    let status = conn.header("Subject", "hello...").await.unwrap();
    assert_eq!(status, Status::Noreply);

    let status = conn.eoh().await.unwrap();
    assert_eq!(status, Status::Noreply);

    conn.close().await.unwrap();

    shutdown_milter.send(()).unwrap();

    milter.await.unwrap().unwrap();
}

#[tokio::test]
async fn disabled_stage() {
    fn milter_callbacks() -> Callbacks<()> {
        use indymilter::Status;

        Callbacks::new().on_mail(|_cx, _| Box::pin(async { Status::Continue }))
    }

    let _ = tracing_subscriber::fmt::try_init();

    let listener = TcpListener::bind(LOCALHOST).await.unwrap();
    let callbacks = milter_callbacks();
    let config = Default::default();
    let (shutdown_milter, shutdown) = oneshot::channel();

    let milter_addr = listener.local_addr().unwrap();

    let milter = tokio::spawn(indymilter::run(listener, callbacks, config, shutdown));

    let mut conn = TestConnection::open(milter_addr).await.unwrap();

    let status = conn.mail(["<from@example.org>"]).await.unwrap();
    assert_eq!(status, Status::Continue);

    // indymilter, like libmilter, automatically disables stages (eg `NO_RCPT`)
    // for which no callback was registered. Therefore sending command Rcpt here
    // is not allowed.
    let error = conn.rcpt(["<to@example.org>"]).await.unwrap_err();
    assert!(matches!(error, TestError::MilterUsage));

    conn.close().await.unwrap();

    shutdown_milter.send(()).unwrap();

    milter.await.unwrap().unwrap();
}