Dstack-TEE / Dstack-TEE/dstack
Support PROXY protocol for source IP preservation in dstack-gateway
- Dominant language
- Rust
- Stars
- 544
- Forks
- 96
- Avg merge
- 17h 57m
- Merged PRs (30d)
- 117
Description
## 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: `[-][s|g].`
- 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](https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt) 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:
```json
{
"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`
```protobuf
message RegisterCvmRequest {
string client_public_key = 1;
repeated uint32 proxy_protocol_ports = 2; // NEW
}
```
### 2. Update AppCompose Struct
**File:** `dstack-types/src/lib.rs`
```rust
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct AppCompose {
// ... existing fields ...
#[serde(default)]
pub proxy_protocol_ports: Vec, // NEW
}
```
### 3. Update InstanceInfo
**File:** `gateway/src/models.rs` or `gateway/src/main_service.rs`
```rust
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,
pub proxy_protocol_ports: Vec, // NEW
}
```
### 4. Pass Config During Registration
**File:** `dstack-util/src/system_setup.rs:742`
```rust
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)
```rust
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:
```rust
fn generate_proxy_v2_header(
src_ip: IpAddr,
src_port: u16,
dst_ip: IpAddr,
dst_port: u16,
) -> Vec {
// 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:
```rust
pub(crate) async fn proxy_to_app(
state: Proxy,
inbound: TcpStream,
buffer: Vec,
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:**
```nginx
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](https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt)
- [HAProxy Documentation](https://www.haproxy.com/blog/use-the-proxy-protocol-to-preserve-a-clients-ip-address)
- [AWS NLB PROXY Protocol v2](https://aws.amazon.com/blogs/networking-and-content-delivery/preserving-client-ip-address-with-proxy-protocol-v2-and-network-load-balancer/)
- [nginx PROXY Protocol Support](https://docs.nginx.com/nginx/admin-guide/load-balancer/using-proxy-protocol/)
## 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
Contributor guide
Assessment
This issue has not been assessed yet.