netget/server/milter/
mod.rs

1//! Milter (mail filter) server. Rust owns the packets, option negotiation, macros and the
2//! collection of headers and body; the handler decides at each SMTP stage and may modify the
3//! message at its end.
4pub mod actions;
5pub mod wire;
6
7use crate::llm::action_helper::call_llm;
8use crate::llm::actions::protocol_trait::ActionResult;
9use crate::logging::emit::Log;
10use crate::protocol::{Event, SpawnContext};
11use crate::server::connection::ConnectionId;
12use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
13use anyhow::{ensure, Result};
14use serde_json::{json, Map, Value};
15use std::{net::SocketAddr, time::Duration};
16use tokio::net::{TcpListener, TcpStream};
17
18pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
19pub const MAX_IDLE_TIMEOUT_SECS: u64 = 86400;
20/// Most macros remembered for one connection.
21pub const MAX_MACROS: usize = 256;
22
23pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
24    let secs = ctx
25        .startup_params
26        .as_ref()
27        .map(|p| p.get_optional_u64("idle_timeout_secs"))
28        .transpose()?
29        .flatten()
30        .unwrap_or(IDLE_TIMEOUT.as_secs());
31    ensure!(
32        (1..=MAX_IDLE_TIMEOUT_SECS).contains(&secs),
33        "idle_timeout_secs must be between 1 and {MAX_IDLE_TIMEOUT_SECS}"
34    );
35    let idle = Duration::from_secs(secs);
36    let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
37    let addr = listener.local_addr()?;
38    Log::new(Some(&ctx.status_tx)).info(format!("Milter listening on {addr}"));
39    let state = ctx.state.clone();
40    let server_id = ctx.server_id;
41    let accept = tokio::spawn(async move {
42        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(
43            crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
44        );
45        loop {
46            let (socket, peer, permit) = match crate::server::accept_bounded::accept_bounded(
47                &listener,
48                &limiter,
49                b"",
50                "Milter",
51                Some(&ctx.status_tx),
52            )
53            .await
54            {
55                Ok(v) => v,
56                Err(_) => break,
57            };
58            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
59            let now = crate::utils::clock::Instant::now();
60            ctx.state
61                .add_connection_to_server(
62                    server_id,
63                    ConnectionState {
64                        id,
65                        remote_addr: peer,
66                        local_addr: addr,
67                        bytes_sent: 0,
68                        bytes_received: 0,
69                        packets_sent: 0,
70                        packets_received: 0,
71                        last_activity: now,
72                        status: ConnectionStatus::Active,
73                        status_changed_at: now,
74                        protocol_info: ProtocolConnectionInfo::empty(),
75                    },
76                )
77                .await;
78            let child = ctx.clone();
79            ctx.state
80                .spawn_server_task(server_id, async move {
81                    let _permit = permit;
82                    if let Err(e) = session(&child, id, socket, idle).await {
83                        Log::new(Some(&child.status_tx))
84                            .warn(format!("Milter connection {id} ended: {e:#}"));
85                    }
86                    child
87                        .state
88                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
89                        .await;
90                    let _ = child.status_tx.send("__UPDATE_UI__".into());
91                })
92                .await;
93        }
94    });
95    state.register_server_task(server_id, accept).await;
96    Ok(addr)
97}
98
99#[derive(Default)]
100struct Message {
101    sender: String,
102    recipients: Vec<String>,
103    headers: Vec<Value>,
104    body: Vec<u8>,
105    truncated: bool,
106}
107
108struct Session {
109    actions: u32,
110    macros: Map<String, Value>,
111    connection: Value,
112    message: Message,
113}
114
115fn outcome(ctx: &SpawnContext, id: ConnectionId, stage: &str, decision: &str) {
116    let line = format!("Milter connection {id} stage={stage} decision={decision}");
117    let log = Log::new(Some(&ctx.status_tx));
118    if decision.starts_with("fail_closed") {
119        log.error(line)
120    } else {
121        log.info(line)
122    }
123}
124
125/// The handler's actions for one stage, or `None` when it failed (the caller tempfails).
126async fn ask(
127    ctx: &SpawnContext,
128    id: ConnectionId,
129    event: Event,
130    stage: &str,
131) -> Option<Vec<Value>> {
132    match call_llm(
133        &ctx.llm_client,
134        &ctx.state,
135        ctx.server_id,
136        Some(id),
137        &event,
138        &actions::MilterProtocol,
139    )
140    .await
141    {
142        Ok(r) if r.failures.is_empty() => {
143            let mut out = Vec::new();
144            let mut stack = r.protocol_results;
145            while let Some(item) = stack.pop() {
146                match item {
147                    ActionResult::Custom { name, data } if name.starts_with("milter_") => {
148                        out.push(data)
149                    }
150                    ActionResult::Multiple(items) => stack.extend(items),
151                    _ => {}
152                }
153            }
154            out.reverse();
155            Some(out)
156        }
157        Ok(_) => {
158            outcome(ctx, id, stage, "fail_closed_invalid_reply");
159            None
160        }
161        Err(_) => {
162            outcome(ctx, id, stage, "fail_closed_llm_error");
163            None
164        }
165    }
166}
167
168const DECISIONS: [&str; 6] = [
169    "milter_continue",
170    "milter_accept",
171    "milter_reject",
172    "milter_tempfail",
173    "milter_discard",
174    "milter_reply",
175];
176
177/// Write the stage's replies: modifications (end of message only, and only those the MTA
178/// allowed), then one decision. Returns the decision's name.
179async fn answer<W: tokio::io::AsyncWrite + Unpin>(
180    ctx: &SpawnContext,
181    id: ConnectionId,
182    w: &mut W,
183    s: &Session,
184    stage: &str,
185    answers: Option<Vec<Value>>,
186    end_of_message: bool,
187) -> Result<&'static str> {
188    let Some(answers) = answers else {
189        wire::write(w, wire::R_TEMPFAIL, &[]).await?;
190        return Ok("tempfail");
191    };
192    let log = Log::new(Some(&ctx.status_tx));
193    let mut decision = None;
194    for a in &answers {
195        let name = a["type"].as_str().unwrap_or_default();
196        if DECISIONS.contains(&name) {
197            decision = Some(a.clone());
198            continue;
199        }
200        if !end_of_message {
201            log.warn(format!(
202                "Milter: {name} is only possible at end of message; ignored at {stage}"
203            ));
204            continue;
205        }
206        let needed = match name {
207            "milter_add_header" => wire::F_ADDHDRS,
208            "milter_change_header" => wire::F_CHGHDRS,
209            "milter_add_rcpt" => wire::F_ADDRCPT,
210            "milter_del_rcpt" => wire::F_DELRCPT,
211            "milter_replace_body" => wire::F_CHGBODY,
212            "milter_quarantine" => wire::F_QUARANTINE,
213            _ => 0,
214        };
215        if s.actions & needed == 0 {
216            log.warn(format!("Milter: the MTA did not allow {name}; not sent"));
217            continue;
218        }
219        let (cmd, data) = actions::reply_packet(a)?;
220        wire::write(w, cmd, &data).await?;
221    }
222    let (cmd, data) = match &decision {
223        Some(d) => actions::reply_packet(d)?,
224        // Saying nothing lets the message through this stage.
225        None if end_of_message => (wire::R_ACCEPT, vec![]),
226        None => (wire::R_CONTINUE, vec![]),
227    };
228    wire::write(w, cmd, &data).await?;
229    let name = wire::reply_name(cmd);
230    outcome(
231        ctx,
232        id,
233        stage,
234        if decision.is_some() {
235            "model_answer"
236        } else {
237            "model_silent_continue"
238        },
239    );
240    Ok(name)
241}
242
243async fn session(
244    ctx: &SpawnContext,
245    id: ConnectionId,
246    socket: TcpStream,
247    idle: Duration,
248) -> Result<()> {
249    let (mut r, mut w) = tokio::io::split(socket);
250    let mut s = Session {
251        actions: 0,
252        macros: Map::new(),
253        connection: json!({}),
254        message: Message::default(),
255    };
256    while let Some((cmd, data)) = wire::read_packet(&mut r, idle).await? {
257        ctx.state
258            .update_connection_stats(
259                ctx.server_id,
260                id,
261                Some(data.len() as u64 + 5),
262                None,
263                Some(1),
264                None,
265            )
266            .await;
267        let stage_event = |e: &'static std::sync::LazyLock<crate::protocol::EventType>,
268                           mut v: Value,
269                           s: &Session| {
270            v["connection"] = s.connection.clone();
271            v["macros"] = Value::Object(s.macros.clone());
272            Event::new(e, v)
273        };
274        match cmd {
275            wire::C_OPTNEG => {
276                let (version, actions, _protocol) = wire::parse_optneg(&data)?;
277                ensure!(version >= 2, "milter protocol version {version} is too old");
278                s.actions = actions & wire::ALL_ACTIONS;
279                wire::write(
280                    &mut w,
281                    wire::C_OPTNEG,
282                    &wire::optneg(version.min(wire::VERSION), s.actions, 0),
283                )
284                .await?;
285            }
286            wire::C_MACRO => {
287                let fields = wire::strings(data.get(1..).unwrap_or_default());
288                for pair in fields.chunks(2) {
289                    if let [k, v] = pair {
290                        if s.macros.len() < MAX_MACROS || s.macros.contains_key(k) {
291                            s.macros.insert(k.clone(), json!(v));
292                        }
293                    }
294                }
295            }
296            wire::C_CONNECT => {
297                let (hostname, family, port, address) = wire::parse_connect(&data)?;
298                s.connection = json!({"hostname": hostname, "family": family.to_string(), "address": address, "port": port});
299                let e = stage_event(
300                    &actions::CONNECT_EVENT,
301                    json!({"hostname": hostname, "address": address, "port": port}),
302                    &s,
303                );
304                let a = ask(ctx, id, e, "connect").await;
305                answer(ctx, id, &mut w, &s, "connect", a, false).await?;
306            }
307            wire::C_HELO => {
308                let helo = wire::strings(&data).into_iter().next().unwrap_or_default();
309                s.connection["helo"] = json!(helo);
310                let e = stage_event(&actions::HELO_EVENT, json!({"helo": helo}), &s);
311                let a = ask(ctx, id, e, "helo").await;
312                answer(ctx, id, &mut w, &s, "helo", a, false).await?;
313            }
314            wire::C_MAIL => {
315                let mut f = wire::strings(&data);
316                ensure!(!f.is_empty(), "MAIL without a sender");
317                let sender = f.remove(0);
318                s.message = Message {
319                    sender: sender.clone(),
320                    ..Message::default()
321                };
322                let e = stage_event(
323                    &actions::MAIL_EVENT,
324                    json!({"sender": sender, "esmtp_args": f}),
325                    &s,
326                );
327                let a = ask(ctx, id, e, "mail").await;
328                answer(ctx, id, &mut w, &s, "mail", a, false).await?;
329            }
330            wire::C_RCPT => {
331                let mut f = wire::strings(&data);
332                ensure!(!f.is_empty(), "RCPT without a recipient");
333                let rcpt = f.remove(0);
334                let e = stage_event(
335                    &actions::RCPT_EVENT,
336                    json!({"recipient": rcpt, "esmtp_args": f}),
337                    &s,
338                );
339                let a = ask(ctx, id, e, "rcpt").await;
340                if answer(ctx, id, &mut w, &s, "rcpt", a, false).await? == "continue" {
341                    s.message.recipients.push(rcpt);
342                }
343            }
344            wire::C_HEADER => {
345                let f = wire::strings(&data);
346                if s.message.headers.len() < wire::MAX_HEADERS {
347                    s.message.headers.push(json!({"name": f.first().cloned().unwrap_or_default(), "value": f.get(1).cloned().unwrap_or_default()}));
348                }
349                wire::write(&mut w, wire::R_CONTINUE, &[]).await?;
350            }
351            wire::C_BODY => {
352                let room = wire::MAX_BODY.saturating_sub(s.message.body.len());
353                s.message.truncated |= data.len() > room;
354                s.message
355                    .body
356                    .extend_from_slice(&data[..data.len().min(room)]);
357                wire::write(&mut w, wire::R_CONTINUE, &[]).await?;
358            }
359            wire::C_BODYEOB => {
360                if !data.is_empty() {
361                    let room = wire::MAX_BODY.saturating_sub(s.message.body.len());
362                    s.message.truncated |= data.len() > room;
363                    s.message
364                        .body
365                        .extend_from_slice(&data[..data.len().min(room)]);
366                }
367                let m = &s.message;
368                let e = stage_event(
369                    &actions::MESSAGE_EVENT,
370                    json!({"sender": m.sender, "recipients": m.recipients, "headers": m.headers,
371                           "body": String::from_utf8_lossy(&m.body), "body_truncated": m.truncated}),
372                    &s,
373                );
374                let a = ask(ctx, id, e, "message").await;
375                answer(ctx, id, &mut w, &s, "message", a, true).await?;
376                s.message = Message::default();
377            }
378            wire::C_DATA | wire::C_EOH | wire::C_UNKNOWN => {
379                wire::write(&mut w, wire::R_CONTINUE, &[]).await?
380            }
381            wire::C_ABORT => s.message = Message::default(),
382            wire::C_QUIT_NC => {
383                s.message = Message::default();
384                s.macros.clear();
385                s.connection = json!({});
386            }
387            wire::C_QUIT => return Ok(()),
388            other => anyhow::bail!("unknown milter command {:?}", other as char),
389        }
390    }
391    Ok(())
392}
netget/server/milter/mod.rs

netget/server/milter/
mod.rs
