fast_pull/reqwest/
prefetch.rs

1extern crate alloc;
2use crate::{UrlInfo, reqwest::ReqwestReader};
3use alloc::string::{String, ToString};
4use content_disposition;
5use reqwest::{
6    Client, IntoUrl, StatusCode, Url,
7    header::{self, HeaderMap},
8};
9use sanitize_filename;
10
11fn get_file_size(headers: &HeaderMap, status: &StatusCode) -> u64 {
12    if *status == StatusCode::PARTIAL_CONTENT {
13        headers
14            .get(header::CONTENT_RANGE)
15            .and_then(|hv| hv.to_str().ok())
16            .and_then(|s| s.rsplit('/').next())
17            .and_then(|total| total.parse().ok())
18            .unwrap_or(0)
19    } else {
20        headers
21            .get(header::CONTENT_LENGTH)
22            .and_then(|hv| hv.to_str().ok())
23            .and_then(|s| s.parse().ok())
24            .unwrap_or(0)
25    }
26}
27
28fn get_header_str(headers: &HeaderMap, header_name: &header::HeaderName) -> Option<String> {
29    headers
30        .get(header_name)
31        .and_then(|hv| hv.to_str().ok())
32        .map(String::from)
33}
34
35fn get_filename(headers: &HeaderMap, final_url: &Url) -> String {
36    let from_disposition = headers
37        .get(header::CONTENT_DISPOSITION)
38        .and_then(|hv| hv.to_str().ok())
39        .and_then(|s| content_disposition::parse_content_disposition(s).filename_full())
40        .filter(|s| !s.trim().is_empty());
41
42    let from_url = final_url
43        .path_segments()
44        .and_then(|mut segments| segments.next_back())
45        .and_then(|s| urlencoding::decode(s).ok())
46        .filter(|s| !s.trim().is_empty())
47        .map(|s| ToString::to_string(&s));
48
49    let raw_name = from_disposition
50        .or(from_url)
51        .unwrap_or_else(|| final_url.to_string());
52
53    sanitize_filename::sanitize_with_options(
54        &raw_name,
55        sanitize_filename::Options {
56            windows: true,
57            truncate: true,
58            replacement: "_",
59        },
60    )
61}
62
63pub trait Prefetch {
64    fn prefetch(
65        &self,
66        url: impl IntoUrl + Send,
67    ) -> impl Future<Output = Result<UrlInfo, reqwest::Error>> + Send;
68}
69
70impl Prefetch for Client {
71    async fn prefetch(&self, url: impl IntoUrl + Send) -> Result<UrlInfo, reqwest::Error> {
72        let url = url.into_url()?;
73        let resp = self.head(url.clone()).send().await?;
74        let resp = match resp.error_for_status() {
75            Ok(resp) => resp,
76            Err(_) => return prefetch_fallback(url, self).await,
77        };
78
79        let status = resp.status();
80        let final_url = resp.url();
81
82        let resp_headers = resp.headers();
83        let size = get_file_size(resp_headers, &status);
84
85        let supports_range = match resp.headers().get(header::ACCEPT_RANGES) {
86            Some(accept_ranges) => accept_ranges
87                .to_str()
88                .ok()
89                .map(|v| v.split(' '))
90                .and_then(|supports| supports.into_iter().find(|&ty| ty == "bytes"))
91                .is_some(),
92            None => return prefetch_fallback(url, self).await,
93        };
94
95        Ok(UrlInfo {
96            final_url: final_url.clone(),
97            name: get_filename(resp_headers, final_url),
98            size,
99            supports_range,
100            fast_download: size > 0 && supports_range,
101            etag: get_header_str(resp_headers, &header::ETAG),
102            last_modified: get_header_str(resp_headers, &header::LAST_MODIFIED),
103        })
104    }
105}
106
107impl Prefetch for ReqwestReader {
108    fn prefetch(
109        &self,
110        url: impl IntoUrl + Send,
111    ) -> impl Future<Output = Result<UrlInfo, reqwest::Error>> + Send {
112        self.client.prefetch(url)
113    }
114}
115
116async fn prefetch_fallback(url: Url, client: &Client) -> Result<UrlInfo, reqwest::Error> {
117    let resp = client
118        .get(url)
119        .header(header::RANGE, "bytes=0-")
120        .send()
121        .await?
122        .error_for_status()?;
123    let status = resp.status();
124    let final_url = resp.url();
125
126    let resp_headers = resp.headers();
127    let size = get_file_size(resp_headers, &status);
128    let supports_range = status == StatusCode::PARTIAL_CONTENT;
129    Ok(UrlInfo {
130        final_url: final_url.clone(),
131        name: get_filename(resp_headers, final_url),
132        size,
133        supports_range,
134        fast_download: size > 0 && supports_range,
135        etag: get_header_str(resp_headers, &header::ETAG),
136        last_modified: get_header_str(resp_headers, &header::LAST_MODIFIED),
137    })
138}
139
140#[cfg(test)]
141mod tests {
142    use alloc::{format, vec};
143
144    use super::*;
145
146    #[tokio::test]
147    async fn test_redirect_and_content_range() {
148        let mut server = mockito::Server::new_async().await;
149
150        let mock_redirect = server
151            .mock("GET", "/redirect")
152            .with_status(301)
153            .with_header("Location", "/real-file.txt")
154            .create_async()
155            .await;
156
157        let mock_file = server
158            .mock("GET", "/real-file.txt")
159            .with_status(206)
160            .with_header("Content-Range", "bytes 0-1023/2048")
161            .with_body(vec![0; 1024])
162            .create_async()
163            .await;
164
165        let client = Client::new();
166        let url_info = client
167            .prefetch(&format!("{}/redirect", server.url()))
168            .await
169            .expect("Request should succeed");
170
171        assert_eq!(url_info.size, 2048);
172        assert_eq!(url_info.name, "real-file.txt");
173        assert_eq!(
174            url_info.final_url.as_str(),
175            format!("{}/real-file.txt", server.url())
176        );
177        assert!(url_info.supports_range);
178
179        mock_redirect.assert_async().await;
180        mock_file.assert_async().await;
181    }
182
183    #[tokio::test]
184    async fn test_content_range_priority() {
185        let mut server = mockito::Server::new_async().await;
186        let mock = server
187            .mock("GET", "/file")
188            .with_status(206)
189            .with_header("Content-Range", "bytes 0-1023/2048")
190            .create_async()
191            .await;
192
193        let client = Client::new();
194        let url_info = client
195            .prefetch(&format!("{}/file", server.url()))
196            .await
197            .expect("Request should succeed");
198
199        assert_eq!(url_info.size, 2048);
200        mock.assert_async().await;
201    }
202
203    #[tokio::test]
204    async fn test_filename_sources() {
205        let mut server = mockito::Server::new_async().await;
206
207        // Test Content-Disposition source
208        let mock1 = server
209            .mock("GET", "/test1")
210            .with_header("Content-Disposition", "attachment; filename=\"test.txt\"")
211            .create_async()
212            .await;
213        let url_info = Client::new()
214            .prefetch(&format!("{}/test1", server.url()))
215            .await
216            .unwrap();
217        assert_eq!(url_info.name, "test.txt");
218        mock1.assert_async().await;
219
220        // Test URL path source
221        let mock2 = server.mock("GET", "/test2/file.pdf").create_async().await;
222        let url_info = Client::new()
223            .prefetch(&format!("{}/test2/file.pdf", server.url()))
224            .await
225            .unwrap();
226        assert_eq!(url_info.name, "file.pdf");
227        mock2.assert_async().await;
228
229        // Test sanitization
230        let mock3 = server
231      .mock("GET", "/test3")
232      .with_header(
233        "Content-Disposition",
234        "attachment; filename*=UTF-8''%E6%82%AA%E3%81%84%3C%3E%E3%83%95%E3%82%A1%E3%82%A4%E3%83%AB%3F%E5%90%8D.txt",
235      )
236      .create_async()
237      .await;
238        let url_info = Client::new()
239            .prefetch(&format!("{}/test3", server.url()))
240            .await
241            .unwrap();
242        assert_eq!(url_info.name, "悪い__ファイル_名.txt");
243        mock3.assert_async().await;
244    }
245
246    #[tokio::test]
247    async fn test_error_handling() {
248        let mut server = mockito::Server::new_async().await;
249        let mock1 = server
250            .mock("GET", "/404")
251            .with_status(404)
252            .create_async()
253            .await;
254
255        let client = Client::new();
256
257        match client.prefetch(&format!("{}/404", server.url())).await {
258            Ok(info) => panic!("404 status code should not success: {info:?}"),
259            Err(err) => {
260                assert!(err.is_status(), "should be error about status code");
261                assert_eq!(err.status(), Some(StatusCode::NOT_FOUND));
262            }
263        }
264
265        mock1.assert_async().await;
266    }
267}
fast_pull/reqwest/prefetch.rs

fast_pull/reqwest/
prefetch.rs
