RustProxy/src/acl/dynamic.rs

//! Bounded, fail-closed DNS/DDNS ACL address cache. No resolver runs in the request path.
use crate::config::Config;
use anyhow::{ensure, Result};
use arc_swap::ArcSwap;
use futures_util::{stream, StreamExt};
use std::{
    collections::{BTreeSet, HashMap},
    net::IpAddr,
    sync::{Arc, LazyLock, Weak},
    time::{Duration, Instant},
};
use tokio::{sync::Notify, time::MissedTickBehavior};
use tokio_util::sync::CancellationToken;

pub const REFRESH_SECONDS: u64 = 60;
const MAX_AGE_SECONDS: u64 = 180;
const MAX_NAMES: usize = 128;
const MAX_ADDRESSES: usize = 32;
const PARALLEL_LOOKUPS: usize = 8;
const LOOKUP_SECONDS: u64 = 4;

#[derive(Clone)]
struct Addresses {
    ips: Vec<IpAddr>,
    valid_until: Instant,
}

static CACHE: LazyLock<ArcSwap<HashMap<String, Addresses>>> =
    LazyLock::new(|| ArcSwap::from_pointee(HashMap::new()));
static CHANGED: LazyLock<Notify> = LazyLock::new(Notify::new);

pub(crate) fn matches(name: &str, ip: IpAddr) -> bool {
    let cache = CACHE.load();
    cache.get(name).is_some_and(|entry| {
        entry.valid_until > Instant::now() && entry.ips.binary_search(&ip).is_ok()
    })
}

/// Strict ASCII DNS hostnames only (IDNs can be supplied in punycode).
/// A and AAAA results are eligible. IPs and CIDRs remain separate static rules.
pub(crate) fn hostname(value: &str) -> Result<String> {
    let candidate = value.strip_suffix('.').unwrap_or(value);
    ensure!(candidate.len() <= 253 && candidate.contains('.'), "ACL DNS name requires a fully qualified hostname");
    ensure!(candidate.is_ascii() && !candidate.contains('/') && !candidate.contains(':'), "invalid ACL DNS name");
    let labels: Vec<_> = candidate.split('.').collect();
    ensure!(labels.last().is_some_and(|part| part.chars().any(|ch| ch.is_ascii_alphabetic())), "invalid ACL DNS name suffix");
    for label in labels {
        ensure!(!label.is_empty() && label.len() <= 63, "invalid ACL DNS label length");
        ensure!(label.as_bytes()[0].is_ascii_alphanumeric()
            && label.as_bytes()[label.len() - 1].is_ascii_alphanumeric()
            && label.bytes().all(|byte| byte.is_ascii_alphanumeric() || byte == b'-'), "invalid ACL DNS label");
    }
    Ok(candidate.to_ascii_lowercase())
}

/// Walk every ACL scope, including standalone host policies, ACL profiles,
/// inline vhost/location ACLs and Docker overrides. Limit process-wide DNS work.
pub(crate) fn configured_hostnames(config: &Config) -> Result<Vec<String>> {
    let mut names = BTreeSet::new();
    let mut inspect = |acl: Option<&crate::acl::Acl>| -> Result<()> {
        if let Some(acl) = acl {
            for rule in &acl.rules {
                if rule.network.parse::<ipnet::IpNet>().is_ok() || rule.network.parse::<IpAddr>().is_ok() {
                    continue;
                }
                // A slash is always a CIDR error, never a DNS name.
                let name = hostname(&rule.network)?;
                ensure!(rule.action == crate::acl::Action::Allow,
                    "dynamic DNS ACL rules only support allow; use static IP/CIDR for deny");
                names.insert(name);
                ensure!(names.len() <= MAX_NAMES, "too many distinct dynamic DNS ACL names (max 128)");
            }
        }
        Ok(())
    };
    for policy in &config.host_policies { inspect(Some(&policy.acl))?; }
    for profile in &config.acl_profiles { inspect(profile.acl.as_ref())?; }
    for vhost in &config.vhosts {
        inspect(vhost.acl.as_ref())?;
        for route in &vhost.routes { inspect(route.acl.as_ref())?; }
    }
    for agent in &config.docker_agents.agents {
        for assignment in &agent.assignments { inspect(assignment.acl.as_ref())?; }
    }
    Ok(names.into_iter().collect())
}

async fn resolve(name: String) -> (String, Option<Addresses>) {
    // Keep the returned cache key owned independently of the DNS lookup future.
    // The lookup result can retain the borrowed hostname until it is dropped.
    let lookup_name = name.clone();
    let answer = tokio::time::timeout(
        Duration::from_secs(LOOKUP_SECONDS),
        tokio::net::lookup_host((lookup_name.as_str(), 0u16)),
    ).await;
    let mut ips = match answer {
        Ok(Ok(sockets)) => sockets.map(|socket| socket.ip()).collect::<Vec<_>>(),
        _ => {
            tracing::warn!(hostname = %name, "dynamic ACL DNS lookup failed; no addresses permitted");
            return (name, None);
        }
    };
    ips.sort_unstable();
    ips.dedup();
    if ips.is_empty() || ips.len() > MAX_ADDRESSES || ips.iter().any(|ip| {
        matches!(ip, IpAddr::V6(v6) if v6.to_ipv4_mapped().is_some())
    }) {
        tracing::warn!(hostname = %name, "dynamic ACL DNS returned an invalid or excessive address set");
        return (name, None);
    }
    (name, Some(Addresses { ips, valid_until: Instant::now() + Duration::from_secs(MAX_AGE_SECONDS) }))
}

pub(crate) fn config_changed() { CHANGED.notify_one(); }

pub(crate) async fn poll_loop(
    app: Weak<crate::runtime::App>,
    shutdown: CancellationToken,
) {
    let mut interval = tokio::time::interval(Duration::from_secs(REFRESH_SECONDS));
    interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
    loop {
        tokio::select! {
            _ = shutdown.cancelled() => break,
            _ = interval.tick() => {},
            _ = CHANGED.notified() => {},
        }
        let Some(state) = app.upgrade() else { break; };
        let snapshot = state.config.load();
        let Ok(names) = configured_hostnames(&snapshot.envelope.config) else { continue; };
        drop(snapshot);
        // Resolve concurrently with hard bounds; write one immutable snapshot.
        let results: HashMap<_, _> = tokio::select! {
            _ = shutdown.cancelled() => break,
            result = stream::iter(names.iter().cloned())
                .map(resolve)
                .buffer_unordered(PARALLEL_LOOKUPS)
                .filter_map(|(name, answer)| async move { answer.map(|value| (name, value)) })
                .collect() => result,
        };
        // A concurrently changed configuration must not publish obsolete DNS data.
        let current = state.config.load();
        if configured_hostnames(&current.envelope.config).ok().as_ref() != Some(&names) {
            CHANGED.notify_one();
            continue;
        }
        CACHE.store(Arc::new(results));
    }
    // Do not retain permit decisions when the resolver lifecycle stops.
    CACHE.store(Arc::new(HashMap::new()));
}

#[cfg(test)]
include!("dynamic/tests.rs");