indymilter 0.3.0

Asynchronous milter library
Documentation 
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

#[cfg(feature = "miltertest-tests")]
mod miltertest;

#[cfg(feature = "miltertest-tests")]
pub use self::miltertest::*;

use indymilter::{
    message::{self, command::Command, reply::Reply},
    Callbacks, Config,
};
use std::{
    io,
    net::{Ipv4Addr, SocketAddr},
    time::Duration,
};
use tokio::{
    io::{AsyncWriteExt, BufStream},
    net::{TcpListener, TcpStream, ToSocketAddrs},
    sync::oneshot,
    task::{self, JoinHandle},
};

pub fn init_tracing_subscriber() {
    let _ = tracing_subscriber::fmt::try_init();
}

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

pub fn default_config() -> Config {
    // Override very long default connection timeout for tests.
    Config {
        connection_timeout: Duration::from_secs(30),
        ..Default::default()
    }
}

pub struct Milter {
    milter_handle: JoinHandle<io::Result<()>>,
    shutdown: oneshot::Sender<()>,
    addr: SocketAddr,
}

impl Milter {
    pub async fn spawn<T: Send + 'static>(
        addr: impl ToSocketAddrs,
        callbacks: Callbacks<T>,
        config: Config,
    ) -> io::Result<Self> {
        let listener = TcpListener::bind(addr).await?;

        let addr = listener.local_addr()?;

        let (shutdown, shutdown_rx) = oneshot::channel();

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

        Ok(Self {
            milter_handle: milter,
            shutdown,
            addr,
        })
    }

    pub fn addr(&self) -> SocketAddr {
        self.addr
    }

    pub async fn shutdown(self) -> io::Result<()> {
        // The milter task was spawned with `tokio::spawn`. Sometimes the test
        // function can proceed at each await point, giving the milter task no
        // chance to run to completion before the shutdown signal is received.
        // So we first yield a few times, hoping to give the milter time to
        // proceed a bit further (eg, terminate an open session).
        //
        // Note: This is not intended to impact test correctness, rather to
        // match the behaviour of a real milter (eg in terms of logging).
        yield_a_few_times().await;

        let _ = self.shutdown.send(());

        self.milter_handle.await?
    }
}

pub async fn yield_a_few_times() {
    for _ in 0..10 {
        task::yield_now().await;
    }
}

pub struct Client {
    stream: BufStream<TcpStream>,
}

impl Client {
    pub async fn connect(addr: impl ToSocketAddrs) -> io::Result<Self> {
        let stream = TcpStream::connect(addr).await?;

        Ok(Self {
            stream: BufStream::new(stream),
        })
    }

    pub async fn write_command(&mut self, cmd: Command) -> io::Result<()> {
        let msg = cmd.into_message();

        message::write(&mut self.stream, msg).await?;

        Ok(())
    }

    pub async fn write_bytes(&mut self, bytes: &[u8]) -> io::Result<()> {
        self.stream.write_all(bytes).await?;
        self.stream.flush().await?;

        Ok(())
    }

    pub async fn read_reply(&mut self) -> io::Result<Reply> {
        let msg = message::read(&mut self.stream).await?;

        // Assume that indymilter always responds with a well-formed reply.
        let reply = Reply::parse_reply(msg).unwrap();

        Ok(reply)
    }

    // Note: consumes and therefore drops this client and connection.
    pub async fn disconnect(mut self) -> io::Result<()> {
        self.stream.shutdown().await
    }
}