Dstack-TEE / Dstack-TEE/dstack

Support PROXY protocol for source IP preservation in dstack-gateway

Open
#353 2 comments 2 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.