tls_passthough.rs

// SPDX-FileCopyrightText: © 2024-2025 Phala Network <dstack@phala.network>
//
// SPDX-License-Identifier: Apache-2.0

use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Debug;
use std::sync::atomic::Ordering;

use anyhow::{bail, Context, Result};
use proxy_protocol::ProxyHeader;
use tokio::{io::AsyncWriteExt, net::TcpStream, task::JoinSet, time::timeout};
use tracing::{debug, info, warn};

use crate::{
    main_service::Proxy,
    models::{Counting, EnteredCounter},
};

use super::{
    io_bridge::bridge,
    port_policy::{filter_allowed_addresses, should_send_pp},
    AddressGroup,
};

#[derive(Debug)]
struct AppAddress {
    app_id: String,
    port: u16,
}

impl AppAddress {
    fn parse(data: &[u8]) -> Result<Self> {
        // format: "3327603e03f5bd1f830812ca4a789277fc31f577:555"
        let data = String::from_utf8(data.to_vec()).context("invalid app address")?;
        let (app_id, port) = data.split_once(':').context("invalid app address")?;
        Ok(Self {
            app_id: app_id.to_string(),
            port: port.parse().context("invalid port")?,
        })
    }
}

/// find app by app id in current memory box
fn select_app_address(items: &[Box<[u8]>], known_apps: &BTreeMap<String, BTreeSet<String>>) -> Result<AppAddress> {
    let mut fallback = None;
    for data in items {
        if let Ok(addr) = AppAddress::parse(data) {
            if known_apps.contains_key(&addr.app_id) {
                return Ok(addr);
            }
            if fallback.is_none() {
                fallback = Some(addr);
            }
        }
    }
    fallback.context("no app address found in txt record")
}

/// resolve app address by sni
async fn resolve_app_address(prefix: &str, sni: &str, compat: bool, state: &Proxy) -> Result<AppAddress> {
    let txt_domain = format!("{prefix}.{sni}");
    let resolver = hickory_resolver::AsyncResolver::tokio_from_system_conf()
        .context("failed to create dns resolver")?;

    if compat && prefix != "_tapp-address" {
        let txt_domain_legacy = format!("_tapp-address.{sni}");
        let (lookup, lookup_legacy) = tokio::join!(
            resolver.txt_lookup(txt_domain),
            resolver.txt_lookup(txt_domain_legacy),
        );
        for lookup in [lookup, lookup_legacy] {
            let Ok(lookup) = lookup else {
                continue;
            };
            let Some(txt_record) = lookup.iter().next() else {
                continue;
            };
            let locked = state.lock();
            if let Ok(addr) = select_app_address(txt_record.txt_data(), &locked.state.apps) {
                return Ok(addr);
            }
        }
    } else if let Ok(lookup) = resolver.txt_lookup(txt_domain).await {
        if let Some(txt_record) = lookup.iter().next() {
            let locked = state.lock();
            if let Ok(addr) = select_app_address(txt_record.txt_data(), &locked.state.apps) {
                return Ok(addr);
            }
        }
    }

    // wildcard fallback: try {prefix}-wildcard.{parent_domain}
    if let Some((_, parent)) = sni.split_once('.') {
        let wildcard_domain = format!("{prefix}-wildcard.{parent}");
        let lookup = resolver
            .txt_lookup(&wildcard_domain)
            .await
            .with_context(|| {
                format!("failed to lookup wildcard app address for {sni} via {wildcard_domain}")
            })?;
        let txt_record = lookup
            .iter()
            .next()
            .with_context(|| format!("no txt record found for {sni} via {wildcard_domain}"))?;
        let locked = state.lock();
        return select_app_address(txt_record.txt_data(), &locked.state.apps)
            .with_context(|| {
                format!("failed to parse app address for {sni} via {wildcard_domain}")
            });
    }

    anyhow::bail!("failed to resolve app address for {sni}");
}

pub(crate) async fn proxy_with_sni(
    state: Proxy,
    inbound: TcpStream,
    pp_header: ProxyHeader,
    buffer: Vec<u8>,
    sni: &str,
) -> Result<()> {
    let ns_prefix = &state.config.proxy.app_address_ns_prefix;
    let compat = state.config.proxy.app_address_ns_compat;
    let dns_timeout = state.config.proxy.timeouts.dns_resolve;
    let addr = timeout(dns_timeout, resolve_app_address(ns_prefix, sni, compat, &state))
        .await
        .with_context(|| format!("DNS TXT resolve timeout for {sni}"))?
        .with_context(|| format!("failed to resolve app address for {sni}"))?;
    debug!("target address is {}:{}", addr.app_id, addr.port);
    proxy_to_app(state, inbound, pp_header, buffer, &addr.app_id, addr.port).await
}

/// Check if app has reached max connections limit
fn check_connection_limit(
    addresses: &AddressGroup,
    max_connections: u64,
    app_id: &str,
) -> Result<()> {
    if max_connections == 0 {
        return Ok(());
    }
    let total: u64 = addresses
        .iter()
        .map(|a| a.counter.load(Ordering::Relaxed))
        .sum();
    if total >= max_connections {
        warn!(
            app_id,
            total, max_connections, "app connection limit exceeded"
        );
        bail!("app connection limit exceeded: {total}/{max_connections}");
    }
    Ok(())
}

/// connect to multiple hosts simultaneously and return the first successful connection
/// along with the instance_id of the winning address.
pub(crate) async fn connect_multiple_hosts(
    addresses: AddressGroup,
    port: u16,
    max_connections: u64,
    app_id: &str,
) -> Result<(TcpStream, EnteredCounter, String)> {
    check_connection_limit(&addresses, max_connections, app_id)?;

    let mut join_set = JoinSet::new();
    for addr in addresses {
        let counter = addr.counter.enter();
        let ip = addr.ip;
        let instance_id = addr.instance_id;
        debug!("connecting to {ip}:{port}");
        let future = TcpStream::connect((ip, port));
        join_set.spawn(async move {
            (
                future.await.map_err(|e| (e, ip, port)),
                counter,
                instance_id,
            )
        });
    }
    // select the first successful connection
    let (connection, counter, instance_id) = loop {
        let (result, counter, instance_id) = join_set
            .join_next()
            .await
            .context("No connection success")?
            .context("Failed to join the connect task")?;
        match result {
            Ok(connection) => break (connection, counter, instance_id),
            Err((e, addr, port)) => {
                info!("failed to connect to app@{addr}:{port}: {e}");
            }
        }
    };
    debug!("connected to {:?}", connection.peer_addr());
    Ok((connection, counter, instance_id))
}

pub(crate) async fn proxy_to_app(
    state: Proxy,
    inbound: TcpStream,
    pp_header: ProxyHeader,
    buffer: Vec<u8>,
    app_id: &str,
    port: u16,
) -> Result<()> {
    let addresses = state.lock().select_top_n_hosts(app_id)?;
    let addresses = filter_allowed_addresses(&state, addresses, app_id, port)?;
    let max_connections = state.config.proxy.max_connections_per_app;
    let (mut outbound, _counter, instance_id) = timeout(
        state.config.proxy.timeouts.connect,
        connect_multiple_hosts(addresses.clone(), port, max_connections, app_id),
    )
    .await
    .with_context(|| format!("connecting timeout to app {app_id}: {addresses:?}:{port}"))?
    .with_context(|| format!("failed to connect to app {app_id}: {addresses:?}:{port}"))?;
    if should_send_pp(&state, &instance_id, port) {
        let pp_header_bin =
            proxy_protocol::encode(pp_header).context("failed to encode pp header")?;
        outbound.write_all(&pp_header_bin).await?;
    }
    outbound
        .write_all(&buffer)
        .await
        .context("failed to write to app")?;
    bridge(inbound, outbound, &state.config.proxy)
        .await
        .context("failed to copy between inbound and outbound")?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        config::{load_config_figment, Config, MutualConfig, TlsConfig},
        main_service::ProxyOptions,
    };
    use tempfile::TempDir;

    fn boxed(s: &[u8]) -> Box<[u8]> {
        s.to_vec().into_boxed_slice()
    }

    async fn create_test_proxy() -> (Proxy, TempDir) {
        let figment = load_config_figment(None);
        let mut config = figment.focus("core").extract::<Config>().unwrap();
        let temp_dir = TempDir::new().expect("failed to create temp dir");
        config.sync.data_dir = temp_dir.path().to_string_lossy().to_string();
        let proxy = Proxy::new(ProxyOptions {
            config,
            my_app_id: None,
            tls_config: TlsConfig {
                certs: "".to_string(),
                key: "".to_string(),
                mutual: MutualConfig { ca_certs: "".to_string() },
            },
        })
        .await
        .expect("failed to create proxy");
        (proxy, temp_dir)
    }

    #[tokio::test]
    async fn test_resolve_app_address() {
        let (state, _dir) = create_test_proxy().await;
        let app_addr = resolve_app_address(
            "_dstack-app-address",
            "3327603e03f5bd1f830812ca4a789277fc31f577.app.dstack.org",
            false,
            &state,
        )
        .await
        .unwrap();
        assert_eq!(app_addr.app_id, "3327603e03f5bd1f830812ca4a789277fc31f577");
        assert_eq!(app_addr.port, 8090);
    }

    #[test]
    fn test_select_app_address_prefers_local() {
        let items = vec![boxed(b"aaaaaa:443"), boxed(b"bbbbbb:8080")];
        let mut apps = BTreeMap::new();
        apps.insert("bbbbbb".to_string(), BTreeSet::new());
        let addr = select_app_address(&items, &apps).unwrap();
        assert_eq!(addr.app_id, "bbbbbb");
        assert_eq!(addr.port, 8080);
    }

    #[test]
    fn test_select_app_address_fallback_when_none_local() {
        let items = vec![boxed(b"aaaaaa:443"), boxed(b"bbbbbb:8080")];
        let addr = select_app_address(&items, &BTreeMap::new()).unwrap();
        assert_eq!(addr.app_id, "aaaaaa");
        assert_eq!(addr.port, 443);
    }
}