gw: implement PROXY protocol

Cargo.toml

@@ -247,6 +247,7 @@ yaml-rust2 = "0.10.4"
 luks2 = "0.5.0"
 scopeguard = "1.2.0"
 tar = "0.4"
+proxy-protocol = "0.5.0"
 
 [profile.release]
 panic = "abort"

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/Cargo.toml

@@ -54,6 +54,7 @@ hyper-rustls.workspace = true
 http-body-util.workspace = true
 x509-parser.workspace = true
 jemallocator.workspace = true
+proxy-protocol.workspace = true
 wavekv.workspace = true
 tdx-attest.workspace = true
 flate2.workspace = true

gateway/gateway.toml

@@ -58,6 +58,8 @@ workers = 32
 external_port = 443
 # Maximum concurrent connections per app. 0 means unlimited.
 max_connections_per_app = 2000
+# Whether to read PROXY protocol from inbound connections (e.g. from Cloudflare).
+inbound_pp_enabled = false
 
 [core.proxy.timeouts]
 # Timeout for establishing a connection to the target app.
@@ -81,6 +83,8 @@ write = "5s"
 shutdown = "5s"
 # Timeout for total connection duration.
 total = "5h"
+# Timeout for proxy protocol header.
+pp_header = "5s"
 
 [core.recycle]
 enabled = true

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,6 +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,
 }
 
 #[derive(Debug, Clone, Deserialize)]
@@ -142,6 +149,9 @@ pub struct Timeouts {
     pub write: Duration,
     #[serde(with = "serde_duration")]
     pub shutdown: Duration,
+    /// Timeout for reading the proxy protocol header from inbound connections.
+    #[serde(with = "serde_duration")]
+    pub pp_header: Duration,
 }
 
 #[derive(Debug, Clone, Deserialize, Serialize)]

gateway/src/debug_service.rs

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

gateway/src/kv/mod.rs

@@ -42,13 +42,31 @@ 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>>,
+    /// Hex-encoded compose_hash that `port_attrs` was learned against.
+    /// When a re-registration presents a different compose_hash (app upgrade),
+    /// the cache is invalidated and re-fetched lazily.
+    #[serde(default)]
+    pub port_attrs_hash: String,
 }
 
 /// Gateway node status (stored separately for independent updates)

gateway/src/main.rs

@@ -32,6 +32,7 @@ mod distributed_certbot;
 mod kv;
 mod main_service;
 mod models;
+mod pp;
 mod proxy;
 mod web_routes;