gw: per-instance PROXY protocol via app-compose port_attrs

Replace the global outbound_pp_enabled switch with a per-(instance, port)
lookup so different ports of the same backend can have different PP
behaviour. PP is declared by the app and reported to the gateway through
authenticated channels — never by client SNI.

Pipeline:

1. dstack-types::AppCompose grows a "ports" array. Each entry carries a
   port number and a "pp" flag. Because it's part of app-compose.json it
   is measured into compose_hash and attested.

2. RegisterCvmRequest grows an optional PortAttrsList. New CVMs include
   their port_attrs at WireGuard registration time. The optional wrapper
   lets the gateway distinguish "not reported" (legacy CVM) from
   "reported empty" (new CVM with no PP-enabled port).

3. The gateway stores port_attrs on InstanceInfo and persists/syncs it
   via WaveKV (InstanceData), keyed by instance_id (different instances
   of the same app may run different code).

4. AddressInfo now carries instance_id, and connect_multiple_hosts
   returns the winner's instance_id. The proxy looks up that instance's
   port_attrs to decide whether to send a PROXY header.

5. Backward compat: if an instance has no port_attrs (legacy CVM), the
   gateway lazily fetches them via the agent's Info() RPC, parses
   tcb_info.app_compose, and caches the result in WaveKV.

PROXY protocol module is unchanged; only the *decision* of whether to
send a header moves from a global config to a per-instance lookup.

Author: kvinwang

dstack-types/src/lib.rs

@@ -45,6 +45,18 @@ pub struct AppCompose {
     pub storage_fs: Option<String>,
     #[serde(default, with = "human_size")]
     pub swap_size: u64,
+    /// Per-port attributes consumed by the gateway (e.g. PROXY protocol).
+    #[serde(default)]
+    pub ports: Vec<PortAttrs>,
+}
+
+#[derive(Deserialize, Serialize, Debug, Clone)]
+pub struct PortAttrs {
+    pub port: u16,
+    /// Whether the gateway should send a PROXY protocol header on outbound
+    /// connections to this port.
+    #[serde(default)]
+    pub pp: bool,
 }
 
 fn default_true() -> bool {

dstack-util/src/system_setup.rs

@@ -53,7 +53,8 @@ use crate::{
 use cert_client::CertRequestClient;
 use cmd_lib::run_fun as cmd;
 use dstack_gateway_rpc::{
-    gateway_client::GatewayClient, RegisterCvmRequest, RegisterCvmResponse, WireGuardPeer,
+    gateway_client::GatewayClient, PortAttrs as RpcPortAttrs, PortAttrsList, RegisterCvmRequest,
+    RegisterCvmResponse, WireGuardPeer,
 };
 use ra_tls::rcgen::{KeyPair, PKCS_ECDSA_P256_SHA256};
 use serde_human_bytes as hex_bytes;
@@ -446,11 +447,24 @@ impl<'a> GatewayContext<'a> {
         gateway_url: &str,
         key_store: &GatewayKeyStore,
     ) -> Result<RegisterCvmResponse> {
+        let port_attrs = PortAttrsList {
+            attrs: self
+                .shared
+                .app_compose
+                .ports
+                .iter()
+                .map(|p| RpcPortAttrs {
+                    port: p.port as u32,
+                    pp: p.pp,
+                })
+                .collect(),
+        };
         let client =
             self.create_gateway_client(gateway_url, &key_store.client_key, &key_store.client_cert)?;
         let result = client
             .register_cvm(RegisterCvmRequest {
                 client_public_key: key_store.wg_pk.clone(),
+                port_attrs: Some(port_attrs.clone()),
             })
             .await
             .context("Failed to register CVM");
@@ -471,6 +485,7 @@ impl<'a> GatewayContext<'a> {
         client
             .register_cvm(RegisterCvmRequest {
                 client_public_key: key_store.wg_pk.clone(),
+                port_attrs: Some(port_attrs),
             })
             .await
             .context("Failed to register CVM")

gateway/gateway.toml

@@ -60,8 +60,6 @@ external_port = 443
 max_connections_per_app = 2000
 # Whether to read PROXY protocol from inbound connections (e.g. from Cloudflare).
 inbound_pp_enabled = false
-# Whether to send PROXY protocol headers to backend apps.
-outbound_pp_enabled = false
 
 [core.proxy.timeouts]
 # Timeout for establishing a connection to the target app.

gateway/rpc/proto/gateway_rpc.proto

@@ -12,6 +12,25 @@ package gateway;
 message RegisterCvmRequest {
   // The public key of the WireGuard interface of the CVM.
   string client_public_key = 1;
+  // Per-port attributes the gateway should apply when proxying to this CVM.
+  // Wrapped in a message so we can distinguish "not reported" (old CVM →
+  // gateway falls back to fetching app-compose via Info()) from "reported
+  // empty" (new CVM with no special port behaviour).
+  optional PortAttrsList port_attrs = 2;
+}
+
+// PortAttrsList wraps a list of PortAttrs so it can be optional on the wire.
+message PortAttrsList {
+  repeated PortAttrs attrs = 1;
+}
+
+// PortAttrs declares per-port behaviour for the gateway.
+message PortAttrs {
+  // The CVM port these attributes apply to.
+  uint32 port = 1;
+  // Whether the gateway should send a PROXY protocol header on outbound
+  // connections to this port.
+  bool pp = 2;
 }
 
 // DebugRegisterCvmRequest is the request for DebugRegisterCvm (only works when debug_mode is enabled).

gateway/src/config.rs

@@ -117,14 +117,13 @@ pub struct ProxyConfig {
     pub app_address_ns_compat: bool,
     /// Maximum concurrent connections per app. 0 means unlimited.
     pub max_connections_per_app: u64,
+    /// Port the dstack guest-agent listens on inside each CVM. Used by the
+    /// gateway to fetch app metadata (e.g. port_attrs for legacy CVMs).
+    pub agent_port: u16,
     /// Whether to read PROXY protocol headers from inbound connections
     /// (e.g. when behind a PP-aware load balancer like Cloudflare).
     #[serde(default)]
     pub inbound_pp_enabled: bool,
-    /// Whether to send PROXY protocol headers on outbound connections to backend apps.
-    /// This is a server-side setting; it must NOT be controlled by client input (e.g. SNI).
-    #[serde(default)]
-    pub outbound_pp_enabled: bool,
 }
 
 #[derive(Debug, Clone, Deserialize)]

gateway/src/debug_service.rs

@@ -35,6 +35,7 @@ impl DebugRpc for DebugRpcHandler {
             &request.app_id,
             &request.instance_id,
             &request.client_public_key,
+            None,
         )
     }
 

gateway/src/kv/mod.rs

@@ -42,13 +42,26 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::watch;
 use wavekv::{node::NodeState, types::NodeId, Node};
 
+/// Per-port flags applied by the gateway when proxying to a CVM port.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
+pub struct PortFlags {
+    /// Send a PROXY protocol header on outbound connections to this port.
+    #[serde(default)]
+    pub pp: bool,
+}
+
 /// Instance core data (persistent)
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct InstanceData {
     pub app_id: String,
     pub ip: Ipv4Addr,
     pub public_key: String,
     pub reg_time: u64,
+    /// Per-port flags reported at registration. `None` means "not reported"
+    /// (legacy CVM); the gateway will fall back to fetching app-compose via
+    /// Info() on first connection and populate this lazily.
+    #[serde(default)]
+    pub port_attrs: Option<BTreeMap<u16, PortFlags>>,
 }
 
 /// Gateway node status (stored separately for independent updates)

gateway/src/main_service.rs

@@ -37,7 +37,7 @@ use crate::{
     config::{Config, TlsConfig},
     kv::{
         fetch_peers_from_bootnode, AppIdValidator, HttpsClientConfig, InstanceData, KvStore,
-        NodeData, NodeStatus, WaveKvSyncService,
+        NodeData, NodeStatus, PortFlags, WaveKvSyncService,
     },
     models::{InstanceInfo, WgConf},
     proxy::{create_acceptor_with_cert_resolver, AddressGroup, AddressInfo},
@@ -378,12 +378,16 @@ impl Proxy {
         })
     }
 
-    /// Register a CVM with the given app_id, instance_id and client_public_key
+    /// Register a CVM with the given app_id, instance_id and client_public_key.
+    ///
+    /// `port_attrs = None` means the CVM didn't report port attributes (legacy
+    /// CVM). The gateway will lazily fetch them via Info() on first connection.
     pub fn do_register_cvm(
         &self,
         app_id: &str,
         instance_id: &str,
         client_public_key: &str,
+        port_attrs: Option<BTreeMap<u16, PortFlags>>,
     ) -> Result<RegisterCvmResponse> {
         let mut state = self.lock();
 
@@ -403,7 +407,7 @@ impl Proxy {
             bail!("[{instance_id}] client public key is empty");
         }
         let client_info = state
-            .new_client_by_id(instance_id, app_id, client_public_key)
+            .new_client_by_id(instance_id, app_id, client_public_key, port_attrs)
             .context("failed to allocate IP address for client")?;
         if let Err(err) = state.reconfigure() {
             error!("failed to reconfigure: {err:?}");
@@ -425,7 +429,7 @@ impl Proxy {
             }),
             agent: Some(GuestAgentConfig {
                 external_port: port.into(),
-                internal_port: 8090,
+                internal_port: state.config.proxy.agent_port.into(),
                 domain: base_domain,
                 app_address_ns_prefix: state.config.proxy.app_address_ns_prefix.clone(),
             }),
@@ -449,6 +453,7 @@ fn build_state_from_kv_store(instances: BTreeMap<String, InstanceData>) -> Proxy
             reg_time: UNIX_EPOCH
                 .checked_add(Duration::from_secs(data.reg_time))
                 .unwrap_or(UNIX_EPOCH),
+            port_attrs: data.port_attrs,
             connections: Default::default(),
         };
         state.allocated_addresses.insert(data.ip);
@@ -742,6 +747,7 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()>
             reg_time: UNIX_EPOCH
                 .checked_add(Duration::from_secs(data.reg_time))
                 .unwrap_or(UNIX_EPOCH),
+            port_attrs: data.port_attrs.clone(),
             connections: Default::default(),
         };
 
@@ -823,6 +829,7 @@ impl ProxyState {
         id: &str,
         app_id: &str,
         public_key: &str,
+        port_attrs: Option<BTreeMap<u16, PortFlags>>,
     ) -> Result<InstanceInfo> {
         if id.is_empty() {
             bail!("instance_id is empty (no_instance_id is set?)");
@@ -841,6 +848,10 @@ impl ProxyState {
                 // Update reg_time so other nodes will pick up the change
                 existing.reg_time = SystemTime::now();
             }
+            // Always update port_attrs from the latest registration so changes
+            // take effect across re-registrations. A `None` here (legacy CVM)
+            // wipes any previously-cached attrs so the lazy fetch path runs again.
+            existing.port_attrs = port_attrs.clone();
             let existing = existing.clone();
             if self.valid_ip(existing.ip) {
                 // Sync existing instance to KvStore (might be from legacy state)
@@ -849,6 +860,7 @@ impl ProxyState {
                     ip: existing.ip,
                     public_key: existing.public_key.clone(),
                     reg_time: encode_ts(existing.reg_time),
+                    port_attrs: existing.port_attrs.clone(),
                 };
                 if let Err(err) = self.kv_store.sync_instance(&existing.id, &data) {
                     error!("failed to sync existing instance to KvStore: {err:?}");
@@ -867,19 +879,58 @@ impl ProxyState {
             ip,
             public_key: public_key.to_string(),
             reg_time: SystemTime::now(),
+            port_attrs,
             connections: Default::default(),
         };
         self.add_instance(host_info.clone());
         Ok(host_info)
     }
 
+    /// Lookup an instance's IP. Returns `None` if the instance is unknown.
+    pub(crate) fn instance_ip(&self, instance_id: &str) -> Option<Ipv4Addr> {
+        self.state.instances.get(instance_id).map(|i| i.ip)
+    }
+
+    /// Lookup an instance's port_attrs. `None` means the CVM never reported
+    /// them (legacy CVM), so the caller should fall back to fetching via Info().
+    pub(crate) fn instance_port_attrs(
+        &self,
+        instance_id: &str,
+    ) -> Option<BTreeMap<u16, PortFlags>> {
+        self.state.instances.get(instance_id)?.port_attrs.clone()
+    }
+
+    /// Update an instance's port_attrs (used after a lazy fetch via Info()).
+    /// Persists to the WaveKV store so other gateway nodes pick it up.
+    pub(crate) fn update_instance_port_attrs(
+        &mut self,
+        instance_id: &str,
+        attrs: BTreeMap<u16, PortFlags>,
+    ) {
+        let Some(info) = self.state.instances.get_mut(instance_id) else {
+            return;
+        };
+        info.port_attrs = Some(attrs.clone());
+        let data = InstanceData {
+            app_id: info.app_id.clone(),
+            ip: info.ip,
+            public_key: info.public_key.clone(),
+            reg_time: encode_ts(info.reg_time),
+            port_attrs: Some(attrs),
+        };
+        if let Err(err) = self.kv_store.sync_instance(instance_id, &data) {
+            error!("failed to sync updated port_attrs to KvStore: {err:?}");
+        }
+    }
+
     fn add_instance(&mut self, info: InstanceInfo) {
         // Sync to KvStore
         let data = InstanceData {
             app_id: info.app_id.clone(),
             ip: info.ip,
             public_key: info.public_key.clone(),
             reg_time: encode_ts(info.reg_time),
+            port_attrs: info.port_attrs.clone(),
         };
         if let Err(err) = self.kv_store.sync_instance(&info.id, &data) {
             error!("failed to sync instance to KvStore: {err:?}");
@@ -921,13 +972,15 @@ impl ProxyState {
             return Ok(smallvec![AddressInfo {
                 ip: Ipv4Addr::new(127, 0, 0, 1),
                 counter: Default::default(),
+                instance_id: "localhost".to_string(),
             }]);
         }
         let n = self.config.proxy.connect_top_n;
         if let Some(instance) = self.state.instances.get(id) {
             return Ok(smallvec![AddressInfo {
                 ip: instance.ip,
                 counter: instance.connections.clone(),
+                instance_id: instance.id.clone(),
             }]);
         };
         let app_instances = self.state.apps.get(id).context("app not found")?;
@@ -955,15 +1008,24 @@ impl ProxyState {
                 .filter_map(|instance_id| {
                     let instance = self.state.instances.get(instance_id)?;
                     let (_, elapsed) = handshakes.get(&instance.public_key)?;
-                    Some((instance.ip, *elapsed, instance.connections.clone()))
+                    Some((
+                        instance.ip,
+                        *elapsed,
+                        instance.connections.clone(),
+                        instance.id.clone(),
+                    ))
                 })
                 .collect::<SmallVec<[_; 4]>>(),
         };
         instances.sort_by(|a, b| a.1.cmp(&b.1));
         instances.truncate(n);
         Ok(instances
             .into_iter()
-            .map(|(ip, _, counter)| AddressInfo { ip, counter })
+            .map(|(ip, _, counter, instance_id)| AddressInfo {
+                ip,
+                counter,
+                instance_id,
+            })
             .collect())
     }
 
@@ -973,6 +1035,7 @@ impl ProxyState {
             return Some(smallvec![AddressInfo {
                 ip: info.ip,
                 counter: info.connections.clone(),
+                instance_id: info.id.clone(),
             }]);
         }
 
@@ -999,6 +1062,7 @@ impl ProxyState {
             smallvec![AddressInfo {
                 ip: info.ip,
                 counter: info.connections.clone(),
+                instance_id: info.id.clone(),
             }]
         })
     }
@@ -1271,8 +1335,14 @@ impl GatewayRpc for RpcHandler {
             .context("App authorization failed")?;
         let app_id = hex::encode(&app_info.app_id);
         let instance_id = hex::encode(&app_info.instance_id);
+        let port_attrs = request.port_attrs.map(|list| {
+            list.attrs
+                .into_iter()
+                .map(|p| (p.port as u16, PortFlags { pp: p.pp }))
+                .collect::<BTreeMap<u16, PortFlags>>()
+        });
         self.state
-            .do_register_cvm(&app_id, &instance_id, &request.client_public_key)
+            .do_register_cvm(&app_id, &instance_id, &request.client_public_key, port_attrs)
     }
 
     async fn acme_info(self) -> Result<AcmeInfoResponse> {

gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config-2.snap

@@ -12,5 +12,6 @@ InstanceInfo {
         tv_sec: 0,
         tv_nsec: 0,
     },
+    port_attrs: None,
     connections: 0,
 }

gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config.snap

@@ -12,5 +12,6 @@ InstanceInfo {
         tv_sec: 0,
         tv_nsec: 0,
     },
+    port_attrs: None,
     connections: 0,
 }