spider_lib/
middleware.rs

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    // Modified to return Result
25    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}