use std::time::Duration;
use futures_util::StreamExt;
use reqwest::Url;
use super::tools::web::USER_AGENT;
pub const MAX_REDIRECTS: usize = 5;
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const NAME_CHARS: usize = 64;
const EXT_CHARS: usize = 10;
pub struct FetchedImage {
pub bytes: Vec<u8>,
pub content_type: String,
pub final_url: Url,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum FetchError {
#[error("only http and https URLs can be fetched")]
Scheme,
#[error("malformed URL")]
Malformed,
#[error("too many redirects")]
TooManyRedirects,
#[error("request failed: {0}")]
Request(String),
#[error("HTTP status {0}")]
Status(u16),
#[error("larger than the size limit")]
TooBig,
#[error("empty response")]
Empty,
}
pub async fn fetch(url: &str, max_bytes: u64) -> Result<FetchedImage, FetchError> {
let client = reqwest::Client::builder()
.connect_timeout(CONNECT_TIMEOUT)
.timeout(REQUEST_TIMEOUT)
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| FetchError::Request(e.to_string()))?;
let mut target = parse_http(url)?;
for _ in 0..=MAX_REDIRECTS {
let resp = client
.get(target.clone())
.header(reqwest::header::USER_AGENT, USER_AGENT)
.header(reqwest::header::ACCEPT, "image/*,*/*;q=0.8")
.send()
.await
.map_err(|e| FetchError::Request(e.to_string()))?;
let status = resp.status();
if status.is_redirection() {
let location = resp
.headers()
.get(reqwest::header::LOCATION)
.and_then(|v| v.to_str().ok())
.ok_or(FetchError::Malformed)?;
let next = target.join(location).map_err(|_| FetchError::Malformed)?;
check_scheme(&next)?;
target = next;
continue;
}
if !status.is_success() {
return Err(FetchError::Status(status.as_u16()));
}
if resp.content_length().is_some_and(|len| len > max_bytes) {
return Err(FetchError::TooBig);
}
let content_type = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.map(|ct| {
ct.split(';')
.next()
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
})
.unwrap_or_default();
let final_url = target.clone();
let mut bytes: Vec<u8> = Vec::new();
let mut stream = resp.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| FetchError::Request(e.to_string()))?;
if bytes.len() as u64 + chunk.len() as u64 > max_bytes {
return Err(FetchError::TooBig);
}
bytes.extend_from_slice(&chunk);
}
if bytes.is_empty() {
return Err(FetchError::Empty);
}
return Ok(FetchedImage {
bytes,
content_type,
final_url,
});
}
Err(FetchError::TooManyRedirects)
}
pub fn looks_like_url(arg: &str) -> bool {
let lower = arg.trim().to_ascii_lowercase();
lower.starts_with("http://") || lower.starts_with("https://")
}
pub fn display_name(url: &Url, ext: &str) -> String {
let segment = url
.path_segments()
.and_then(|mut s| s.next_back())
.unwrap_or_default();
let decoded = percent_encoding::percent_decode_str(segment).decode_utf8_lossy();
let name = decoded.trim();
if name.is_empty() {
return format!("{}.{ext}", url.host_str().unwrap_or("image"));
}
trim_name(name)
}
fn trim_name(name: &str) -> String {
if name.chars().count() <= NAME_CHARS {
return name.to_string();
}
let ext = name
.rsplit_once('.')
.map(|(_, ext)| ext)
.filter(|ext| !ext.is_empty() && ext.chars().count() <= EXT_CHARS)
.unwrap_or_default();
if ext.is_empty() {
return name.chars().take(NAME_CHARS).collect();
}
let stem: String = name
.chars()
.take(NAME_CHARS.saturating_sub(ext.chars().count() + 1))
.collect();
format!("{stem}.{ext}")
}
fn parse_http(url: &str) -> Result<Url, FetchError> {
let parsed = Url::parse(url.trim()).map_err(|_| FetchError::Malformed)?;
check_scheme(&parsed)?;
Ok(parsed)
}
fn check_scheme(url: &Url) -> Result<(), FetchError> {
match url.scheme() {
"http" | "https" => Ok(()),
_ => Err(FetchError::Scheme),
}
}
#[cfg(test)]
pub(crate) mod stub {
use std::io::{Read, Write};
use std::net::TcpListener;
pub(crate) fn serve(responses: Vec<Vec<u8>>) -> (String, std::thread::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let handle = std::thread::spawn(move || {
for response in responses {
let Ok((mut stream, _)) = listener.accept() else {
return;
};
let mut buf = [0u8; 2048];
let _ = stream.read(&mut buf);
let _ = stream.write_all(&response);
let _ = stream.flush();
}
});
(format!("http://{addr}"), handle)
}
pub(crate) fn png_bytes(width: u32, height: u32) -> Vec<u8> {
let buf = image::ImageBuffer::from_fn(width, height, |x, _| {
image::Rgb([(x % 256) as u8, 30, 60])
});
let mut out = Vec::new();
image::DynamicImage::ImageRgb8(buf)
.write_to(&mut std::io::Cursor::new(&mut out), image::ImageFormat::Png)
.unwrap();
out
}
pub(crate) fn ok_response(content_type: &str, body: &[u8], with_length: bool) -> Vec<u8> {
let mut head = format!("HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\n");
if with_length {
head.push_str(&format!("Content-Length: {}\r\n", body.len()));
}
head.push_str("Connection: close\r\n\r\n");
let mut out = head.into_bytes();
out.extend_from_slice(body);
out
}
pub(crate) fn redirect_response(location: &str) -> Vec<u8> {
format!(
"HTTP/1.1 302 Found\r\nLocation: {location}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
)
.into_bytes()
}
}
#[cfg(test)]
mod tests {
use super::stub::*;
use super::*;
#[tokio::test]
async fn downloads_an_image_and_reports_its_content_type() {
let png = png_bytes(20, 10);
let (base, _h) = serve(vec![ok_response("image/png", &png, true)]);
let fetched = fetch(&format!("{base}/pics/shot.png"), 10_000_000)
.await
.unwrap();
assert_eq!(fetched.bytes, png);
assert_eq!(fetched.content_type, "image/png");
assert_eq!(fetched.final_url.path(), "/pics/shot.png");
}
#[tokio::test]
async fn content_type_parameters_are_stripped() {
let (base, _h) = serve(vec![ok_response(
"IMAGE/PNG; charset=binary",
&png_bytes(4, 4),
true,
)]);
let fetched = fetch(&base, 10_000_000).await.unwrap();
assert_eq!(fetched.content_type, "image/png");
}
#[tokio::test]
async fn follows_a_redirect_chain_within_the_cap() {
let png = png_bytes(8, 8);
let (base, _h) = serve(vec![
redirect_response("/one"),
redirect_response("/two"),
ok_response("image/png", &png, true),
]);
let fetched = fetch(&base, 10_000_000).await.unwrap();
assert_eq!(fetched.bytes, png);
assert_eq!(fetched.final_url.path(), "/two");
}
#[tokio::test]
async fn a_chain_over_the_cap_is_refused() {
let responses = (0..=MAX_REDIRECTS + 1)
.map(|i| redirect_response(&format!("/hop{i}")))
.collect();
let (base, _h) = serve(responses);
assert_eq!(
fetch(&base, 10_000_000).await.err(),
Some(FetchError::TooManyRedirects)
);
}
#[tokio::test]
async fn a_redirect_out_of_http_is_refused() {
let (base, _h) = serve(vec![redirect_response("file:///etc/passwd")]);
assert_eq!(
fetch(&base, 10_000_000).await.err(),
Some(FetchError::Scheme)
);
}
#[tokio::test]
async fn a_non_http_url_is_refused_before_any_request() {
assert_eq!(
fetch("file:///C:/secrets/id_rsa", 10_000_000).await.err(),
Some(FetchError::Scheme)
);
assert_eq!(
fetch("ftp://example.com/a.png", 10_000_000).await.err(),
Some(FetchError::Scheme)
);
assert_eq!(
fetch("not a url", 10_000_000).await.err(),
Some(FetchError::Malformed)
);
}
#[tokio::test]
async fn a_failing_status_is_reported_with_its_code() {
let (base, _h) = serve(vec![
b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n".to_vec(),
]);
assert_eq!(
fetch(&base, 10_000_000).await.err(),
Some(FetchError::Status(404))
);
}
#[tokio::test]
async fn an_oversized_body_is_refused_by_its_content_length() {
let (base, _h) = serve(vec![ok_response("image/png", &png_bytes(100, 100), true)]);
assert_eq!(fetch(&base, 64).await.err(), Some(FetchError::TooBig));
}
#[tokio::test]
async fn an_oversized_body_without_a_content_length_is_refused_on_the_stream() {
let (base, _h) = serve(vec![ok_response("image/png", &png_bytes(100, 100), false)]);
assert_eq!(fetch(&base, 64).await.err(), Some(FetchError::TooBig));
}
#[tokio::test]
async fn an_empty_body_is_refused_rather_than_handed_on_as_zero_pixels() {
let (base, _h) = serve(vec![ok_response("image/png", &[], true)]);
assert_eq!(
fetch(&base, 10_000_000).await.err(),
Some(FetchError::Empty)
);
}
#[tokio::test]
async fn html_is_downloaded_and_its_type_is_carried_for_the_message() {
let (base, _h) = serve(vec![ok_response("text/html", b"<html>nope</html>", true)]);
let fetched = fetch(&base, 10_000_000).await.unwrap();
assert_eq!(fetched.content_type, "text/html");
}
#[test]
fn a_url_is_told_apart_from_a_path() {
assert!(looks_like_url("https://example.com/a.png"));
assert!(looks_like_url("HTTP://example.com/a.png"));
assert!(looks_like_url(" https://example.com/a.png "));
assert!(!looks_like_url("D:\\pics\\a.png"));
assert!(!looks_like_url("/home/user/a.png"));
assert!(!looks_like_url("a.png"));
assert!(!looks_like_url("./https://weird.png"));
}
#[test]
fn the_name_comes_from_the_last_path_segment() {
let name = |u: &str| display_name(&Url::parse(u).unwrap(), "png");
assert_eq!(name("https://example.com/pics/chart.png"), "chart.png");
assert_eq!(name("https://example.com/a/b.jpg?w=800&h=600"), "b.jpg");
assert_eq!(name("https://example.com/my%20chart.png"), "my chart.png");
}
#[test]
fn a_url_with_no_usable_segment_falls_back_to_the_host() {
let name = |u: &str| display_name(&Url::parse(u).unwrap(), "jpeg");
assert_eq!(name("https://example.com/"), "example.com.jpeg");
assert_eq!(name("https://example.com"), "example.com.jpeg");
assert_eq!(
name("https://cdn.example.com/?id=42"),
"cdn.example.com.jpeg"
);
}
#[test]
fn a_very_long_name_is_trimmed_but_keeps_its_extension() {
let long = format!("https://e.com/{}.png", "a".repeat(200));
let name = display_name(&Url::parse(&long).unwrap(), "png");
assert_eq!(name.chars().count(), NAME_CHARS);
assert!(name.ends_with(".png"), "{name}");
let long = format!("https://e.com/{}", "b".repeat(200));
assert_eq!(
display_name(&Url::parse(&long).unwrap(), "png")
.chars()
.count(),
NAME_CHARS
);
let long = format!("https://e.com/{}.{}", "c".repeat(100), "d".repeat(40));
let name = display_name(&Url::parse(&long).unwrap(), "png");
assert_eq!(name.chars().count(), NAME_CHARS);
assert!(!name.contains('.'), "{name}");
}
}