RustProxy/src/upstream.rs

use crate::config::{ip_policy, Backend};
use anyhow::{ensure, Context, Result};
use http::Uri;
use hyper_rustls::HttpsConnector;
use hyper_util::client::legacy::connect::{dns::Name, HttpConnector};
use ipnet::IpNet;
use rustls::{
    client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
    pki_types::{CertificateDer, ServerName, UnixTime},
    DigitallySignedStruct, RootCertStore, SignatureScheme,
};
use std::{
    future::Future,
    io,
    net::{IpAddr, SocketAddr, ToSocketAddrs},
    pin::Pin,
    sync::{Arc, OnceLock},
    task::{Context as TaskContext, Poll},
    time::Duration,
};
use tokio::sync::Semaphore;
use tower::Service;

const DNS_TIMEOUT: Duration = Duration::from_secs(5);
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_DNS_ANSWERS: usize = 32;
static DNS_SLOTS: OnceLock<Arc<Semaphore>> = OnceLock::new();
pub type BoxError = Box<dyn std::error::Error + Send + Sync>;

mod client;
pub use client::{
    client, client_with_keepalive, client_with_timeout, is_graceful_h2_goaway, ClientRotationStats,
    ProxyClient,
};

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct AddressPolicy {
    pub allow_private: bool,
    pub accept_local_listener_risk: bool,
    pub allowlist: Vec<IpNet>,
    port: u16,
    listeners: Vec<SocketAddr>,
    local_ips: Vec<IpAddr>,
    interfaces_known: bool,
}
impl AddressPolicy {
    pub fn new(backend: &Backend) -> Result<Self> {
        ensure!(
            backend.address_allowlist.len() <= 32,
            "at most 32 backend CIDR allowlist entries"
        );
        let allowlist = backend
            .address_allowlist
            .iter()
            .map(|network| {
                ensure!(network.len() <= 64, "CIDR is too long");
                network
                    .parse::<IpNet>()
                    .context("invalid backend CIDR allowlist entry")
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Self {
            allow_private: backend.allow_private,
            accept_local_listener_risk: backend.accept_local_listener_risk,
            allowlist,
            port: url::Url::parse(&backend.url)?
                .port_or_known_default()
                .unwrap_or(0),
            ..Default::default()
        })
    }
    pub fn with_listeners(mut self, policy: &crate::config::ValidationPolicy) -> Self {
        self.listeners = vec![policy.proxy_addr, policy.admin_addr];
        self.listeners.extend(policy.tls_addr);
        let addresses = local_addresses();
        self.interfaces_known = addresses.is_some();
        self.local_ips = addresses.unwrap_or_default();
        self.local_ips.sort_unstable();
        self.local_ips.dedup();
        self
    }
    pub fn with_bindings(mut self, bindings: &[SocketAddr]) -> Self {
        self.listeners.extend_from_slice(bindings);
        self.listeners.sort();
        self.listeners.dedup();
        self
    }
    pub fn check(&self, ip: IpAddr) -> Result<()> {
        let ip = match ip {
            IpAddr::V6(ip) => ip
                .to_ipv4_mapped()
                .map(IpAddr::V4)
                .unwrap_or(IpAddr::V6(ip)),
            ip => ip,
        };
        ip_policy(ip, self.allow_private)?;
        for listener in &self.listeners {
            if listener.port() != self.port {
                continue;
            }
            let local = match listener.ip() {
                IpAddr::V6(ip) => ip
                    .to_ipv4_mapped()
                    .map(IpAddr::V4)
                    .unwrap_or(IpAddr::V6(ip)),
                ip => ip,
            };
            let local_interface =
                self.local_ips.contains(&ip) || (!self.interfaces_known && locally_bindable(ip));
            let self_target = ip == local
                || (ip.is_loopback() && local.is_loopback())
                || (local.is_unspecified() && (ip.is_loopback() || local_interface));
            ensure!(
                !self_target || self.accept_local_listener_risk,
                "backend resolves to a local RustProxy listener ({ip}:{} via {listener}); set allow_private=true and accept_local_listener_risk=true only after accepting proxy-loop/control-plane exposure risk",
                self.port
            );
        }
        ensure!(
            self.allowlist.is_empty() || self.allowlist.iter().any(|network| network.contains(&ip)),
            "backend address is outside its CIDR allowlist"
        );
        Ok(())
    }
    fn validate_answers(
        &self,
        answers: impl Iterator<Item = SocketAddr>,
    ) -> io::Result<Vec<SocketAddr>> {
        let mut approved = Vec::new();
        for (index, address) in answers.enumerate() {
            if index >= MAX_DNS_ANSWERS {
                return Err(io::Error::other("too many DNS answers"));
            }
            self.check(address.ip())
                .map_err(|e| io::Error::new(io::ErrorKind::PermissionDenied, e.to_string()))?;
            if !approved.contains(&address) {
                approved.push(address);
            }
        }
        if approved.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::NotFound,
                "DNS returned no usable addresses",
            ));
        }
        Ok(approved)
    }
}
#[derive(Clone, Debug)]
pub struct GuardedResolver {
    policy: AddressPolicy,
    connect_host: Option<String>,
}
impl Service<Name> for GuardedResolver {
    type Response = std::vec::IntoIter<SocketAddr>;
    type Error = io::Error;
    type Future = Pin<Box<dyn Future<Output = io::Result<Self::Response>> + Send>>;
    fn poll_ready(&mut self, _: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
    fn call(&mut self, name: Name) -> Self::Future {
        let policy = self.policy.clone();
        let hostname = self
            .connect_host
            .clone()
            .unwrap_or_else(|| name.as_str().to_owned());
        let permit = DNS_SLOTS
            .get_or_init(|| Arc::new(Semaphore::new(16)))
            .clone()
            .try_acquire_owned();
        Box::pin(async move {
            let permit = permit.map_err(|_| {
                io::Error::new(io::ErrorKind::WouldBlock, "DNS worker limit reached")
            })?;
            let lookup = tokio::task::spawn_blocking(move || {
                let _permit = permit;
                policy.validate_answers((hostname.as_str(), 0).to_socket_addrs()?)
            });
            let result = tokio::time::timeout(DNS_TIMEOUT, lookup)
                .await
                .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "DNS lookup timed out"))?
                .map_err(|e| io::Error::other(e.to_string()))??;
            Ok(result.into_iter())
        })
    }
}
type InnerConnector = HttpsConnector<HttpConnector<GuardedResolver>>;
#[derive(Clone, Debug)]
pub struct TimedConnector {
    inner: InnerConnector,
    policy: AddressPolicy,
    timeout: Duration,
    destination: Uri,
}
impl Service<Uri> for TimedConnector {
    type Response = <InnerConnector as Service<Uri>>::Response;
    type Error = BoxError;
    type Future =
        Pin<Box<dyn Future<Output = std::result::Result<Self::Response, BoxError>> + Send>>;
    fn poll_ready(&mut self, cx: &mut TaskContext<'_>) -> Poll<std::result::Result<(), BoxError>> {
        self.inner.poll_ready(cx).map_err(Into::into)
    }
    fn call(&mut self, _request_uri: Uri) -> Self::Future {
        let uri = self.destination.clone();
        if let Some(ip) = uri.host().and_then(|host| {
            host.trim_start_matches('[')
                .trim_end_matches(']')
                .parse::<IpAddr>()
                .ok()
        }) {
            if let Err(error) = self.policy.check(ip) {
                return Box::pin(async move {
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::PermissionDenied,
                        error.to_string(),
                    )) as BoxError)
                });
            }
        }
        let future = self.inner.call(uri);
        let timeout = self.timeout;
        Box::pin(async move {
            match tokio::time::timeout(timeout, future).await {
                Ok(result) => result,
                Err(error) => Err(Box::new(error) as BoxError),
            }
        })
    }
}
fn public_roots() -> RootCertStore {
    let mut roots = RootCertStore::empty();
    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
    roots
}
fn roots(backend: &Backend) -> Result<RootCertStore> {
    let mut roots = public_roots();
    if let Some(pem) = &backend.tls_ca_pem {
        ensure!(pem.len() <= 65_536, "additional CA PEM exceeds 64 KiB");
        let certs = rustls_pemfile::certs(&mut pem.as_bytes()).collect::<io::Result<Vec<_>>>()?;
        ensure!(
            !certs.is_empty() && certs.len() <= 16,
            "additional CA PEM must contain 1..16 certificates"
        );
        for cert in certs {
            roots.add(cert).context("invalid additional trust anchor")?;
        }
    }
    Ok(roots)
}

#[derive(Debug)]
struct BackendCertificateVerificationDisabled {
    supported_schemes: Vec<SignatureScheme>,
}
impl ServerCertVerifier for BackendCertificateVerificationDisabled {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp_response: &[u8],
        _now: UnixTime,
    ) -> std::result::Result<ServerCertVerified, rustls::Error> {
        Ok(ServerCertVerified::assertion())
    }
    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }
    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }
    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        self.supported_schemes.clone()
    }
}
pub async fn tcp_probe(origin: &str, policy: &AddressPolicy) -> Result<()> {
    let url = url::Url::parse(origin)?;
    let host = url
        .host_str()
        .context("probe host missing")?
        .trim_start_matches('[')
        .trim_end_matches(']');
    let port = url.port_or_known_default().context("probe port missing")?;
    let addresses = if let Ok(ip) = host.parse::<IpAddr>() {
        policy.check(ip)?;
        vec![SocketAddr::new(ip, port)]
    } else {
        let mut resolver = GuardedResolver {
            policy: policy.clone(),
            connect_host: None,
        };
        resolver
            .call(host.parse::<Name>()?)
            .await?
            .map(|mut a| {
                a.set_port(port);
                a
            })
            .collect()
    };
    for address in addresses {
        if tokio::net::TcpStream::connect(address).await.is_ok() {
            return Ok(());
        }
    }
    anyhow::bail!("no approved backend address accepted TCP")
}
#[cfg(test)]
include!("upstream/tests.rs");

#[cfg(unix)]
fn local_addresses() -> Option<Vec<IpAddr>> {
    let mut first: *mut libc::ifaddrs = std::ptr::null_mut();
    unsafe {
        if libc::getifaddrs(&mut first) != 0 {
            return None;
        }
        let mut result = Vec::new();
        let mut current = first;
        while !current.is_null() {
            let item = &*current;
            if !item.ifa_addr.is_null() {
                match (*item.ifa_addr).sa_family as i32 {
                    libc::AF_INET => {
                        let addr = &*(item.ifa_addr as *const libc::sockaddr_in);
                        result.push(IpAddr::V4(std::net::Ipv4Addr::from(
                            addr.sin_addr.s_addr.to_ne_bytes(),
                        )));
                    }
                    libc::AF_INET6 => {
                        let addr = &*(item.ifa_addr as *const libc::sockaddr_in6);
                        result.push(IpAddr::V6(std::net::Ipv6Addr::from(addr.sin6_addr.s6_addr)));
                    }
                    _ => {}
                }
            }
            current = item.ifa_next;
        }
        libc::freeifaddrs(first);
        Some(result)
    }
}
#[cfg(not(unix))]
fn local_addresses() -> Option<Vec<IpAddr>> {
    None
}

fn locally_bindable(ip: IpAddr) -> bool {
    std::net::UdpSocket::bind(SocketAddr::new(ip, 0)).is_ok()
}