Server-Side Request Forgery (SSRF) in Rust (reqwest, Axum)

// fetch.rs: reqwest may only connect to public addresses
use axum::extract::State;
use axum::http::StatusCode;
use axum::Json;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use reqwest::redirect::Policy;
use reqwest::{Client, Url};
use serde::Deserialize;
use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
use std::sync::Arc;
use std::time::Duration;
use url::Host;

#[derive(Deserialize)]
pub struct FetchRequest {
    url: String,
}

// reject private, loopback, link-local, CGNAT and non-global IPv6
fn is_public(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            let o = v4.octets();
            !(o[0] == 0 || v4.is_private() || v4.is_loopback() || v4.is_link_local() || v4.is_multicast()
                || v4.is_broadcast() || v4.is_documentation() || (o[0] == 100 && (o[1] & 0xc0) == 64))
        }
        IpAddr::V6(v6) => match v6.to_ipv4_mapped() {
            Some(v4) => is_public(IpAddr::V4(v4)),
            None => (v6.segments()[0] & 0xe000) == 0x2000,
        },
    }
}

// reqwest connects only to the addresses returned here
struct PublicResolver;

impl Resolve for PublicResolver {
    fn resolve(&self, name: Name) -> Resolving {
        let host = name.as_str().to_owned();
        Box::pin(async move {
            let addrs: Vec<SocketAddr> =
                tokio::task::spawn_blocking(move || (host.as_str(), 0).to_socket_addrs().map(Iterator::collect))
                    .await??;
            if addrs.is_empty() || !addrs.iter().all(|a| is_public(a.ip())) {
                return Err("blocked destination".into());
            }
            let addrs: Addrs = Box::new(addrs.into_iter());
            Ok(addrs)
        })
    }
}

// no redirects, no proxy, timeouts; IP literals are checked in checked_url
pub fn guarded_client() -> reqwest::Result<Client> {
    Client::builder()
        .dns_resolver(Arc::new(PublicResolver))
        .redirect(Policy::none())
        .no_proxy()
        .connect_timeout(Duration::from_secs(3))
        .timeout(Duration::from_secs(10))
        .build()
}

fn checked_url(raw: &str) -> Option<Url> {
    let url = Url::parse(raw).ok()?;
    if !matches!(url.scheme(), "http" | "https") {
        return None;
    }
    match url.host()? {
        Host::Ipv4(ip) if !is_public(IpAddr::V4(ip)) => None,
        Host::Ipv6(ip) if !is_public(IpAddr::V6(ip)) => None,
        _ => Some(url),
    }
}

async fn first_bytes(mut resp: reqwest::Response, limit: usize) -> reqwest::Result<Vec<u8>> {
    let mut out = Vec::new();
    while let Some(chunk) = resp.chunk().await? {
        out.extend_from_slice(&chunk);
        if out.len() >= limit {
            out.truncate(limit);
            break;
        }
    }
    Ok(out)
}

// at most 5000 bytes are read from any response
pub async fn preview(State(http): State<Client>, Json(input): Json<FetchRequest>) -> Result<String, StatusCode> {
    let url = checked_url(&input.url).ok_or(StatusCode::BAD_REQUEST)?;
    let resp = http.get(url).send().await.map_err(|_| StatusCode::BAD_GATEWAY)?;
    let body = first_bytes(resp, 5000).await.map_err(|_| StatusCode::BAD_GATEWAY)?;
    Ok(String::from_utf8_lossy(&body).into_owned())
}

pub async fn test_webhook(
    State(http): State<Client>,
    Json(input): Json<FetchRequest>,
) -> Result<Json<serde_json::Value>, StatusCode> {
    let url = checked_url(&input.url).ok_or(StatusCode::BAD_REQUEST)?;
    let resp = http
        .post(url)
        .json(&serde_json::json!({ "ping": true }))
        .send()
        .await
        .map_err(|_| StatusCode::BAD_GATEWAY)?;
    Ok(Json(serde_json::json!({ "status": resp.status().as_u16() })))
}