use std::sync::{Arc, Mutex};
use super::{Action, Attempt, Policy};
#[derive(Clone, Debug)]
pub struct Limited {
remaining: Arc<Mutex<usize>>,
}
impl Limited {
pub fn new(max: usize) -> Self {
Limited {
remaining: Arc::new(Mutex::new(max)),
}
}
}
impl Default for Limited {
fn default() -> Self {
Limited::new(20)
}
}
impl<B, E> Policy<B, E> for Limited {
fn redirect(&self, _: &Attempt<'_>) -> Result<Action, E> {
let mut remaining = self.remaining.lock().unwrap();
if *remaining > 0 {
*remaining -= 1;
Ok(Action::Follow)
} else {
Ok(Action::Stop)
}
}
}
#[cfg(test)]
mod tests {
use http::{Request, Uri};
use super::*;
#[test]
fn works() {
let uri = Uri::from_static("https://example.com/");
let policy = Limited::new(2);
for _ in 0..2 {
let mut request = Request::builder().uri(uri.clone()).body(()).unwrap();
Policy::<(), ()>::on_request(&policy, &mut request);
let attempt = Attempt {
status: Default::default(),
location: &uri,
previous: &uri,
};
assert!(Policy::<(), ()>::redirect(&policy, &attempt)
.unwrap()
.is_follow());
}
let mut request = Request::builder().uri(uri.clone()).body(()).unwrap();
Policy::<(), ()>::on_request(&policy, &mut request);
let attempt = Attempt {
status: Default::default(),
location: &uri,
previous: &uri,
};
assert!(Policy::<(), ()>::redirect(&policy, &attempt)
.unwrap()
.is_stop());
}
}