Skip to main content

leyline/core/robots/
fetch.rs

1use http::StatusCode;
2
3use super::{ROBOTS_PATH, Robots, RobotsSource};
4use crate::core::config::RedirectPolicy;
5use crate::core::error::Result;
6use crate::core::into_url::IntoUrl;
7use crate::core::retry::RetryPolicy;
8use crate::core::session::Session;
9
10const MAX_REDIRECTS: usize = 5;
11const MAX_BYTES: usize = 512 * 1024;
12
13impl Session {
14    pub async fn robots(&self, url: impl IntoUrl, agent: &str) -> Result<Robots> {
15        self.robots_with(url, agent, RetryPolicy::transient(), None)
16            .await
17    }
18
19    pub(crate) async fn robots_with(
20        &self,
21        url: impl IntoUrl,
22        agent: &str,
23        retry: RetryPolicy,
24        trusted: Option<url::Url>,
25    ) -> Result<Robots> {
26        let target = url.into_url_with_base(self.base_url())?.join(ROBOTS_PATH)?;
27        let mut builder = self
28            .get(target.as_str())
29            .redirect(RedirectPolicy::limited(MAX_REDIRECTS))
30            .retry(retry);
31        builder.trusted_origin = trusted;
32        let response = match builder.stream().send().await {
33            Ok(response) => response,
34            Err(err) if err.is_shut_down() => return Err(err),
35            Err(_) => return Ok(Robots::empty(RobotsSource::Unreachable)),
36        };
37        let status = response.status();
38        if status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS {
39            return Ok(Robots::empty(RobotsSource::Unreachable));
40        }
41        if !status.is_success() {
42            return Ok(Robots::empty(RobotsSource::Unavailable));
43        }
44        let body = response.read_until(MAX_BYTES, |_, _| false);
45        match self.deadline(None).total(body).await {
46            Ok(bytes) => {
47                let body = String::from_utf8_lossy(&bytes);
48                let whole = match bytes.len() >= MAX_BYTES {
49                    true => body.rfind('\n').map_or("", |end| &body[..end]),
50                    false => &body,
51                };
52                Ok(Robots::parse(whole, agent))
53            }
54            Err(err) if err.is_shut_down() => Err(err),
55            Err(_) => Ok(Robots::empty(RobotsSource::Unreachable)),
56        }
57    }
58}