use std::net::{IpAddr, ToSocketAddrs};
pub fn url_host(url: &str) -> Option<String> {
let after = url.split_once("://").map(|(_, r)| r).unwrap_or(url);
let host = after.split(['/', '?', '#']).next().unwrap_or("");
let host = host.rsplit('@').next().unwrap_or(host);
if let Some(end) = host.strip_prefix('[').and_then(|h| h.split_once(']')).map(|(h, _)| h) {
return Some(end.to_lowercase());
}
let host = host.split(':').next().unwrap_or(host);
if host.is_empty() { None } else { Some(host.to_lowercase()) }
}
fn ip_is_blocked(ip: &IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => {
v4.is_loopback() || v4.is_private() || v4.is_link_local()
|| v4.is_unspecified() || v4.is_multicast() || v4.is_broadcast()
}
IpAddr::V6(v6) => {
if let Some(v4) = v6.to_ipv4_mapped() {
return ip_is_blocked(&IpAddr::V4(v4));
}
v6.is_loopback() || v6.is_unspecified() || v6.is_multicast()
|| (v6.segments()[0] & 0xfe00) == 0xfc00
|| (v6.segments()[0] & 0xffc0) == 0xfe80
}
}
}
pub fn url_is_allowed(url: &str, allow_hosts: &[String]) -> Result<(), String> {
let host = url_host(url).ok_or_else(|| format!("no host in url: {url}"))?;
if allow_hosts.iter().any(|h| h.to_lowercase() == host) {
return Ok(());
}
if let Ok(ip) = host.parse::<IpAddr>() {
return if ip_is_blocked(&ip) { Err(format!("blocked ip: {host}")) } else { Ok(()) };
}
let addrs = (host.as_str(), 80u16)
.to_socket_addrs()
.map_err(|e| format!("cannot resolve {host}: {e}"))?;
let mut any = false;
for a in addrs {
any = true;
if ip_is_blocked(&a.ip()) {
return Err(format!("blocked resolved ip {} for {host}", a.ip()));
}
}
if !any {
return Err(format!("no addresses for {host}"));
}
Ok(())
}
pub fn resolve_proxy(
config_proxy: Option<&str>,
lookup: impl Fn(&str) -> Option<String>,
) -> Option<String> {
if let Some(p) = config_proxy {
if !p.is_empty() {
return Some(p.to_string());
}
}
for key in ["ALL_PROXY", "HTTPS_PROXY", "HTTP_PROXY"] {
if let Some(v) = lookup(key) {
if !v.is_empty() {
return Some(v);
}
}
}
None
}
pub const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
pub fn build_client(proxy: Option<&str>) -> reqwest::Result<reqwest::Client> {
let mut builder = reqwest::Client::builder().timeout(REQUEST_TIMEOUT);
if let Some(p) = proxy {
builder = builder.proxy(reqwest::Proxy::all(p)?);
}
builder.build()
}
pub fn build_client_ua(proxy: Option<&str>, user_agent: &str) -> reqwest::Result<reqwest::Client> {
let mut builder = reqwest::Client::builder().user_agent(user_agent).timeout(REQUEST_TIMEOUT);
if let Some(p) = proxy {
builder = builder.proxy(reqwest::Proxy::all(p)?);
}
builder.build()
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
fn env<'a>(map: &'a HashMap<&'a str, &'a str>) -> impl Fn(&str) -> Option<String> + 'a {
move |k| map.get(k).map(|v| v.to_string())
}
#[test]
fn config_proxy_wins_over_env() {
let mut m = HashMap::new();
m.insert("ALL_PROXY", "socks5://env:1");
let got = resolve_proxy(Some("http://cfg:2"), env(&m));
assert_eq!(got.as_deref(), Some("http://cfg:2"));
}
#[test]
fn falls_back_to_env_in_priority_order() {
let mut m = HashMap::new();
m.insert("HTTP_PROXY", "http://low:1");
m.insert("HTTPS_PROXY", "http://mid:1");
assert_eq!(resolve_proxy(None, env(&m)).as_deref(), Some("http://mid:1"));
m.insert("ALL_PROXY", "socks5://top:1");
assert_eq!(resolve_proxy(None, env(&m)).as_deref(), Some("socks5://top:1"));
}
#[test]
fn none_when_unset() {
let m: HashMap<&str, &str> = HashMap::new();
assert_eq!(resolve_proxy(None, env(&m)), None);
}
#[test]
fn builds_client_with_and_without_proxy() {
assert!(build_client(None).is_ok());
assert!(build_client(Some("socks5://127.0.0.1:9050")).is_ok());
}
#[test]
fn builds_client_ua_with_and_without_proxy() {
assert!(build_client_ua(None, "kibble-crawler/1.0").is_ok());
assert!(build_client_ua(Some("socks5://127.0.0.1:9050"), "kibble-crawler/1.0").is_ok());
}
#[test]
fn rejects_loopback_and_private_literals() {
assert!(url_is_allowed("http://127.0.0.1/x", &[]).is_err());
assert!(url_is_allowed("http://10.0.0.5/x", &[]).is_err());
assert!(url_is_allowed("http://169.254.169.254/latest/meta-data/", &[]).is_err());
assert!(url_is_allowed("http://[::1]/x", &[]).is_err());
}
#[test]
fn allow_hosts_bypasses_guard() {
assert!(url_is_allowed("http://127.0.0.1:8080/x", &["127.0.0.1".to_string()]).is_ok());
}
#[test]
fn rejects_unresolvable_host() {
assert!(url_is_allowed("http://no.such.host.invalid./x", &[]).is_err());
}
#[test]
fn rejects_ipv4_mapped_ipv6_loopback() {
assert!(url_is_allowed("http://[::ffff:127.0.0.1]/x", &[]).is_err());
assert!(url_is_allowed("http://[::ffff:10.0.0.1]/x", &[]).is_err());
}
#[test]
fn userinfo_at_host_uses_connect_target() {
assert!(url_is_allowed("http://evil.com@127.0.0.1/x", &[]).is_err());
}
}