gw: add restrict_mode port whitelist (fail-close)

Apps can now declare `restrict_mode: true` in app-compose to make the
gateway reject connections to any port not listed under `ports`. Combined
with this PR's existing PROXY protocol support, this gives apps explicit
control over which ports the gateway will forward.

Implementation:

- Collapse the previous `port_attrs: Option<BTreeMap>` and
  `port_attrs_hash: String` fields into a single `Option<PortPolicy>`
  carrying `{ ports, restrict_mode }`. One Option distinguishes
  "not reported" from "reported empty" cleanly.
- RPC: `RegisterCvmRequest.port_attrs` (PortAttrsList) becomes
  `port_policy` (PortPolicy { ports, restrict_mode }).
- Enforcement: `filter_allowed_addresses` runs before
  `connect_multiple_hosts` in both TLS-terminate and TLS-passthrough
  paths. A denied connection bubbles up as a normal error and the
  proxy closes the inbound TCP stream — no special HTTP response.
- Failure mode: fail-close. An unknown policy (cache miss) denies the
  connection and triggers a background fetch, so subsequent connections
  proceed once the policy is known. Apps that opt into restrict_mode
  must run a CVM that reports policy at registration time; legacy
  CVMs that fall back to the lazy `Info()` fetch get the open default
  (`restrict_mode = false`) so they keep working.
- Unknown instance_ids (e.g. the `localhost` shortcut) bypass the
  check — the policy machinery only applies to registered CVMs.
- Rename the lazy-fetch config block from `port_attrs_fetch` to
  `port_policy_fetch` for consistency (only ever shipped on this
  unmerged branch).

Tests:

- 4 new unit tests covering allow/deny, disabled mode, fail-close on
  unknown policy, and the unknown-instance bypass.
- Snapshot tests updated for the renamed fields.

Author: kvinwang

dstack-types/src/lib.rs

@@ -48,6 +48,10 @@ pub struct AppCompose {
     /// Per-port attributes consumed by the gateway (e.g. PROXY protocol).
     #[serde(default)]
     pub ports: Vec<PortAttrs>,
+    /// When true, the gateway only forwards traffic to ports listed in `ports`.
+    /// All other ports are rejected at TCP-accept time.
+    #[serde(default)]
+    pub restrict_mode: bool,
 }
 
 #[derive(Deserialize, Serialize, Debug, Clone)]

dstack-util/src/system_setup.rs

@@ -53,8 +53,8 @@ use crate::{
 use cert_client::CertRequestClient;
 use cmd_lib::run_fun as cmd;
 use dstack_gateway_rpc::{
-    gateway_client::GatewayClient, PortAttrs as RpcPortAttrs, PortAttrsList, RegisterCvmRequest,
-    RegisterCvmResponse, WireGuardPeer,
+    gateway_client::GatewayClient, PortAttrs as RpcPortAttrs, PortPolicy as RpcPortPolicy,
+    RegisterCvmRequest, RegisterCvmResponse, WireGuardPeer,
 };
 use ra_tls::rcgen::{KeyPair, PKCS_ECDSA_P256_SHA256};
 use serde_human_bytes as hex_bytes;
@@ -447,8 +447,8 @@ impl<'a> GatewayContext<'a> {
         gateway_url: &str,
         key_store: &GatewayKeyStore,
     ) -> Result<RegisterCvmResponse> {
-        let port_attrs = PortAttrsList {
-            attrs: self
+        let port_policy = RpcPortPolicy {
+            ports: self
                 .shared
                 .app_compose
                 .ports
@@ -458,13 +458,14 @@ impl<'a> GatewayContext<'a> {
                     pp: p.pp,
                 })
                 .collect(),
+            restrict_mode: self.shared.app_compose.restrict_mode,
         };
         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()),
+                port_policy: Some(port_policy.clone()),
             })
             .await
             .context("Failed to register CVM");
@@ -485,7 +486,7 @@ impl<'a> GatewayContext<'a> {
         client
             .register_cvm(RegisterCvmRequest {
                 client_public_key: key_store.wg_pk.clone(),
-                port_attrs: Some(port_attrs),
+                port_policy: Some(port_policy),
             })
             .await
             .context("Failed to register CVM")

gateway/dstack-app/builder/entrypoint.sh

@@ -125,11 +125,11 @@ shutdown = "${TIMEOUT_SHUTDOWN:-5s}"
 total = "${TIMEOUT_TOTAL:-5h}"
 pp_header = "${TIMEOUT_PP_HEADER:-5s}"
 
-[core.proxy.port_attrs_fetch]
-timeout = "${PORT_ATTRS_FETCH_TIMEOUT:-10s}"
-max_retries = ${PORT_ATTRS_FETCH_MAX_RETRIES:-5}
-backoff_initial = "${PORT_ATTRS_FETCH_BACKOFF_INITIAL:-1s}"
-backoff_max = "${PORT_ATTRS_FETCH_BACKOFF_MAX:-30s}"
+[core.proxy.port_policy_fetch]
+timeout = "${PORT_POLICY_FETCH_TIMEOUT:-10s}"
+max_retries = ${PORT_POLICY_FETCH_MAX_RETRIES:-5}
+backoff_initial = "${PORT_POLICY_FETCH_BACKOFF_INITIAL:-1s}"
+backoff_max = "${PORT_POLICY_FETCH_BACKOFF_MAX:-30s}"
 
 [core.recycle]
 enabled = true

gateway/dstack-app/docker-compose.yaml

@@ -43,10 +43,10 @@ services:
       - ADMIN_LISTEN_PORT=${ADMIN_LISTEN_PORT:-8001}
       - INBOUND_PP_ENABLED=${INBOUND_PP_ENABLED:-false}
       - TIMEOUT_PP_HEADER=${TIMEOUT_PP_HEADER:-5s}
-      - PORT_ATTRS_FETCH_TIMEOUT=${PORT_ATTRS_FETCH_TIMEOUT:-10s}
-      - PORT_ATTRS_FETCH_MAX_RETRIES=${PORT_ATTRS_FETCH_MAX_RETRIES:-5}
-      - PORT_ATTRS_FETCH_BACKOFF_INITIAL=${PORT_ATTRS_FETCH_BACKOFF_INITIAL:-1s}
-      - PORT_ATTRS_FETCH_BACKOFF_MAX=${PORT_ATTRS_FETCH_BACKOFF_MAX:-30s}
+      - PORT_POLICY_FETCH_TIMEOUT=${PORT_POLICY_FETCH_TIMEOUT:-10s}
+      - PORT_POLICY_FETCH_MAX_RETRIES=${PORT_POLICY_FETCH_MAX_RETRIES:-5}
+      - PORT_POLICY_FETCH_BACKOFF_INITIAL=${PORT_POLICY_FETCH_BACKOFF_INITIAL:-1s}
+      - PORT_POLICY_FETCH_BACKOFF_MAX=${PORT_POLICY_FETCH_BACKOFF_MAX:-30s}
     restart: always
 
 volumes:

gateway/gateway.toml

@@ -61,8 +61,8 @@ max_connections_per_app = 2000
 # Whether to read PROXY protocol from inbound connections (e.g. from Cloudflare).
 inbound_pp_enabled = false
 
-[core.proxy.port_attrs_fetch]
-# Background lazy-fetch of port_attrs from legacy CVM agents.
+[core.proxy.port_policy_fetch]
+# Background lazy-fetch of port_policy from legacy CVM agents.
 # Single Info() RPC timeout.
 timeout = "10s"
 # Retries cover the WireGuard / agent warmup window after registration.

gateway/rpc/proto/gateway_rpc.proto

@@ -12,16 +12,22 @@ 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.
+  // Per-port policy 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;
+  optional PortPolicy port_policy = 2;
+}
+
+// PortPolicy carries the gateway-relevant per-port configuration declared by
+// the app in its compose file. Keeping `ports` and `restrict_mode` together
+// lets a single Option distinguish "not reported" from "reported".
+message PortPolicy {
+  // Per-port attributes (PROXY protocol opt-in, etc.).
+  repeated PortAttrs ports = 1;
+  // When true, the gateway only forwards traffic to ports listed in `ports`
+  // and rejects connections to any other port at TCP-accept time.
+  bool restrict_mode = 2;
 }
 
 // PortAttrs declares per-port behaviour for the gateway.

gateway/src/config.rs

@@ -118,18 +118,18 @@ pub struct ProxyConfig {
     /// 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).
+    /// gateway to fetch app metadata (e.g. port_policy 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,
-    /// Background lazy-fetch behaviour for `port_attrs` (legacy CVMs).
-    pub port_attrs_fetch: PortAttrsFetchConfig,
+    /// Background lazy-fetch behaviour for `port_policy` (legacy CVMs).
+    pub port_policy_fetch: PortPolicyFetchConfig,
 }
 
 #[derive(Debug, Clone, Deserialize)]
-pub struct PortAttrsFetchConfig {
+pub struct PortPolicyFetchConfig {
     /// Timeout for a single `Info()` RPC attempt.
     #[serde(with = "serde_duration")]
     pub timeout: Duration,

gateway/src/kv/mod.rs

@@ -50,23 +50,37 @@ pub struct PortFlags {
     pub pp: bool,
 }
 
+/// Gateway-relevant per-port policy declared by the app in its compose file.
+/// Reported atomically at CVM registration; `Option<PortPolicy>` distinguishes
+/// "not reported" (legacy CVM) from "reported with no entries".
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
+pub struct PortPolicy {
+    /// Per-port flags (PROXY protocol opt-in, etc.).
+    #[serde(default)]
+    pub ports: BTreeMap<u16, PortFlags>,
+    /// When true, only ports listed in `ports` are forwarded; connections to
+    /// any other port are rejected at TCP-accept time.
+    #[serde(default)]
+    pub restrict_mode: 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"
+    /// Port policy 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.
+    pub port_policy: Option<PortPolicy>,
+    /// Hex-encoded compose_hash that `port_policy` 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,
+    pub port_policy_hash: String,
 }
 
 /// Gateway node status (stored separately for independent updates)