1use async_trait::async_trait;
2use reqwest::header::{HeaderValue, USER_AGENT};
3use tracing::{debug, instrument};
4
5use crate::error::SpiderError;
6use crate::request::Request;
7use crate::response::Response;
8
9#[async_trait]
10pub trait Middleware: Send + Sync + 'static {
11 async fn process_request(&self, request: Request) -> Result<Request, SpiderError> {
12 Ok(request)
13 }
14 async fn process_response(&self, response: Response) -> Result<Response, SpiderError> {
15 Ok(response)
16 }
17}
18
19pub struct UserAgentMiddleware {
20 user_agent_value: HeaderValue,
21}
22
23impl UserAgentMiddleware {
24 pub fn new(user_agent: &str) -> Result<Self, SpiderError> {
26 let user_agent_value = HeaderValue::from_str(user_agent).map_err(|e| {
27 SpiderError::HeaderValueError(format!(
28 "Invalid User-Agent string '{}': {}",
29 user_agent, e
30 ))
31 })?;
32 Ok(UserAgentMiddleware { user_agent_value })
33 }
34}
35
36#[async_trait]
37impl Middleware for UserAgentMiddleware {
38 #[instrument(skip(self, request))]
39 async fn process_request(&self, mut request: Request) -> Result<Request, SpiderError> {
40 debug!("Applying User-Agent: {:?}", self.user_agent_value);
41 request
42 .headers
43 .insert(USER_AGENT, self.user_agent_value.clone());
44 Ok(request)
45 }
46}