gw: admin RPC to override per-instance port policy

Operators can now override an instance's port policy through the Admin
service, taking precedence over what the instance reports for itself.

Three new methods on the Admin service:

- SetInstancePortPolicy(instance_id, policy)
  Persists an admin override on the instance record. Errors with
  "instance not found" (404 semantics) if the instance is not
  registered — operators don't pre-create policy rows.

- ClearInstancePortPolicy(instance_id)
  Removes the override. Effective policy reverts to whatever the
  instance reported (or none, fail-close again).

- GetInstancePortPolicy(instance_id) -> { effective, source,
  instance_reported, admin_override }
  Returns both layers and the resolved effective policy. The `source`
  string ("admin" | "instance" | "none") makes "why was port X
  rejected" answerable in one call.

Storage and resolution:

- `InstanceData.admin_port_policy: Option<PortPolicy>` lives next to
  the existing `port_policy` field, persisted to WaveKV under the
  same `inst/{id}` record so overrides sync across nodes for free.
- `instance_port_policy()` (used by both `is_port_allowed` and
  `should_send_pp`) now returns `admin_port_policy.or(port_policy)`.
- App upgrade (compose_hash change) clears the instance-reported
  policy as before, but the admin override survives. Operator intent
  is stronger than app updates; they can re-evaluate explicitly.
- Lazy fetch only ever writes `port_policy`, never touches the
  override.

Tests:

- 5 new unit tests: override beats instance, override can open what
  the instance restricts (and vice versa), clearing reverts, unknown
  instance errors, override survives compose_hash change.

Author: kvinwang

gateway/rpc/proto/gateway_rpc.proto

@@ -439,6 +439,18 @@ service Admin {
   rpc GetCertbotConfig(google.protobuf.Empty) returns (CertbotConfigResponse) {}
   // Set global certbot configuration (includes ACME URL)
   rpc SetCertbotConfig(SetCertbotConfigRequest) returns (google.protobuf.Empty) {}
+
+  // ==================== Per-Instance Port Policy Override ====================
+  // Set an admin override for an instance's port policy. Takes precedence
+  // over the policy reported by the instance itself, and survives app
+  // upgrades. Errors if the instance is not registered.
+  rpc SetInstancePortPolicy(SetInstancePortPolicyRequest) returns (google.protobuf.Empty) {}
+  // Clear the admin override for an instance, reverting to the
+  // instance-reported policy. Errors if the instance is not registered.
+  rpc ClearInstancePortPolicy(ClearInstancePortPolicyRequest) returns (google.protobuf.Empty) {}
+  // Inspect both the admin override and the instance-reported policy for an
+  // instance, plus the effective policy the proxy will enforce.
+  rpc GetInstancePortPolicy(GetInstancePortPolicyRequest) returns (GetInstancePortPolicyResponse) {}
 }
 
 // ==================== DNS Credential Messages ====================
@@ -648,3 +660,36 @@ message SetCertbotConfigRequest {
   // ACME server URL (empty means use default Let's Encrypt production)
   optional string acme_url = 4;
 }
+
+// ==================== Per-Instance Port Policy Override Messages ====================
+
+// Set an admin override for an instance.
+message SetInstancePortPolicyRequest {
+  // The instance to override.
+  string instance_id = 1;
+  // The policy to apply. An empty `ports` list with `restrict_mode = true`
+  // is a valid "deny everything" lockdown.
+  PortPolicy policy = 2;
+}
+
+// Clear the admin override for an instance.
+message ClearInstancePortPolicyRequest {
+  string instance_id = 1;
+}
+
+// Inspect an instance's port-policy state.
+message GetInstancePortPolicyRequest {
+  string instance_id = 1;
+}
+
+message GetInstancePortPolicyResponse {
+  // The policy the proxy will actually enforce. Absent when neither admin
+  // nor instance has set anything (fail-close until populated).
+  optional PortPolicy effective = 1;
+  // Where `effective` came from: "admin", "instance", or "none".
+  string source = 2;
+  // The policy reported by the instance itself, if any.
+  optional PortPolicy instance_reported = 3;
+  // The admin override, if any.
+  optional PortPolicy admin_override = 4;
+}

gateway/src/admin_service.rs

@@ -8,16 +8,18 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use anyhow::{bail, Context, Result};
 use dstack_gateway_rpc::{
     admin_server::{AdminRpc, AdminServer},
-    CertAttestationInfo, CertbotConfigResponse, CreateDnsCredentialRequest,
-    DeleteDnsCredentialRequest, DeleteZtDomainRequest, DnsCredentialInfo,
-    ForceReleaseCertLockRequest, GetDefaultDnsCredentialResponse, GetDnsCredentialRequest,
-    GetInfoRequest, GetInfoResponse, GetInstanceHandshakesRequest, GetInstanceHandshakesResponse,
+    CertAttestationInfo, CertbotConfigResponse, ClearInstancePortPolicyRequest,
+    CreateDnsCredentialRequest, DeleteDnsCredentialRequest, DeleteZtDomainRequest,
+    DnsCredentialInfo, ForceReleaseCertLockRequest, GetDefaultDnsCredentialResponse,
+    GetDnsCredentialRequest, GetInfoRequest, GetInfoResponse, GetInstanceHandshakesRequest,
+    GetInstanceHandshakesResponse, GetInstancePortPolicyRequest, GetInstancePortPolicyResponse,
     GetMetaResponse, GetNodeStatusesResponse, GetZtDomainRequest, GlobalConnectionsStats,
     HandshakeEntry, HostInfo, LastSeenEntry, ListCertAttestationsRequest,
     ListCertAttestationsResponse, ListDnsCredentialsResponse, ListZtDomainsResponse,
-    NodeStatusEntry, PeerSyncStatus as ProtoPeerSyncStatus, RenewCertResponse,
-    RenewZtDomainCertRequest, RenewZtDomainCertResponse, SetCertbotConfigRequest,
-    SetDefaultDnsCredentialRequest, SetNodeStatusRequest, SetNodeUrlRequest, StatusResponse,
+    NodeStatusEntry, PeerSyncStatus as ProtoPeerSyncStatus, PortAttrs as RpcPortAttrs,
+    PortPolicy as RpcPortPolicy, RenewCertResponse, RenewZtDomainCertRequest,
+    RenewZtDomainCertResponse, SetCertbotConfigRequest, SetDefaultDnsCredentialRequest,
+    SetInstancePortPolicyRequest, SetNodeStatusRequest, SetNodeUrlRequest, StatusResponse,
     StoreSyncStatus, UpdateDnsCredentialRequest, WaveKvStatusResponse, ZtDomainCertStatus,
     ZtDomainConfig as ProtoZtDomainConfig, ZtDomainInfo,
 };
@@ -26,8 +28,9 @@ use tracing::info;
 use wavekv::node::NodeStatus as WaveKvNodeStatus;
 
 use crate::{
-    kv::{DnsCredential, DnsProvider, NodeStatus, ZtDomainConfig},
+    kv::{DnsCredential, DnsProvider, NodeStatus, PortFlags, PortPolicy, ZtDomainConfig},
     main_service::Proxy,
+    models::PortPolicyView,
     proxy::NUM_CONNECTIONS,
 };
 
@@ -625,6 +628,73 @@ impl AdminRpc for AdminRpcHandler {
         );
         Ok(())
     }
+
+    async fn set_instance_port_policy(self, request: SetInstancePortPolicyRequest) -> Result<()> {
+        let proto = request.policy.context("port policy is required")?;
+        let policy = port_policy_from_proto(proto)?;
+        self.state
+            .lock()
+            .set_admin_port_policy(&request.instance_id, policy)
+    }
+
+    async fn clear_instance_port_policy(
+        self,
+        request: ClearInstancePortPolicyRequest,
+    ) -> Result<()> {
+        self.state
+            .lock()
+            .clear_admin_port_policy(&request.instance_id)
+    }
+
+    async fn get_instance_port_policy(
+        self,
+        request: GetInstancePortPolicyRequest,
+    ) -> Result<GetInstancePortPolicyResponse> {
+        let view = self
+            .state
+            .lock()
+            .instance_port_policy_view(&request.instance_id)
+            .with_context(|| format!("instance {} not found", request.instance_id))?;
+        Ok(port_policy_view_to_proto(view))
+    }
+}
+
+fn port_policy_from_proto(proto: RpcPortPolicy) -> Result<PortPolicy> {
+    let mut ports = std::collections::BTreeMap::new();
+    for attr in proto.ports {
+        let port = u16::try_from(attr.port)
+            .with_context(|| format!("port {} out of u16 range", attr.port))?;
+        ports.insert(port, PortFlags { pp: attr.pp });
+    }
+    Ok(PortPolicy {
+        ports,
+        restrict_mode: proto.restrict_mode,
+    })
+}
+
+fn port_policy_to_proto(policy: &PortPolicy) -> RpcPortPolicy {
+    RpcPortPolicy {
+        ports: policy
+            .ports
+            .iter()
+            .map(|(port, flags)| RpcPortAttrs {
+                port: u32::from(*port),
+                pp: flags.pp,
+            })
+            .collect(),
+        restrict_mode: policy.restrict_mode,
+    }
+}
+
+fn port_policy_view_to_proto(view: PortPolicyView) -> GetInstancePortPolicyResponse {
+    let source = view.source().to_string();
+    let effective = view.effective().map(port_policy_to_proto);
+    GetInstancePortPolicyResponse {
+        effective,
+        source,
+        instance_reported: view.instance_reported.as_ref().map(port_policy_to_proto),
+        admin_override: view.admin_override.as_ref().map(port_policy_to_proto),
+    }
 }
 
 fn build_store_status(

gateway/src/kv/mod.rs

@@ -81,6 +81,12 @@ pub struct InstanceData {
     /// the cache is invalidated and re-fetched lazily.
     #[serde(default)]
     pub port_policy_hash: String,
+    /// Operator-set override applied via the Admin RPC. Takes precedence over
+    /// the instance-reported `port_policy` when set, and survives app upgrades
+    /// (compose_hash changes do not clear it). Cleared explicitly via
+    /// ClearInstancePortPolicy.
+    #[serde(default)]
+    pub admin_port_policy: Option<PortPolicy>,
 }
 
 /// Gateway node status (stored separately for independent updates)

gateway/src/main_service.rs

@@ -42,7 +42,7 @@ use crate::{
         fetch_peers_from_bootnode, AppIdValidator, HttpsClientConfig, InstanceData, KvStore,
         NodeData, NodeStatus, PortPolicy, WaveKvSyncService,
     },
-    models::{InstanceInfo, WgConf},
+    models::{InstanceInfo, PortPolicyView, WgConf},
     proxy::{create_acceptor_with_cert_resolver, AddressGroup, AddressInfo},
 };
 
@@ -78,10 +78,12 @@ pub struct ProxyInner {
     pub(crate) wavekv_sync: Option<Arc<WaveKvSyncService>>,
     /// HTTPS client config for mTLS (used for bootnode peer discovery)
     https_config: Option<HttpsClientConfig>,
-    /// Sender for the background port_policy lazy-fetch worker. The proxy
-    /// data path enqueues unknown instance_ids and waits for the cache to
-    /// populate (fail-close): without a known policy, restrict_mode is
-    /// indeterminate and we cannot safely allow traffic.
+    /// Sender for the background port_policy lazy-fetch worker. On a cache
+    /// miss the proxy data path enqueues the instance_id and immediately
+    /// rejects the connection (fail-close); the fetch populates the cache
+    /// asynchronously so subsequent connections can proceed. Without a
+    /// known policy, `restrict_mode` is indeterminate and we cannot safely
+    /// allow traffic.
     pub(crate) port_policy_tx: UnboundedSender<String>,
 }
 
@@ -491,6 +493,7 @@ fn build_state_from_kv_store(instances: BTreeMap<String, InstanceData>) -> Proxy
                 .unwrap_or(UNIX_EPOCH),
             port_policy: data.port_policy,
             port_policy_hash: data.port_policy_hash,
+            admin_port_policy: data.admin_port_policy,
             connections: Default::default(),
         };
         state.allocated_addresses.insert(data.ip);
@@ -786,6 +789,7 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()>
                 .unwrap_or(UNIX_EPOCH),
             port_policy: data.port_policy.clone(),
             port_policy_hash: data.port_policy_hash.clone(),
+            admin_port_policy: data.admin_port_policy.clone(),
             connections: Default::default(),
         };
 
@@ -914,6 +918,7 @@ impl ProxyState {
                     reg_time: encode_ts(existing.reg_time),
                     port_policy: existing.port_policy.clone(),
                     port_policy_hash: existing.port_policy_hash.clone(),
+                    admin_port_policy: existing.admin_port_policy.clone(),
                 };
                 if let Err(err) = self.kv_store.sync_instance(&existing.id, &data) {
                     error!("failed to sync existing instance to KvStore: {err:?}");
@@ -934,6 +939,7 @@ impl ProxyState {
             reg_time: SystemTime::now(),
             port_policy,
             port_policy_hash: compose_hash.to_string(),
+            admin_port_policy: None,
             connections: Default::default(),
         };
         self.add_instance(host_info.clone());
@@ -945,10 +951,14 @@ impl ProxyState {
         self.state.instances.get(instance_id).map(|i| i.ip)
     }
 
-    /// Lookup an instance's port_policy. `None` means the CVM never reported
-    /// it (legacy CVM), so the caller should fall back to fetching via Info().
+    /// Lookup the effective port_policy for an instance: admin override wins,
+    /// otherwise fall back to the instance-reported policy. `None` means
+    /// neither is set — caller fails closed and may schedule a lazy fetch.
     pub(crate) fn instance_port_policy(&self, instance_id: &str) -> Option<&PortPolicy> {
-        self.state.instances.get(instance_id)?.port_policy.as_ref()
+        let info = self.state.instances.get(instance_id)?;
+        info.admin_port_policy
+            .as_ref()
+            .or(info.port_policy.as_ref())
     }
 
     /// Update an instance's port_policy (used after a lazy fetch via Info()).
@@ -957,17 +967,71 @@ impl ProxyState {
         let Some(info) = self.state.instances.get_mut(instance_id) else {
             return;
         };
-        info.port_policy = Some(policy.clone());
+        info.port_policy = Some(policy);
+        self.persist_instance_record(instance_id);
+    }
+
+    /// Snapshot view of an instance's port-policy state for inspection.
+    pub(crate) fn instance_port_policy_view(&self, instance_id: &str) -> Option<PortPolicyView> {
+        let info = self.state.instances.get(instance_id)?;
+        Some(PortPolicyView {
+            instance_reported: info.port_policy.clone(),
+            admin_override: info.admin_port_policy.clone(),
+        })
+    }
+
+    /// Set an admin override. Errors if the instance is not registered.
+    pub(crate) fn set_admin_port_policy(
+        &mut self,
+        instance_id: &str,
+        policy: PortPolicy,
+    ) -> Result<()> {
+        let Some(info) = self.state.instances.get_mut(instance_id) else {
+            bail!("instance {instance_id} not found");
+        };
+        let prev = info.admin_port_policy.take();
+        info.admin_port_policy = Some(policy.clone());
+        info!(
+            "admin set port_policy for instance {instance_id}: \
+             restrict_mode={}, ports={} (prev: {})",
+            policy.restrict_mode,
+            policy.ports.len(),
+            prev.is_some(),
+        );
+        self.persist_instance_record(instance_id);
+        Ok(())
+    }
+
+    /// Clear any admin override. Errors if the instance is not registered.
+    /// No-op (still Ok) if no override was set.
+    pub(crate) fn clear_admin_port_policy(&mut self, instance_id: &str) -> Result<()> {
+        let Some(info) = self.state.instances.get_mut(instance_id) else {
+            bail!("instance {instance_id} not found");
+        };
+        let had_override = info.admin_port_policy.take().is_some();
+        info!("admin cleared port_policy for instance {instance_id} (was set: {had_override})");
+        if had_override {
+            self.persist_instance_record(instance_id);
+        }
+        Ok(())
+    }
+
+    /// Persist the current in-memory `InstanceInfo` snapshot to WaveKV.
+    fn persist_instance_record(&self, instance_id: &str) {
+        let Some(info) = self.state.instances.get(instance_id) else {
+            return;
+        };
         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_policy: Some(policy),
+            port_policy: info.port_policy.clone(),
             port_policy_hash: info.port_policy_hash.clone(),
+            admin_port_policy: info.admin_port_policy.clone(),
         };
         if let Err(err) = self.kv_store.sync_instance(instance_id, &data) {
-            error!("failed to sync updated port_policy to KvStore: {err:?}");
+            error!("failed to sync instance {instance_id} to KvStore: {err:?}");
         }
     }
 
@@ -980,6 +1044,7 @@ impl ProxyState {
             reg_time: encode_ts(info.reg_time),
             port_policy: info.port_policy.clone(),
             port_policy_hash: info.port_policy_hash.clone(),
+            admin_port_policy: info.admin_port_policy.clone(),
         };
         if let Err(err) = self.kv_store.sync_instance(&info.id, &data) {
             error!("failed to sync instance to KvStore: {err:?}");

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

@@ -14,5 +14,6 @@ InstanceInfo {
     },
     port_policy: None,
     port_policy_hash: "",
+    admin_port_policy: None,
     connections: 0,
 }

gateway/src/main_service/snapshots/dstack_gateway__main_service__tests__config.snap

@@ -14,5 +14,6 @@ InstanceInfo {
     },
     port_policy: None,
     port_policy_hash: "",
+    admin_port_policy: None,
     connections: 0,
 }