spider_lib/middlewares/

user_agent.rs

use async_trait::async_trait;
use dashmap::DashMap;
use moka::sync::Cache;
use reqwest::header::{HeaderValue, USER_AGENT};
use serde::{Deserialize, Serialize};
use std::fmt::Debug;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tracing::{debug, info, warn};

use rand::seq::SliceRandom;

use crate::error::SpiderError;
use crate::middleware::{Middleware, MiddlewareAction};
use crate::request::Request;

/// Defines the strategy for rotating User-Agents.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum UserAgentRotationStrategy {
    /// Randomly selects a User-Agent from the available pool.
    #[default]
    Random,
    /// Sequentially cycles through the available User-Agents.
    Sequential,
}

/// Predefined lists of User-Agents for common scenarios.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum BuiltinUserAgentList {
    /// A list of popular desktop browser User-Agents.
    Desktop,
    /// A list of popular mobile browser User-Agents.
    Mobile,
    /// A mixed list of popular desktop and mobile browser User-Agents.
    RandomPopular,
}

/// Defines the source from which User-Agents are loaded.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)] // Allow flexible deserialization
pub enum UserAgentSource {
    /// A direct list of User-Agent strings.
    List(Vec<String>),
    /// Path to a file containing User-Agent strings, one per line.
    File(PathBuf),
    /// Use a predefined, built-in list of User-Agents.
    Builtin(BuiltinUserAgentList),
    /// No User-Agent source specified, will fallback to a default if available.
    None,
}

impl Default for UserAgentSource {
    fn default() -> Self {
        // Default to a builtin popular list for convenience
        UserAgentSource::Builtin(BuiltinUserAgentList::RandomPopular)
    }
}

/// Represents a User-Agent profile, including the User-Agent string and other associated headers.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserAgentProfile {
    /// The User-Agent string.
    pub user_agent: Arc<String>,
    /// Additional headers that should be sent with this User-Agent to mimic a real browser.
    #[serde(default)]
    pub headers: DashMap<String, String>,
}

impl From<String> for UserAgentProfile {
    fn from(user_agent: String) -> Self {
        UserAgentProfile {
            user_agent: Arc::new(user_agent),
            headers: DashMap::new(),
        }
    }
}

impl From<&str> for UserAgentProfile {
    fn from(user_agent: &str) -> Self {
        UserAgentProfile {
            user_agent: Arc::new(user_agent.to_string()),
            headers: DashMap::new(),
        }
    }
}

/// Builder for creating a `UserAgentMiddleware`.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct UserAgentMiddlewareBuilder {
    source: UserAgentSource,
    strategy: UserAgentRotationStrategy,
    fallback_user_agent: Option<String>,
    per_domain: DashMap<String, UserAgentSource>,
}

impl UserAgentMiddlewareBuilder {
    /// Sets the primary source for User-Agents.
    pub fn source(mut self, source: UserAgentSource) -> Self {
        self.source = source;
        self
    }

    /// Sets the strategy to use for rotating User-Agents.
    pub fn strategy(mut self, strategy: UserAgentRotationStrategy) -> Self {
        self.strategy = strategy;
        self
    }

    /// Sets a fallback User-Agent to use if no other User-Agents are available.
    pub fn fallback_user_agent(mut self, fallback_user_agent: String) -> Self {
        self.fallback_user_agent = Some(fallback_user_agent);
        self
    }

    /// Adds a domain-specific User-Agent source.
    pub fn per_domain(self, domain: String, source: UserAgentSource) -> Self {
        self.per_domain.insert(domain, source);
        self
    }

    /// Builds the `UserAgentMiddleware`.
    /// This can fail if a User-Agent source file is specified but cannot be read.
    pub fn build(self) -> Result<UserAgentMiddleware, SpiderError> {
        let default_pool = Arc::new(UserAgentMiddleware::load_user_agents(&self.source)?);

        let domain_cache = Cache::builder()
            .time_to_live(Duration::from_secs(30 * 60)) // 30 minutes
            .build();

        for entry in self.per_domain.iter() {
            let domain = entry.key().clone();
            let source = entry.value().clone();
            let pool = Arc::new(UserAgentMiddleware::load_user_agents(&source)?);
            domain_cache.insert(domain, pool);
        }

        let middleware = UserAgentMiddleware {
            strategy: self.strategy,
            fallback_user_agent: self.fallback_user_agent,
            domain_cache,
            default_pool,
            current_index: AtomicUsize::new(0),
        };

        info!(
            "Initializing UserAgentMiddleware with config: {:?}",
            middleware
        );

        Ok(middleware)
    }
}

pub struct UserAgentMiddleware {
    strategy: UserAgentRotationStrategy,
    fallback_user_agent: Option<String>,
    domain_cache: Cache<String, Arc<Vec<UserAgentProfile>>>,
    default_pool: Arc<Vec<UserAgentProfile>>,
    current_index: AtomicUsize,
}

impl Debug for UserAgentMiddleware {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UserAgentMiddleware")
            .field("strategy", &self.strategy)
            .field("fallback_user_agent", &self.fallback_user_agent)
            .field(
                "domain_cache",
                &format!("Cache({})", self.domain_cache.weighted_size()),
            )
            .field(
                "default_pool",
                &format!("Pool({})", self.default_pool.len()),
            )
            .field("current_index", &self.current_index)
            .finish()
    }
}

impl UserAgentMiddleware {
    /// Creates a new `UserAgentMiddlewareBuilder` to start building a `UserAgentMiddleware`.
    pub fn builder() -> UserAgentMiddlewareBuilder {
        UserAgentMiddlewareBuilder::default()
    }

    fn load_user_agents(source: &UserAgentSource) -> Result<Vec<UserAgentProfile>, SpiderError> {
        match source {
            UserAgentSource::List(list) => Ok(list
                .iter()
                .map(|ua| UserAgentProfile::from(ua.clone()))
                .collect()),
            UserAgentSource::File(path) => Self::load_from_file(path),
            UserAgentSource::Builtin(builtin_list) => {
                Ok(Self::load_builtin_user_agents(builtin_list))
            }
            UserAgentSource::None => Ok(Vec::new()),
        }
    }

    fn load_from_file(path: &Path) -> Result<Vec<UserAgentProfile>, SpiderError> {
        if !path.exists() {
            return Err(SpiderError::IoError(std::io::Error::new(
                std::io::ErrorKind::NotFound,
                format!("User-agent file not found: {}", path.display()),
            )));
        }
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let user_agents: Vec<UserAgentProfile> = reader
            .lines()
            .map_while(Result::ok)
            .filter(|line| !line.trim().is_empty())
            .map(UserAgentProfile::from)
            .collect();

        if user_agents.is_empty() {
            warn!(
                "User-Agent file {:?} is empty or contains no valid User-Agents.",
                path
            );
        }
        Ok(user_agents)
    }

    //TODO: provide a list of user agents from third parties
    fn load_builtin_user_agents(list_type: &BuiltinUserAgentList) -> Vec<UserAgentProfile> {
        match list_type {
            BuiltinUserAgentList::Desktop => vec![
                UserAgentProfile::from(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
                ),
                UserAgentProfile::from(
                    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
                ),
                UserAgentProfile::from(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/121.0",
                ),
            ],
            BuiltinUserAgentList::Mobile => vec![
                UserAgentProfile::from(
                    "Mozilla/5.0 (Linux; Android 10) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.6099.193 Mobile Safari/537.36",
                ),
                UserAgentProfile::from(
                    "Mozilla/5.0 (iPhone; CPU iPhone OS 17_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Mobile/15E148 Safari/604.1",
                ),
            ],
            BuiltinUserAgentList::RandomPopular => vec![
                UserAgentProfile::from(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
                ),
                UserAgentProfile::from(
                    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
                ),
                UserAgentProfile::from(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/121.0",
                ),
                UserAgentProfile::from(
                    "Mozilla/5.0 (Linux; Android 10) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.6099.193 Mobile Safari/537.36",
                ),
                UserAgentProfile::from(
                    "Mozilla/5.0 (iPhone; CPU iPhone OS 17_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Mobile/15E148 Safari/604.1",
                ),
            ],
        }
    }

    fn get_user_agent(&self, domain: Option<&str>) -> Option<UserAgentProfile> {
        let mut rng = rand::thread_rng();

        let pool = domain
            .and_then(|d| self.domain_cache.get(d))
            .unwrap_or_else(|| self.default_pool.clone());

        if pool.is_empty() {
            debug!("User-Agent pool is empty.");
            return self
                .fallback_user_agent
                .as_ref()
                .map(|ua| UserAgentProfile::from(ua.clone()));
        }

        match self.strategy {
            UserAgentRotationStrategy::Random => pool.choose(&mut rng).cloned(),
            UserAgentRotationStrategy::Sequential => {
                let current = self.current_index.fetch_add(1, Ordering::SeqCst);
                let index = current % pool.len();
                pool.get(index).cloned()
            }
        }
    }
}

#[async_trait]
impl<C: Send + Sync> Middleware<C> for UserAgentMiddleware {
    fn name(&self) -> &str {
        "UserAgentMiddleware"
    }

    async fn process_request(
        &mut self,
        _client: &C,
        mut request: Request,
    ) -> Result<MiddlewareAction<Request>, SpiderError> {
        let domain = request.url.domain();
        if let Some(profile) = self.get_user_agent(domain) {
            debug!("Applying User-Agent: {}", profile.user_agent);
            request.headers.insert(
                USER_AGENT,
                HeaderValue::from_str(&profile.user_agent).map_err(|e| {
                    SpiderError::HeaderValueError(format!(
                        "Invalid User-Agent string '{}': {}",
                        profile.user_agent, e
                    ))
                })?,
            );
            for header in profile.headers.iter() {
                request.headers.insert(
                    reqwest::header::HeaderName::from_bytes(header.key().as_bytes()).map_err(
                        |e| SpiderError::HeaderValueError(format!("Invalid header name: {}", e)),
                    )?,
                    HeaderValue::from_str(header.value().as_str()).map_err(|e| {
                        SpiderError::HeaderValueError(format!(
                            "Invalid header value for {}: {}",
                            header.key(),
                            e
                        ))
                    })?,
                );
            }
        } else {
            debug!("No User-Agent applied.");
        }
        Ok(MiddlewareAction::Continue(request))
    }
}