Skip to main content

rspamd_client/backend/
async_client.rs

1use crate::backend::traits::*;
2use crate::config::{Config, EnvelopeData};
3use crate::error::RspamdError;
4use crate::protocol::commands::{RspamdCommand, RspamdEndpoint};
5use crate::protocol::encryption::{httpcrypt_decrypt, HttpCryptContext};
6use crate::protocol::scan::parse_scan_reply;
7use crate::protocol::{RspamdLearnReply, RspamdScanReply};
8use bytes::{Bytes, BytesMut};
9use reqwest::header::{HeaderName, HeaderValue};
10use reqwest::Client;
11use std::str::FromStr;
12use std::sync::Arc;
13use url::Url;
14
15/// Persistent asynchronous Rspamd client.
16///
17/// Owns the underlying `reqwest::Client`, so TCP (and TLS, when enabled)
18/// connections are pooled and reused across calls, unlike the one-shot
19/// [`scan_async`] wrapper which pays a fresh connection setup on every call.
20///
21/// Cheap to clone (`reqwest::Client` is internally reference-counted and the
22/// configuration is shared through an `Arc`); `Send + Sync + 'static`, so it
23/// is suitable for stashing in a static or an LRU cache keyed by [`Config`].
24///
25/// ```rust,no_run
26/// use rspamd_client::config::Config;
27/// use rspamd_client::RspamdAsyncClient;
28///
29/// # #[tokio::main]
30/// # async fn main() -> Result<(), rspamd_client::error::RspamdError> {
31/// let config = Config::builder()
32///     .base_url("http://localhost:11333".to_string())
33///     .build();
34/// let client = RspamdAsyncClient::new(config)?;
35/// for email in ["first message", "second message"] {
36///     let reply = client.scan(email, Default::default()).await?;
37///     println!("action: {}", reply.action);
38/// }
39/// # Ok(())
40/// # }
41/// ```
42#[derive(Clone)]
43pub struct RspamdAsyncClient {
44    inner: Client,
45    config: Arc<Config>,
46    crypt: Option<Arc<HttpCryptContext>>,
47}
48
49impl RspamdAsyncClient {
50    /// Create a persistent client from a configuration.
51    pub fn new(config: Config) -> Result<Self, RspamdError> {
52        let client = Client::builder().timeout(config.timeout);
53
54        let client = if let Some(ref proxy) = config.proxy_config {
55            let proxy = reqwest::Proxy::all(proxy.proxy_url.clone())
56                .map_err(|e| RspamdError::HttpError(e.to_string()))?;
57            client.proxy(proxy)
58        } else {
59            client
60        };
61        let client = if let Some(ref tls) = config.tls_settings {
62            if let Some(ca_path) = tls.ca_path.as_ref() {
63                client.add_root_certificate(
64                    reqwest::Certificate::from_pem(
65                        &std::fs::read(std::fs::canonicalize(ca_path.as_str()).unwrap())
66                            .map_err(|e| RspamdError::ConfigError(e.to_string()))?,
67                    )
68                    .map_err(|e| RspamdError::HttpError(e.to_string()))?,
69                )
70            } else {
71                client
72            }
73        } else {
74            client
75        };
76
77        let crypt = config
78            .encryption_key
79            .as_deref()
80            .map(HttpCryptContext::new)
81            .transpose()?
82            .map(Arc::new);
83
84        Ok(Self {
85            inner: client
86                .build()
87                .map_err(|e| RspamdError::HttpError(e.to_string()))?,
88            config: Arc::new(config),
89            crypt,
90        })
91    }
92
93    /// The configuration this client was built from.
94    pub fn config(&self) -> &Config {
95        &self.config
96    }
97
98    /// Scan an email, returning the parsed scan reply.
99    pub async fn scan<B: AsRef<[u8]> + Send>(
100        &self,
101        body: B,
102        envelope_data: EnvelopeData,
103    ) -> Result<RspamdScanReply, RspamdError> {
104        let (headers, body) = self
105            .execute(RspamdCommand::Scan, body, envelope_data)
106            .await?;
107
108        // Check for Message-Offset header to handle body_block feature
109        let message_offset = headers
110            .get("Message-Offset")
111            .map(|v| {
112                v.to_str().map_err(|e| {
113                    RspamdError::HttpError(format!("Invalid Message-Offset header: {}", e))
114                })
115            })
116            .transpose()?;
117        parse_scan_reply(message_offset, body.as_ref())
118    }
119
120    /// Learn an email as spam (requires the controller endpoint and password).
121    pub async fn learn_spam<B: AsRef<[u8]> + Send>(
122        &self,
123        body: B,
124        envelope_data: EnvelopeData,
125    ) -> Result<RspamdLearnReply, RspamdError> {
126        self.learn(RspamdCommand::Learnspam, body, envelope_data)
127            .await
128    }
129
130    /// Learn an email as ham (requires the controller endpoint and password).
131    pub async fn learn_ham<B: AsRef<[u8]> + Send>(
132        &self,
133        body: B,
134        envelope_data: EnvelopeData,
135    ) -> Result<RspamdLearnReply, RspamdError> {
136        self.learn(RspamdCommand::Learnham, body, envelope_data)
137            .await
138    }
139
140    /// Add an email to fuzzy storage; pass `Flag`/`Weight` via
141    /// `EnvelopeData::additional_headers`.
142    pub async fn fuzzy_add<B: AsRef<[u8]> + Send>(
143        &self,
144        body: B,
145        envelope_data: EnvelopeData,
146    ) -> Result<RspamdLearnReply, RspamdError> {
147        self.learn(RspamdCommand::FuzzyAdd, body, envelope_data)
148            .await
149    }
150
151    /// Remove an email from fuzzy storage; pass `Flag` via
152    /// `EnvelopeData::additional_headers`.
153    pub async fn fuzzy_del<B: AsRef<[u8]> + Send>(
154        &self,
155        body: B,
156        envelope_data: EnvelopeData,
157    ) -> Result<RspamdLearnReply, RspamdError> {
158        self.learn(RspamdCommand::FuzzyDel, body, envelope_data)
159            .await
160    }
161
162    async fn learn<B: AsRef<[u8]> + Send>(
163        &self,
164        command: RspamdCommand,
165        body: B,
166        envelope_data: EnvelopeData,
167    ) -> Result<RspamdLearnReply, RspamdError> {
168        let (_headers, body) = self.execute(command, body, envelope_data).await?;
169        Ok(serde_json::from_slice(body.as_ref())?)
170    }
171
172    async fn execute<B: AsRef<[u8]> + Send>(
173        &self,
174        command: RspamdCommand,
175        body: B,
176        envelope_data: EnvelopeData,
177    ) -> Result<(reqwest::header::HeaderMap, Bytes), RspamdError> {
178        let request = ReqwestRequest::new(self.clone(), body, command, envelope_data).await?;
179        request.response().await
180    }
181}
182
183#[cfg(feature = "async")]
184#[deprecated(
185    since = "0.7.0",
186    note = "use RspamdAsyncClient::new for connection pooling"
187)]
188pub fn async_client(options: &Config) -> Result<RspamdAsyncClient, RspamdError> {
189    RspamdAsyncClient::new(options.clone())
190}
191
192/// Deprecated alias kept for backwards compatibility; the client no longer
193/// borrows the configuration, so the lifetime parameter is ignored.
194#[deprecated(since = "0.7.0", note = "use RspamdAsyncClient instead")]
195pub type AsyncClient<'a> = RspamdAsyncClient;
196
197// Temporary structure for making a request
198pub struct ReqwestRequest<B> {
199    endpoint: RspamdEndpoint<'static>,
200    client: RspamdAsyncClient,
201    body: B,
202    envelope_data: Option<EnvelopeData>,
203}
204
205#[maybe_async::maybe_async]
206impl<B: AsRef<[u8]> + Send> Request for ReqwestRequest<B> {
207    type Body = Bytes;
208    type HeaderMap = reqwest::header::HeaderMap;
209
210    async fn response(mut self) -> Result<(Self::HeaderMap, Self::Body), RspamdError> {
211        let mut retry_cnt = self.client.config.retries;
212        let mut maybe_sk = Default::default();
213        let extra_hdrs: Vec<(String, String)> =
214            self.envelope_data.take().unwrap().into_iter().collect();
215
216        let response = loop {
217            // Check if File header is present - if so, we don't need to send the body
218            let has_file_header = extra_hdrs.iter().any(|(k, _)| k == "File");
219            let need_body = self.endpoint.need_body && !has_file_header;
220            let method = if need_body {
221                reqwest::Method::POST
222            } else {
223                reqwest::Method::GET
224            };
225
226            let mut url = Url::from_str(self.client.config.base_url.as_str())
227                .map_err(|e| RspamdError::HttpError(e.to_string()))?;
228            url.set_path(self.endpoint.url);
229            let mut req = self.client.inner.request(method, url.clone());
230
231            if let Some(ref password) = self.client.config.password {
232                req = req.header("Password", password);
233            }
234
235            if self.client.config.zstd && need_body {
236                req = req.header("Content-Encoding", "zstd");
237                req = req.header("Compression", "zstd");
238            }
239
240            for (k, v) in extra_hdrs.iter() {
241                req = req.header(k, v);
242            }
243
244            if let Some(ref crypt) = self.client.crypt {
245                let inner_req = req
246                    .build()
247                    .map_err(|e| RspamdError::HttpError(e.to_string()))?;
248                let body = if need_body {
249                    if self.client.config.zstd {
250                        zstd::encode_all(self.body.as_ref(), 0)?
251                    } else {
252                        self.body.as_ref().to_vec()
253                    }
254                } else {
255                    Vec::new()
256                };
257                let encrypted = crypt.encrypt(url.path(), body.as_slice(), inner_req.headers())?;
258                req = self.client.inner.request(reqwest::Method::POST, url);
259                req = req.header("Key", crypt.key_header(encrypted.peer_key.as_str()));
260                req = req.body(encrypted.body);
261                maybe_sk = Some(encrypted.shared_key);
262            } else if need_body {
263                req = if self.client.config.zstd {
264                    req.body(reqwest::Body::from(zstd::encode_all(
265                        self.body.as_ref(),
266                        0,
267                    )?))
268                } else {
269                    req.body(Bytes::copy_from_slice(self.body.as_ref()))
270                };
271            }
272
273            let req = req.timeout(self.client.config.timeout);
274            let req = req
275                .build()
276                .map_err(|e| RspamdError::HttpError(e.to_string()))?;
277
278            match self.client.inner.execute(req).await {
279                Ok(v) => break Ok(v),
280                Err(e) => {
281                    if (retry_cnt - 1) == 0 {
282                        break Err(e);
283                    }
284                    retry_cnt -= 1;
285                    tokio::time::sleep(self.client.config.timeout).await;
286                    continue;
287                }
288            };
289        }
290        .map_err(|e| RspamdError::HttpError(e.to_string()))?;
291
292        if !response.status().is_success() {
293            return Err(RspamdError::HttpError(format!(
294                "Status: {}",
295                response.status()
296            )));
297        }
298
299        if let Some(sk) = maybe_sk {
300            let mut body = BytesMut::from(
301                response
302                    .bytes()
303                    .await
304                    .map_err(|e| RspamdError::HttpError(e.to_string()))?,
305            );
306            let decrypted_offset = httpcrypt_decrypt(body.as_mut(), sk)?;
307            let mut hdrs = [httparse::EMPTY_HEADER; 64];
308            let mut parsed = httparse::Response::new(&mut hdrs);
309
310            let body_offset = parsed
311                .parse(&body[decrypted_offset..])
312                .map_err(|s| RspamdError::HttpError(s.to_string()))?;
313            let mut output_hdrs = reqwest::header::HeaderMap::with_capacity(parsed.headers.len());
314            for hdr in parsed.headers.iter_mut() {
315                output_hdrs.insert(
316                    HeaderName::from_str(hdr.name)?,
317                    HeaderValue::from_str(std::str::from_utf8(hdr.value)?)?,
318                );
319            }
320            let body = if output_hdrs
321                .get("Compression")
322                .is_some_and(|hv| hv == "zstd")
323            {
324                zstd::decode_all(&body[body_offset.unwrap() + decrypted_offset..])?
325            } else {
326                body[body_offset.unwrap() + decrypted_offset..].to_vec()
327            };
328            Ok((output_hdrs, body.into()))
329        } else {
330            Ok((response.headers().clone(), response.bytes().await?))
331        }
332    }
333}
334
335#[maybe_async::maybe_async]
336impl<B: AsRef<[u8]> + Send> ReqwestRequest<B> {
337    pub async fn new(
338        client: RspamdAsyncClient,
339        body: B,
340        command: RspamdCommand,
341        envelope_data: EnvelopeData,
342    ) -> Result<ReqwestRequest<B>, RspamdError> {
343        Ok(Self {
344            endpoint: RspamdEndpoint::from_command(command),
345            client,
346            body,
347            envelope_data: Some(envelope_data),
348        })
349    }
350}
351
352/// Scan an email asynchronously, returning the parsed reply or error.
353///
354/// This is a one-shot helper: it builds a new client (and therefore a new
355/// connection) on every call. Prefer [`RspamdAsyncClient`] to reuse
356/// connections across scans.
357///
358/// Example:
359/// ```rust,no_run
360/// use rspamd_client::config::Config;
361/// use rspamd_client::scan_async;
362/// use rspamd_client::error::RspamdError;
363///
364///	#[tokio::main]
365/// async fn main() -> Result<(), RspamdError> {
366/// 	let config = Config::builder()
367/// 		.base_url("http://localhost:11333".to_string())
368/// 		.build();
369/// 	let envelope = Default::default();
370/// 	let email = "...";
371/// 	let response = scan_async(&config, email, envelope).await?;
372/// 	Ok(())
373/// }
374/// ```
375#[deprecated(since = "0.7.0", note = "use RspamdAsyncClient for connection pooling")]
376#[maybe_async::maybe_async]
377pub async fn scan_async<B: AsRef<[u8]> + Send>(
378    options: &Config,
379    body: B,
380    envelope_data: EnvelopeData,
381) -> Result<RspamdScanReply, RspamdError> {
382    let client = RspamdAsyncClient::new(options.clone())?;
383    client.scan(body, envelope_data).await
384}