Dstack-TEE / Dstack-TEE/dstack

Support PROXY protocol for source IP preservation in dstack-gateway

Đang mở
#353 2 bình luận 2 reaction 0 người được giao Xem trên GitHub
Ngôn ngữ chính
Rust
Star
544
Fork
96
Merge trung bình
17 giờ 57 phút
Pull request đã merge (30 ngày)
117

Mô tả

## 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

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.