use async_trait::async_trait;
use crate::error::KumoError;
use super::{FetchRequest, Middleware, RotationStrategy};
pub struct ProxyRotator {
proxies: Vec<String>,
strategy: RotationStrategy,
}
impl ProxyRotator {
pub fn new(proxies: Vec<impl Into<String>>) -> Self {
Self {
proxies: proxies.into_iter().map(Into::into).collect(),
strategy: RotationStrategy::round_robin(),
}
}
pub fn random(proxies: Vec<impl Into<String>>) -> Self {
Self {
proxies: proxies.into_iter().map(Into::into).collect(),
strategy: RotationStrategy::random(),
}
}
fn pick(&self) -> Option<&str> {
if self.proxies.is_empty() {
return None;
}
Some(&self.proxies[self.strategy.pick_index(self.proxies.len())])
}
}
#[async_trait]
impl Middleware for ProxyRotator {
async fn before_request(&self, request: &mut FetchRequest) -> Result<(), KumoError> {
if let Some(proxy) = self.pick() {
request.proxy = Some(proxy.to_string());
}
Ok(())
}
}