leyline/core/
pages.rs

1use std::collections::HashSet;
2use std::future::Future;
3use std::pin::Pin;
4use std::task::{Context, Poll, ready};
5
6use futures_util::Stream;
7use http::Method;
8use url::Url;
9
10use crate::core::config::{Origin, ProxyConfig, RedirectPolicy, TimeoutConfig};
11use crate::core::digest::DigestAuth;
12use crate::core::error::Result;
13use crate::core::headers::HeaderList;
14use crate::core::request::RequestBuilder;
15use crate::core::response::Response;
16use crate::core::retry::RetryPolicy;
17use crate::core::session::Session;
18use crate::profile::Preset;
19use crate::util::sensitive_header;
20
21const NEXT: &str = "next";
22
23type PageFuture = Pin<Box<dyn Future<Output = Result<Response>> + Send>>;
24
25#[must_use = "pages are lazy: nothing happens until `.next()` is awaited"]
26pub struct Pages {
27    template: Template,
28    pending: Option<PageFuture>,
29    visited: HashSet<Url>,
30    remaining: Option<usize>,
31}
32
33struct Template {
34    session: Session,
35    first: Option<Url>,
36    preset: Option<Preset>,
37    preset_user: bool,
38    headers: HeaderList,
39    timeouts: Option<TimeoutConfig>,
40    stream_response: bool,
41    retry_policy: RetryPolicy,
42    digest_auth: Option<DigestAuth>,
43    proxy: Option<ProxyConfig>,
44    header_order: Option<Vec<String>>,
45    redirect: Option<RedirectPolicy>,
46    initiator: Option<Url>,
47    tag: Option<String>,
48    status_errors: bool,
49}
50
51impl RequestBuilder {
52    pub fn pages(self) -> Pages {
53        let template = Template::of(&self);
54        let visited = template.first.iter().cloned().collect();
55        Pages {
56            template,
57            pending: Some(Box::pin(self.send())),
58            visited,
59            remaining: None,
60        }
61    }
62}
63
64impl Template {
65    fn of(first: &RequestBuilder) -> Self {
66        Self {
67            session: first.session.clone(),
68            first: Url::parse(&first.url).ok(),
69            preset: first.preset,
70            preset_user: first.preset_user,
71            headers: first.headers.clone(),
72            timeouts: first.timeouts,
73            stream_response: first.stream_response,
74            retry_policy: first.retry_policy.clone(),
75            digest_auth: first.digest_auth.clone(),
76            proxy: first.proxy.clone(),
77            header_order: first.header_order.clone(),
78            redirect: first.redirect.clone(),
79            initiator: first.initiator.clone(),
80            tag: first.tag.clone(),
81            status_errors: first.status_errors,
82        }
83    }
84
85    fn downgrades(&self, url: &Url) -> bool {
86        self.first
87            .as_ref()
88            .is_some_and(|first| first.scheme() == "https" && url.scheme() == "http")
89    }
90
91    fn same_origin(&self, url: &Url) -> bool {
92        let first = self.first.as_ref().and_then(Origin::of);
93        first.is_some() && first == Origin::of(url)
94    }
95
96    fn follow(&self, url: &Url) -> RequestBuilder {
97        let mut next = RequestBuilder::new(&self.session, Method::GET, url.as_str());
98        if self.preset_user {
99            next.preset = self.preset;
100            next.preset_user = true;
101        }
102        next.headers = self.headers.clone();
103        next.timeouts = self.timeouts;
104        next.stream_response = self.stream_response;
105        next.retry_policy = self.retry_policy.clone();
106        if self.same_origin(url) {
107            next.digest_auth = self.digest_auth.clone();
108        } else {
109            next.headers
110                .remove_where(|name| sensitive_header(name.as_str()));
111        }
112        next.proxy = self.proxy.clone();
113        next.header_order = self.header_order.clone();
114        next.redirect = self.redirect.clone();
115        next.initiator = self.initiator.clone();
116        next.tag = self.tag.clone();
117        next.status_errors = self.status_errors;
118        next.trusted_origin = self.first.clone();
119        next
120    }
121}
122
123impl Pages {
124    pub fn limit(mut self, max: usize) -> Self {
125        self.remaining = Some(max);
126        if max == 0 {
127            self.pending = None;
128        }
129        self
130    }
131
132    pub async fn next(&mut self) -> Option<Result<Response>> {
133        std::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
134    }
135
136    fn queue_next(&mut self, response: &Response) {
137        self.visited.insert(response.url().clone());
138        let Some(url) = response.link(NEXT) else {
139            return;
140        };
141        if self.template.downgrades(&url) {
142            return;
143        }
144        if self.visited.insert(url.clone()) {
145            self.pending = Some(Box::pin(self.template.follow(&url).send()));
146        }
147    }
148}
149
150impl Stream for Pages {
151    type Item = Result<Response>;
152
153    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
154        let this = self.get_mut();
155        let Some(pending) = this.pending.as_mut() else {
156            return Poll::Ready(None);
157        };
158        let result = ready!(pending.as_mut().poll(cx));
159        this.pending = None;
160        if let Some(remaining) = this.remaining.as_mut() {
161            *remaining = remaining.saturating_sub(1);
162        }
163        if let (Ok(response), false) = (&result, this.remaining == Some(0)) {
164            this.queue_next(response);
165        }
166        Poll::Ready(Some(result))
167    }
168}
169
170impl std::fmt::Debug for Pages {
171    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172        f.debug_struct("Pages")
173            .field("visited", &self.visited.len())
174            .field("done", &self.pending.is_none())
175            .finish()
176    }
177}
leyline/core/pages.rs

leyline/core/
pages.rs
