gw: address Copilot review (port range, async fetcher, pp tests)

Three fixes from review:

1. Treat the wire-format `port: uint32` as out-of-range when it can't fit
   in u16 (instead of silently truncating to a different valid port). Use
   `u16::try_from` and skip invalid entries.

2. Move the legacy `Info()` lazy fetch off the connection critical path:
   - `should_send_pp` is now sync. On a cache hit it returns the declared
     value; on a miss it enqueues the instance for the background worker
     and returns `pp = false` immediately, so a slow/missing CVM agent
     never blocks a proxied connection.
   - A single background task (`spawn_fetcher`) drains the queue, dedupes
     in-flight instance ids via a HashSet, applies a configurable
     timeout (`timeouts.port_attrs_fetch`, default 10s), and writes the
     result back to WaveKV.

3. Add unit tests in `pp.rs` for the inbound PROXY parser: v1/v2 IPv4
   happy paths, no-prefix rejection, v1 missing terminator, v2
   over-length cap, and the address synthesis/Display helpers.

Author: kvinwang

gateway/gateway.toml

@@ -85,6 +85,8 @@ shutdown = "5s"
 total = "5h"
 # Timeout for proxy protocol header.
 pp_header = "5s"
+# Timeout for the background lazy fetch of port_attrs from a legacy CVM.
+port_attrs_fetch = "10s"
 
 [core.recycle]
 enabled = true

gateway/src/config.rs

@@ -152,6 +152,10 @@ pub struct Timeouts {
     /// Timeout for reading the proxy protocol header from inbound connections.
     #[serde(with = "serde_duration")]
     pub pp_header: Duration,
+    /// Timeout for the background lazy fetch of `port_attrs` from a legacy CVM
+    /// agent's `Info()` RPC.
+    #[serde(with = "serde_duration")]
+    pub port_attrs_fetch: Duration,
 }
 
 #[derive(Debug, Clone, Deserialize, Serialize)]

gateway/src/main_service.rs

@@ -75,6 +75,10 @@ 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_attrs lazy-fetch worker. The proxy fast
+    /// path enqueues unknown instance_ids and immediately returns `pp=false`
+    /// so a missing cache never blocks a connection.
+    pub(crate) port_attrs_tx: tokio::sync::mpsc::UnboundedSender<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, Default)]
@@ -103,9 +107,13 @@ pub struct ProxyOptions {
 
 impl Proxy {
     pub async fn new(options: ProxyOptions) -> Result<Self> {
-        Ok(Self {
-            _inner: Arc::new(ProxyInner::new(options).await?),
-        })
+        let (port_attrs_tx, port_attrs_rx) = tokio::sync::mpsc::unbounded_channel();
+        let inner = ProxyInner::new(options, port_attrs_tx).await?;
+        let proxy = Self {
+            _inner: Arc::new(inner),
+        };
+        crate::proxy::port_attrs::spawn_fetcher(proxy.clone(), port_attrs_rx);
+        Ok(proxy)
     }
 }
 
@@ -114,7 +122,10 @@ impl ProxyInner {
         self.state.lock().or_panic("Failed to lock AppState")
     }
 
-    pub async fn new(options: ProxyOptions) -> Result<Self> {
+    pub async fn new(
+        options: ProxyOptions,
+        port_attrs_tx: tokio::sync::mpsc::UnboundedSender<String>,
+    ) -> Result<Self> {
         let ProxyOptions {
             config,
             my_app_id,
@@ -270,6 +281,7 @@ impl ProxyInner {
             kv_store,
             wavekv_sync,
             https_config: Some(https_config),
+            port_attrs_tx,
         })
     }
 
@@ -1369,7 +1381,12 @@ impl GatewayRpc for RpcHandler {
         let port_attrs = request.port_attrs.map(|list| {
             list.attrs
                 .into_iter()
-                .map(|p| (p.port as u16, PortFlags { pp: p.pp }))
+                .filter_map(|p| {
+                    // Wire format is uint32 to avoid varint shenanigans, but valid TCP
+                    // ports fit in u16. Drop out-of-range entries instead of truncating.
+                    let port = u16::try_from(p.port).ok()?;
+                    Some((port, PortFlags { pp: p.pp }))
+                })
                 .collect::<BTreeMap<u16, PortFlags>>()
         });
         self.state.do_register_cvm(

gateway/src/pp.rs

@@ -171,8 +171,135 @@ where
         }
     }
     if !end_found {
-        bail!("no valid proxy protocol header detected");
+        bail!("no valid proxy protocol header detected (v1 terminator not found)");
     }
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use proxy_protocol::{version1 as v1, version2 as v2, ProxyHeader};
+
+    fn extract_v4(header: ProxyHeader) -> (std::net::SocketAddrV4, std::net::SocketAddrV4) {
+        match header {
+            ProxyHeader::Version1 {
+                addresses:
+                    v1::ProxyAddresses::Ipv4 {
+                        source,
+                        destination,
+                    },
+                ..
+            } => (source, destination),
+            ProxyHeader::Version2 {
+                addresses:
+                    v2::ProxyAddresses::Ipv4 {
+                        source,
+                        destination,
+                    },
+                ..
+            } => (source, destination),
+            other => panic!("expected ipv4 header, got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn parses_v1_ipv4() {
+        // v1 is ASCII: "PROXY TCP4 <src> <dst> <sport> <dport>\r\n"
+        let header = b"PROXY TCP4 1.2.3.4 5.6.7.8 11111 22222\r\n";
+        let (_stream, parsed) = read_proxy_header(&header[..]).await.expect("v1 parse");
+        let (src, dst) = extract_v4(parsed);
+        assert_eq!(src.ip().octets(), [1, 2, 3, 4]);
+        assert_eq!(src.port(), 11111);
+        assert_eq!(dst.ip().octets(), [5, 6, 7, 8]);
+        assert_eq!(dst.port(), 22222);
+    }
+
+    #[tokio::test]
+    async fn parses_v2_ipv4() {
+        // v2 magic + ver/cmd 0x21 + family/proto 0x11 (TCP/IPv4) + len 12
+        let mut header = Vec::new();
+        header.extend_from_slice(V2_PROTOCOL_PREFIX);
+        header.extend_from_slice(&[0x21, 0x11, 0x00, 0x0c]);
+        header.extend_from_slice(&[1, 2, 3, 4]); // src ip
+        header.extend_from_slice(&[5, 6, 7, 8]); // dst ip
+        header.extend_from_slice(&11111u16.to_be_bytes()); // src port
+        header.extend_from_slice(&22222u16.to_be_bytes()); // dst port
+
+        let (_stream, parsed) = read_proxy_header(&header[..]).await.expect("v2 parse");
+        let (src, dst) = extract_v4(parsed);
+        assert_eq!(src.ip().octets(), [1, 2, 3, 4]);
+        assert_eq!(src.port(), 11111);
+        assert_eq!(dst.ip().octets(), [5, 6, 7, 8]);
+        assert_eq!(dst.port(), 22222);
+    }
+
+    #[tokio::test]
+    async fn rejects_no_prefix() {
+        // Looks neither like v1 ("PROXY") nor v2 magic.
+        let bytes = b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n";
+        let err = read_proxy_header(&bytes[..]).await.unwrap_err();
+        assert!(
+            format!("{err:#}").contains("no valid proxy protocol header"),
+            "unexpected error: {err:#}"
+        );
+    }
+
+    #[tokio::test]
+    async fn rejects_v1_without_terminator() {
+        // PROXY prefix matched but no \r\n terminator within V1_MAX_LENGTH bytes.
+        let bytes = vec![b'P'; V1_MAX_LENGTH + 8]; // all 'P' — never closes
+        let mut head = b"PROXY".to_vec();
+        head.extend(std::iter::repeat(b'A').take(V1_MAX_LENGTH));
+        let err = read_proxy_header(&head[..]).await.unwrap_err();
+        let msg = format!("{err:#}");
+        assert!(
+            msg.contains("v1 terminator not found") || msg.contains("no valid proxy"),
+            "unexpected error: {msg}"
+        );
+        // Sanity: the longer no-terminator buffer would also fail (read past)
+        let _ = bytes;
+    }
+
+    #[tokio::test]
+    async fn rejects_v2_oversize_length() {
+        // v2 prefix with a length field exceeding V2_MAX_LENGTH.
+        let mut header = Vec::new();
+        header.extend_from_slice(V2_PROTOCOL_PREFIX);
+        header.extend_from_slice(&[0x21, 0x11]);
+        // length = V2_MAX_LENGTH bytes -> total = MIN + that, blows the cap
+        header.extend_from_slice(&(V2_MAX_LENGTH as u16).to_be_bytes());
+        let err = read_proxy_header(&header[..]).await.unwrap_err();
+        assert!(
+            format!("{err:#}").contains("too long"),
+            "unexpected error: {err:#}"
+        );
+    }
+
+    #[test]
+    fn synthesizes_unspec_when_no_addrs() {
+        // We can't construct a real TcpStream in a unit test cheaply; just
+        // assert the helper returns Unspec for the all-None branch by going
+        // through the public Display impl.
+        let header = ProxyHeader::Version2 {
+            command: v2::ProxyCommand::Proxy,
+            transport_protocol: v2::ProxyTransportProtocol::Stream,
+            addresses: v2::ProxyAddresses::Unspec,
+        };
+        assert_eq!(format!("{}", DisplayAddr(&header)), "<unspec>");
+    }
+
+    #[test]
+    fn display_v2_ipv4_source() {
+        let header = ProxyHeader::Version2 {
+            command: v2::ProxyCommand::Proxy,
+            transport_protocol: v2::ProxyTransportProtocol::Stream,
+            addresses: v2::ProxyAddresses::Ipv4 {
+                source: "9.8.7.6:1234".parse().unwrap(),
+                destination: "1.2.3.4:80".parse().unwrap(),
+            },
+        };
+        assert_eq!(format!("{}", DisplayAddr(&header)), "9.8.7.6:1234");
+    }
+}

gateway/src/proxy.rs

@@ -42,7 +42,7 @@ pub(crate) struct AddressInfo {
 pub(crate) type AddressGroup = smallvec::SmallVec<[AddressInfo; 4]>;
 
 mod io_bridge;
-mod port_attrs;
+pub(crate) mod port_attrs;
 mod sni;
 mod tls_passthough;
 mod tls_terminate;

gateway/src/proxy/port_attrs.rs

@@ -2,56 +2,93 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-//! Per-port attribute lookup with lazy fetch from legacy CVMs.
+//! Per-port attribute lookup with background lazy fetch from legacy CVMs.
+//!
+//! Two paths:
+//!
+//! - Fast path (`should_send_pp`): a synchronous, non-blocking lookup. On a
+//!   cache miss it enqueues the instance for the background worker and
+//!   optimistically returns `pp = false` so the connection isn't blocked.
+//! - Slow path ([`spawn_fetcher`]): a single background task that drains the
+//!   queue, dedupes in-flight instances, calls the agent's `Info()` RPC with
+//!   a timeout, and writes the result back to WaveKV.
 
-use std::{collections::BTreeMap, net::Ipv4Addr};
+use std::collections::{BTreeMap, HashSet};
+use std::net::Ipv4Addr;
+use std::sync::{Arc, Mutex};
 
 use anyhow::{Context, Result};
 use dstack_guest_agent_rpc::dstack_guest_client::DstackGuestClient;
 use dstack_types::AppCompose;
 use http_client::prpc::PrpcClient;
+use tokio::sync::mpsc::UnboundedReceiver;
 use tracing::{debug, warn};
 
 use crate::{kv::PortFlags, main_service::Proxy};
 
 /// Decide whether the gateway should send a PROXY protocol header on the
 /// outbound connection to (`instance_id`, `port`).
 ///
-/// Lookup order:
-///
-/// 1. In-memory `port_attrs` populated at registration (new CVMs).
-/// 2. Lazy fetch via the agent's `Info()` RPC (legacy CVMs that didn't report
-///    attributes at registration). Result is cached on success.
-/// 3. Default `false` on any failure.
-pub(crate) async fn should_send_pp(state: &Proxy, instance_id: &str, port: u16) -> bool {
+/// Cache hit returns the declared value. Cache miss returns `false` and asks
+/// the background worker to populate the cache for the next request — this
+/// keeps the data path off the critical Info() RPC.
+pub(crate) fn should_send_pp(state: &Proxy, instance_id: &str, port: u16) -> bool {
     if let Some(attrs) = state.lock().instance_port_attrs(instance_id) {
         return attrs.get(&port).map(|f| f.pp).unwrap_or(false);
     }
-    match lazy_fetch(state, instance_id).await {
-        Ok(attrs) => attrs.get(&port).map(|f| f.pp).unwrap_or(false),
-        Err(err) => {
-            warn!("port_attrs lazy fetch for instance {instance_id} failed: {err:#}");
-            false
+    // Best-effort enqueue. If the channel is closed (shutdown) or the worker
+    // is gone, silently drop — `false` is the conservative default anyway.
+    let _ = state.port_attrs_tx.send(instance_id.to_string());
+    false
+}
+
+/// Spawn the background lazy-fetch worker. Should be called once at startup.
+pub(crate) fn spawn_fetcher(state: Proxy, mut rx: UnboundedReceiver<String>) {
+    let in_flight: Arc<Mutex<HashSet<String>>> = Default::default();
+    let timeout = state.config.proxy.timeouts.port_attrs_fetch;
+    tokio::spawn(async move {
+        while let Some(instance_id) = rx.recv().await {
+            // Dedupe: only spawn one fetch per instance at a time. After the
+            // fetch completes, the entry is removed so a later registration
+            // (with new compose_hash) can trigger a fresh fetch via the same
+            // path.
+            {
+                let mut in_flight = in_flight.lock().expect("port_attrs in_flight poisoned");
+                if !in_flight.insert(instance_id.clone()) {
+                    continue;
+                }
+            }
+            let state = state.clone();
+            let in_flight = in_flight.clone();
+            let id = instance_id.clone();
+            tokio::spawn(async move {
+                match tokio::time::timeout(timeout, fetch_and_store(&state, &id)).await {
+                    Ok(Ok(())) => debug!("port_attrs cached for instance {id}"),
+                    Ok(Err(err)) => {
+                        warn!("port_attrs fetch for instance {id} failed: {err:#}")
+                    }
+                    Err(_) => {
+                        warn!("port_attrs fetch for instance {id} timed out after {timeout:?}")
+                    }
+                }
+                in_flight
+                    .lock()
+                    .expect("port_attrs in_flight poisoned")
+                    .remove(&id);
+            });
         }
-    }
+    });
 }
 
-async fn lazy_fetch(state: &Proxy, instance_id: &str) -> Result<BTreeMap<u16, PortFlags>> {
+async fn fetch_and_store(state: &Proxy, instance_id: &str) -> Result<()> {
     let (ip, agent_port) = {
         let guard = state.lock();
         let ip = guard.instance_ip(instance_id).context("unknown instance")?;
         (ip, guard.config.proxy.agent_port)
     };
-
     let attrs = fetch_port_attrs(ip, agent_port).await?;
-    state
-        .lock()
-        .update_instance_port_attrs(instance_id, attrs.clone());
-    debug!(
-        "fetched port_attrs for legacy instance {instance_id} via Info(): {} entries",
-        attrs.len()
-    );
-    Ok(attrs)
+    state.lock().update_instance_port_attrs(instance_id, attrs);
+    Ok(())
 }
 
 async fn fetch_port_attrs(ip: Ipv4Addr, agent_port: u16) -> Result<BTreeMap<u16, PortFlags>> {

gateway/src/proxy/tls_passthough.rs

@@ -195,7 +195,7 @@ pub(crate) async fn proxy_to_app(
     .await
     .with_context(|| format!("connecting timeout to app {app_id}: {addresses:?}:{port}"))?
     .with_context(|| format!("failed to connect to app {app_id}: {addresses:?}:{port}"))?;
-    if should_send_pp(&state, &instance_id, port).await {
+    if should_send_pp(&state, &instance_id, port) {
         let pp_header_bin =
             proxy_protocol::encode(pp_header).context("failed to encode pp header")?;
         outbound.write_all(&pp_header_bin).await?;

gateway/src/proxy/tls_terminate.rs

@@ -299,7 +299,7 @@ impl Proxy {
         .await
         .map_err(|_| anyhow!("connecting timeout"))?
         .context("failed to connect to app")?;
-        if should_send_pp(self, &instance_id, port).await {
+        if should_send_pp(self, &instance_id, port) {
             let pp_header_bin =
                 proxy_protocol::encode(pp_header).context("failed to encode pp header")?;
             outbound.write_all(&pp_header_bin).await?;