Support PROXY protocol for source IP preservation in dstack-gateway

[h4x3rotab]

Problem

All traffic received by CVMs is routed through dstack-gateway, but the gateway cannot pass through the original source IP to the CVM. This causes several critical issues:

1. BitTorrent trackers cannot determine the source IP, breaking peer table maintenance
2. Policy-based load balancing cannot make routing decisions based on client IP
3. Firewall rules and access control lists cannot filter by source IP
4. Application logs show gateway IP instead of actual client IP
5. Rate limiting and abuse prevention become ineffective

Current Architecture

Traffic Flow

Client → dstack-gateway → WireGuard VPN → CVM

Gateway Routing

• Uses SNI-based routing: <app_id>[-<port>][s|g].<base_domain>
• Two modes:
  • TLS termination (default): Gateway terminates TLS, sees plaintext
  • TLS passthrough (s suffix): Gateway forwards encrypted traffic
• Traffic forwarding via bridge() function in gateway/src/proxy/io_bridge.rs
• Bridge does bidirectional TCP copy without preserving source IP metadata

Why Source IP is Lost

• Gateway acts as a reverse proxy / NAT
• Backend connections originate from gateway's WireGuard IP
• No mechanism to communicate original client IP to the CVM

Proposed Solution: PROXY Protocol v2

Implement support for the PROXY protocol to preserve client source IP.

What is PROXY Protocol?

A protocol developed by HAProxy to safely transport connection information (including client source IP) across multiple layers of NAT or TCP proxies. It adds a small header to the TCP connection containing the original client's address.

Versions:

• v1: Human-readable ASCII format (e.g., PROXY TCP4 192.168.0.1 192.168.0.11 56324 443\r\n)
• v2: Binary format with better performance and extensibility (recommended)

Industry adoption: HAProxy, nginx, Stunnel, Varnish, Squid, AWS NLB, etc.

Design Overview

1. Configuration (app-compose.json)

Add a new optional field to specify which ports need PROXY protocol:

{
  "manifest_version": 2,
  "name": "my-app",
  "gateway_enabled": true,
  "proxy_protocol_ports": [8080, 3000],
  "docker_compose_file": "..."
}

Behavior:

• Empty/omitted list = default behavior (no PROXY headers) ✅ Backward compatible
• Ports in list = gateway injects PROXY v2 header before forwarding

2. Registration Flow

CVM Side (dstack-util/src/system_setup.rs):

• Read proxy_protocol_ports from app-compose.json
• Pass it in RegisterCvmRequest during CVM registration

Gateway Side (gateway/src/main_service.rs):

• Store proxy_protocol_ports in InstanceInfo struct
• Use this info when routing traffic

3. Traffic Forwarding

Injection Point (gateway/src/proxy/io_bridge.rs):

When forwarding traffic to a CVM:

1. Check if destination port is in proxy_protocol_ports
2. If yes:
   • Extract real client IP from TCP connection
   • Generate PROXY protocol v2 header
   • Prepend header before forwarding data
3. If no:
   • Forward normally (current behavior)

Both modes supported:

• TLS termination: Inject after TLS handshake, before HTTP data
• TLS passthrough: Inject before TLS ClientHello

Implementation Details

1. Extend Protobuf Definition

File: gateway/rpc/proto/gateway_rpc.proto

message RegisterCvmRequest {
  string client_public_key = 1;
  repeated uint32 proxy_protocol_ports = 2;  // NEW
}

2. Update AppCompose Struct

File: dstack-types/src/lib.rs

#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct AppCompose {
    // ... existing fields ...
    #[serde(default)]
    pub proxy_protocol_ports: Vec<u16>,  // NEW
}

3. Update InstanceInfo

File: gateway/src/models.rs or gateway/src/main_service.rs

pub struct InstanceInfo {
    pub id: String,
    pub app_id: String,
    pub ip: Ipv4Addr,
    pub public_key: String,
    pub reg_time: SystemTime,
    pub last_seen: SystemTime,
    pub connections: Arc<AtomicU64>,
    pub proxy_protocol_ports: Vec<u16>,  // NEW
}

4. Pass Config During Registration

File: dstack-util/src/system_setup.rs:742

client.register_cvm(RegisterCvmRequest {
    client_public_key: wg_pk,
    proxy_protocol_ports: self.shared.app_compose.proxy_protocol_ports.clone(),  // NEW
})

5. Store Config in Gateway

File: gateway/src/main_service.rs:713 (RegisterCvm handler)

let host_info = InstanceInfo {
    // ... existing fields ...
    proxy_protocol_ports: request.proxy_protocol_ports.clone(),  // NEW
};

6. Inject PROXY Header

File: gateway/src/proxy/io_bridge.rs

Add new function to generate PROXY protocol v2 header:

fn generate_proxy_v2_header(
    src_ip: IpAddr,
    src_port: u16,
    dst_ip: IpAddr,
    dst_port: u16,
) -> Vec<u8> {
    // PROXY protocol v2 binary format
    // Signature: \x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A
    // Version + Command: \x21 (v2, PROXY command)
    // Family + Protocol: \x11 (IPv4, TCP) or \x21 (IPv6, TCP)
    // Length + Addresses
    // ... implementation details ...
}

Modify bridge functions:

• gateway/src/proxy/tls_terminate.rs:299 (TLS termination mode)
• gateway/src/proxy/tls_passthough.rs:120 (TLS passthrough mode)

Check port and inject header before forwarding:

pub(crate) async fn proxy_to_app(
    state: Proxy,
    inbound: TcpStream,
    buffer: Vec<u8>,
    app_id: &str,
    port: u16,
) -> Result<()> {
    let addresses = state.lock().select_top_n_hosts(app_id)?;
    let (mut outbound, _counter) = connect_multiple_hosts(addresses, port).await?;
    
    // NEW: Check if PROXY protocol is enabled for this port
    let instance = state.lock().state.instances.get(app_id);
    if let Some(info) = instance {
        if info.proxy_protocol_ports.contains(&port) {
            let src_addr = inbound.peer_addr()?;
            let dst_addr = outbound.peer_addr()?;
            let proxy_header = generate_proxy_v2_header(
                src_addr.ip(), src_addr.port(),
                dst_addr.ip(), dst_addr.port(),
            );
            outbound.write_all(&proxy_header).await?;
        }
    }
    
    outbound.write_all(&buffer).await?;
    bridge(inbound, outbound, &state.config.proxy).await?;
    Ok(())
}

7. CVM Application Configuration

Applications must be configured to accept PROXY protocol:

nginx example:

server {
    listen 8080 proxy_protocol;
    real_ip_header proxy_protocol;
    
    location / {
        # $remote_addr now contains real client IP
    }
}

HAProxy example:

frontend web
    bind :8080 accept-proxy
    # Client IP preserved

Backward Compatibility

✅ Fully backward compatible:

• Default: proxy_protocol_ports is empty → no PROXY headers (current behavior)
• Only apps that explicitly configure it will get PROXY headers
• Existing apps continue working unchanged

⚠️ Important: Apps must explicitly support PROXY protocol on configured ports, otherwise they will see garbage data and fail.

Testing Plan

1. Unit tests: PROXY v2 header generation
2. Integration tests:
   • TLS termination mode with PROXY header
   • TLS passthrough mode with PROXY header
   • Mixed ports (some with PROXY, some without)
3. Real-world test:
   • Deploy nginx with proxy_protocol enabled
   • Verify logs show real client IP
   • Test BitTorrent tracker functionality

References

• PROXY Protocol Specification
• HAProxy Documentation
• AWS NLB PROXY Protocol v2
• nginx PROXY Protocol Support

Related Files

• gateway/src/proxy/io_bridge.rs - Main bridge logic
• gateway/src/proxy/tls_terminate.rs - TLS termination mode
• gateway/src/proxy/tls_passthough.rs - TLS passthrough mode
• gateway/src/main_service.rs:713 - RegisterCvm handler
• gateway/rpc/proto/gateway_rpc.proto - RPC definitions
• dstack-types/src/lib.rs - AppCompose struct
• dstack-util/src/system_setup.rs:742 - CVM registration

[Leechael]

Drafted.

#356

[kvinwang]

As for the configuration approach, adding args to RegisterCvm might not be the best way to go.
Another way is to add a flag in the zt-https domain, similar to https/grpc, we'll get:

• {app_id}-{port}s for tls passthrough
• {app_id}-{port}sp for tls passthrough over proxy protocol
• {app_id}-{port}gp for grpc over proxy protocol

	RegisterCvm Parameter	URL Flag
Architecture Complexity	❌ Stateful, increases Gateway complexity, version state sync is complex	✅ Stateless, simple architecture
Version Compatibility	❌ Requires OS image modification, old versions not supported	✅ Compatible with old OS versions
URL Aesthetics	✅ Clean URLs	❌ zt-https URLs contain flag identifiers
Custom Domain Configuration	❌ Configured via CVM parameters, configuration point far from usage point	✅ Configured via dstack-ingress parameters, close to usage (or even zero config)
Dynamic created Service Support	❌ Not supported	✅ Supported