Merge pull request #277 from Dstack-TEE/gw-compat

gw: Add 0.3.x compatibility custom domain dns prefix

Author: kvinwang

gateway/gateway.toml

@@ -60,6 +60,7 @@ buffer_size = 8192
 connect_top_n = 3
 localhost_enabled = false
 app_address_ns_prefix = "_dstack-app-address"
+app_address_ns_compat = true
 workers = 32
 external_port = 443
 

gateway/rpc/proto/gateway_rpc.proto

@@ -41,8 +41,10 @@ message GuestAgentConfig {
   uint32 external_port = 1;
   // The in CVM port of the guest agent.
   uint32 internal_port = 2;
-  // The domain of the guest agent.
+  // The base domain of the zt-https
   string domain = 3;
+  // The app address namespace prefix
+  string app_address_ns_prefix = 4;
 }
 
 // StatusResponse is the response for Status.
@@ -145,6 +147,8 @@ message InfoResponse {
   string base_domain = 1;
   // The external port of the ZT-HTTPS
   uint32 external_port = 2;
+  // The app address namespace prefix
+  string app_address_ns_prefix = 3;
 }
 
 service Gateway {

gateway/src/config.rs

@@ -80,6 +80,7 @@ pub struct ProxyConfig {
     pub localhost_enabled: bool,
     pub workers: usize,
     pub app_address_ns_prefix: String,
+    pub app_address_ns_compat: bool,
 }
 
 #[derive(Debug, Clone, Deserialize)]

gateway/src/main_service.rs

@@ -735,9 +735,10 @@ impl GatewayRpc for RpcHandler {
                 servers,
             }),
             agent: Some(GuestAgentConfig {
-                external_port: state.config.proxy.listen_port as u32,
+                external_port: state.config.proxy.external_port as u32,
                 internal_port: state.config.proxy.agent_port as u32,
                 domain: state.config.proxy.base_domain.clone(),
+                app_address_ns_prefix: state.config.proxy.app_address_ns_prefix.clone(),
             }),
         };
         self.state.notify_state_updated.notify_one();
@@ -786,6 +787,7 @@ impl GatewayRpc for RpcHandler {
         Ok(InfoResponse {
             base_domain: state.config.proxy.base_domain.clone(),
             external_port: state.config.proxy.external_port as u32,
+            app_address_ns_prefix: state.config.proxy.app_address_ns_prefix.clone(),
         })
     }
 }

gateway/src/proxy/tls_passthough.rs

@@ -29,20 +29,42 @@ impl AppAddress {
 }
 
 /// resolve app address by sni
-async fn resolve_app_address(prefix: &str, sni: &str) -> Result<AppAddress> {
+async fn resolve_app_address(prefix: &str, sni: &str, compat: bool) -> Result<AppAddress> {
     let txt_domain = format!("{prefix}.{sni}");
     let resolver = hickory_resolver::AsyncResolver::tokio_from_system_conf()
         .context("failed to create dns resolver")?;
-    let lookup = resolver
-        .txt_lookup(txt_domain)
-        .await
-        .context("failed to lookup app address")?;
-    let txt_record = lookup.iter().next().context("no txt record found")?;
-    let data = txt_record
-        .txt_data()
-        .first()
-        .context("no data in txt record")?;
-    AppAddress::parse(data).context("failed to parse app address")
+
+    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 Some(data) = txt_record.txt_data().first() else {
+                continue;
+            };
+            return AppAddress::parse(data).context("failed to parse app address");
+        }
+        anyhow::bail!("failed to resolve app address");
+    } else {
+        let lookup = resolver
+            .txt_lookup(txt_domain)
+            .await
+            .context("failed to lookup app address")?;
+        let txt_record = lookup.iter().next().context("no txt record found")?;
+        let data = txt_record
+            .txt_data()
+            .first()
+            .context("no data in txt record")?;
+        AppAddress::parse(data).context("failed to parse app address")
+    }
 }
 
 pub(crate) async fn proxy_with_sni(
@@ -51,7 +73,9 @@ pub(crate) async fn proxy_with_sni(
     buffer: Vec<u8>,
     sni: &str,
 ) -> Result<()> {
-    let addr = resolve_app_address(&state.config.proxy.app_address_ns_prefix, sni)
+    let ns_prefix = &state.config.proxy.app_address_ns_prefix;
+    let compat = state.config.proxy.app_address_ns_compat;
+    let addr = resolve_app_address(ns_prefix, sni, compat)
         .await
         .context("failed to resolve app address")?;
     debug!("target address is {}:{}", addr.app_id, addr.port);
@@ -123,6 +147,7 @@ mod tests {
         let app_addr = resolve_app_address(
             "_dstack-app-address",
             "3327603e03f5bd1f830812ca4a789277fc31f577.app.kvin.wang",
+            false,
         )
         .await
         .unwrap();