tower_http/follow_redirect/policy/
limited.rs1use super::{Action, Attempt, Policy};
2
3#[derive(Clone, Copy, Debug)]
5pub struct Limited {
6 remaining: usize,
7}
8
9impl Limited {
10 pub fn new(max: usize) -> Self {
12 Limited { remaining: max }
13 }
14}
15
16impl Default for Limited {
17 fn default() -> Self {
19 Limited::new(20)
23 }
24}
25
26impl<B, E> Policy<B, E> for Limited {
27 fn redirect(&mut self, _: &Attempt<'_>) -> Result<Action, E> {
28 if self.remaining > 0 {
29 self.remaining -= 1;
30 Ok(Action::Follow)
31 } else {
32 Ok(Action::Stop)
33 }
34 }
35}
36
37#[cfg(test)]
38mod tests {
39 use http::{Request, Uri};
40
41 use super::*;
42
43 #[test]
44 fn works() {
45 let uri = Uri::from_static("https://example.com/");
46 let mut policy = Limited::new(2);
47
48 for _ in 0..2 {
49 let mut request = Request::builder().uri(uri.clone()).body(()).unwrap();
50 Policy::<(), ()>::on_request(&mut policy, &mut request);
51
52 let attempt = Attempt {
53 status: Default::default(),
54 location: &uri,
55 previous: &uri,
56 };
57 assert!(Policy::<(), ()>::redirect(&mut policy, &attempt)
58 .unwrap()
59 .is_follow());
60 }
61
62 let mut request = Request::builder().uri(uri.clone()).body(()).unwrap();
63 Policy::<(), ()>::on_request(&mut policy, &mut request);
64
65 let attempt = Attempt {
66 status: Default::default(),
67 location: &uri,
68 previous: &uri,
69 };
70 assert!(Policy::<(), ()>::redirect(&mut policy, &attempt)
71 .unwrap()
72 .is_stop());
73 }
74}