# ClientConnector Refactor Part 2: UDP Improvements

## Background

This document builds on `CLIENT_CONNECTOR_REFACTOR.md` which established the two-phase connection model (`connect_as_first_hop` vs `setup_on_existing_stream`). This part addresses issues discovered with UDP support in proxy chains.

## Problems to Solve

### Problem 1: UDP Support Check with Pooled Final Hops

Currently `supports_udp_over_tcp()` checks only the first connector in the last hop's pool:

```rust
pub fn supports_udp_over_tcp(&self) -> bool {
    self.hops
        .last()
        .map(|last_pool| last_pool.first().map(|c| c.supports_udp_over_tcp()).unwrap_or(false))
        .unwrap_or(false)
}
```

This is broken for pools where some connectors support UDP and others don't. For example:
```yaml
client_chain:
  - pool_with_mixed_udp_support  # Some support UDP, some don't
```

When round-robin selects a non-UDP connector and we try `connect_udp()`, it fails at runtime.

### Problem 2: Security Risk - Silent Fallback to Direct

In `tcp_server.rs`, when `supports_udp_over_tcp()` returns false, the code falls back to direct UDP:

```rust
if chain.supports_udp_over_tcp() {
    // Use chain's UDP-over-TCP
} else {
    // Direct connection - use native UDP socket  // <-- SECURITY RISK!
    let client_socket = chain.configure_udp_socket(remote_addr.is_ipv6())?;
    client_socket.connect(remote_addr).await?;
}
```

This silently bypasses the proxy chain, leaking traffic directly. If users configure a proxy chain, they expect ALL traffic to go through it.

### Problem 3: UdpStreamType Design Issues

Current design:
```rust
pub enum UdpStreamType {
    Bidirectional,
    SessionBased,
}

async fn connect_udp(
    &self,
    remote_location: NetLocation,  // <-- Always passed, even when not needed
    resolver: &Arc<dyn Resolver>,
    preferred_type: UdpStreamType,
) -> std::io::Result<TcpClientUdpSetupResult>;
```

Issues:
1. `remote_location` (target) is only used by `Bidirectional` - SessionBased gets destinations in XUDP frames
2. `MultiDirectional` type was removed but is still needed
3. For SessionBased, callers pass placeholder `0.0.0.0:0` which is ugly

### Problem 4: LocalConnector UDP Support

`LocalConnector` returns `supports_udp_over_tcp() = false`, but it CAN do native UDP:
- **Bidirectional**: Create connected UDP socket to target
- **SessionBased**: Use `UdpSessionMessageStream` for on-demand session sockets
- **MultiDirectional**: Use `UdpMessageStream` or `UdpMultiMessageStream`

This should be supported so direct chains work with UDP.

## Proposed Solutions

### Solution 1: Separate UDP Final Hop Pool (Using Indices)

Store indices into the final hop pool for UDP-capable connectors, avoiding the need to clone:

```rust
pub struct ClientProxyChain {
    /// All hops for TCP connections
    hops: Vec<Vec<Box<dyn ClientConnector>>>,
    next_indices: Vec<AtomicU32>,

    /// Indices into hops.last() for UDP-capable connectors (filtered at creation)
    /// This avoids needing Clone for Box<dyn ClientConnector>
    udp_final_hop_indices: Vec<usize>,
    udp_final_hop_next_index: AtomicU32,
}

impl ClientProxyChain {
    pub fn new(hops: Vec<Vec<Box<dyn ClientConnector>>>) -> Self {
        assert!(!hops.is_empty(), "ClientProxyChain must have at least one hop");

        // Collect indices of UDP-capable connectors in the final hop
        let udp_final_hop_indices: Vec<usize> = hops
            .last()
            .unwrap()
            .iter()
            .enumerate()
            .filter(|(_, c)| c.supports_udp())
            .map(|(i, _)| i)
            .collect();

        let next_indices = hops.iter().map(|_| AtomicU32::new(0)).collect();

        Self {
            hops,
            next_indices,
            udp_final_hop_indices,
            udp_final_hop_next_index: AtomicU32::new(0),
        }
    }
}
```

Update `select_connectors` to take UDP flag:

```rust
fn select_connectors(&self, is_udp: bool) -> Result<Vec<&dyn ClientConnector>, std::io::Error> {
    let num_hops = self.hops.len();
    let mut connectors = Vec::with_capacity(num_hops);

    // Select from each hop except the last
    for i in 0..num_hops - 1 {
        connectors.push(self.select_from_pool(i));
    }

    // For the last hop, use UDP pool if this is a UDP connection
    if is_udp {
        if self.udp_final_hop_indices.is_empty() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::Unsupported,
                "No connectors in final hop support UDP",
            ));
        }
        // Round-robin through UDP-capable connectors only
        let rr_idx = self.udp_final_hop_next_index.fetch_add(1, Ordering::Relaxed) as usize;
        let pool_idx = self.udp_final_hop_indices[rr_idx % self.udp_final_hop_indices.len()];
        connectors.push(self.hops.last().unwrap()[pool_idx].as_ref());
    } else {
        connectors.push(self.select_from_pool(num_hops - 1));
    }

    Ok(connectors)
}
```

### Solution 2: Remove supports_udp_over_tcp() Check

Remove the `supports_udp_over_tcp()` method from `ClientProxyChain` entirely. Instead:

1. Just call `connect_udp()`
2. Let it fail with a clear error if no UDP-capable connectors exist
3. Remove the fallback-to-direct code in `tcp_server.rs`

```rust
// BEFORE (in tcp_server.rs):
if chain.supports_udp_over_tcp() {
    chain.connect_udp(...).await?;
} else {
    // Direct fallback - SECURITY RISK
}

// AFTER:
chain.connect_udp(...).await?;  // Fails with clear error if no UDP support
```

The error message should be clear: "No connectors in chain support UDP-over-TCP tunneling"

### Solution 3: Rename UdpStreamType to UdpStreamRequest with Embedded Target

```rust
/// Request type for UDP connections through the chain.
/// The target is embedded in variants that need it, making the API type-safe.
#[derive(Debug, Clone)]
pub enum UdpStreamRequest {
    /// Single fixed destination - uses AsyncMessageStream.
    /// Target is required: all packets go to this one destination.
    Bidirectional { target: NetLocation },

    /// Multiple destinations with address per packet - uses AsyncTargetedMessageStream.
    /// No target needed: each packet specifies its own destination.
    MultiDirectional,

    /// Session-based multiplexing (XUDP) - uses AsyncSessionMessageStream.
    /// No target needed: destinations come in XUDP frames.
    SessionBased,
}
```

Update the result type:

```rust
pub enum TcpClientUdpSetupResult {
    /// Simple bidirectional UDP - single fixed destination.
    Bidirectional { stream: Box<dyn AsyncMessageStream> },

    /// Multi-target UDP with per-packet destinations.
    MultiDirectional { stream: Box<dyn AsyncSourcedMessageStream> },

    /// Session-based UDP multiplexing (XUDP).
    SessionBased { stream: Box<dyn AsyncSessionMessageStream> },
}
```

Update `connect_udp` signature:

```rust
// BEFORE:
pub async fn connect_udp(
    &self,
    remote_location: NetLocation,
    resolver: &Arc<dyn Resolver>,
    preferred_type: UdpStreamType,
) -> std::io::Result<TcpClientUdpSetupResult>;

// AFTER:
pub async fn connect_udp(
    &self,
    resolver: &Arc<dyn Resolver>,
    request: UdpStreamRequest,
) -> std::io::Result<TcpClientUdpSetupResult>;
```

### Solution 4: LocalConnector Full UDP Support

Update `LocalConnector` to support all three UDP modes:

```rust
impl LocalConnector {
    fn supports_udp_over_tcp(&self) -> bool {
        true  // We support native UDP, which is just as good
    }
}

#[async_trait]
impl ClientConnector for LocalConnector {
    async fn connect_udp_as_first_hop(
        &self,
        resolver: &Arc<dyn Resolver>,
        request: UdpStreamRequest,
    ) -> std::io::Result<TcpClientUdpSetupResult> {
        match request {
            UdpStreamRequest::Bidirectional { target } => {
                // Connected UDP socket to single target
                let remote_addr = resolve_single_address(resolver, &target).await?;
                let client_socket = new_udp_socket(remote_addr.is_ipv6(), self.bind_interface.clone())?;
                client_socket.connect(remote_addr).await?;

                Ok(TcpClientUdpSetupResult::Bidirectional {
                    stream: Box::new(client_socket),
                })
            }

            UdpStreamRequest::MultiDirectional => {
                // Unconnected socket that can send to any destination
                let client_socket = new_udp_socket(true, self.bind_interface.clone())?; // IPv6 for dual-stack

                Ok(TcpClientUdpSetupResult::MultiDirectional {
                    stream: Box::new(UdpMessageStream::new(client_socket, resolver.clone())),
                })
            }

            UdpStreamRequest::SessionBased => {
                // Session manager creates sockets on-demand per session
                let config = UdpSocketConfig {
                    bind_interface: self.bind_interface.clone(),
                };
                let session_manager = UdpSessionMessageStream::new(config);

                Ok(TcpClientUdpSetupResult::SessionBased {
                    stream: Box::new(session_manager),
                })
            }
        }
    }
}
```

## Updated ClientConnector Trait

```rust
#[async_trait]
pub trait ClientConnector: Send + Sync + Debug {
    /// Returns the bind interface for this connector, if any.
    fn bind_interface(&self) -> &Option<String>;

    /// Check if this connector supports UDP connections.
    ///
    /// For proxy connectors, this means UDP-over-TCP tunneling (VLESS, VMess, etc).
    /// For LocalConnector, this means native UDP socket support.
    fn supports_udp(&self) -> bool;

    /// Returns the proxy server address, if this is a proxy connector.
    /// Local/direct connectors return None.
    fn proxy_location(&self) -> Option<&NetLocation>;

    // ... TCP methods unchanged ...

    /// For hop 0 UDP: Create connection and setup UDP stream.
    async fn connect_udp_as_first_hop(
        &self,
        resolver: &Arc<dyn Resolver>,
        request: UdpStreamRequest,
    ) -> std::io::Result<TcpClientUdpSetupResult>;

    /// For hop 1+ UDP: Setup UDP on existing stream.
    async fn setup_udp_on_existing_stream(
        &self,
        stream: Box<dyn AsyncStream>,
        request: UdpStreamRequest,
    ) -> std::io::Result<TcpClientUdpSetupResult>;
}
```

## Updated tcp_server.rs

### BidirectionalUdp Case

```rust
TcpServerSetupResult::BidirectionalUdp {
    stream: mut server_stream,
    need_initial_flush: server_need_initial_flush,
    remote_location,
    override_proxy_provider,
} => {
    // ... get chain from provider ...

    let client_udp_result = chain
        .connect_udp(
            &resolver,
            UdpStreamRequest::Bidirectional { target: remote_location },
        )
        .await?;

    match client_udp_result {
        TcpClientUdpSetupResult::Bidirectional { stream: mut client_stream } => {
            copy_bidirectional_message(&mut server_stream, &mut client_stream, ...).await?;
            Ok(())
        }
        _ => unreachable!("connect_udp must return matching type"),
    }
}
```

### MultiDirectionalUdp Case

```rust
TcpServerSetupResult::MultiDirectionalUdp {
    stream: mut server_stream,
    need_initial_flush: server_need_initial_flush,
    override_proxy_provider,
    num_sockets: _,  // No longer needed - chain handles socket creation
} => {
    // ... get chain from provider ...

    let client_udp_result = chain
        .connect_udp(&resolver, UdpStreamRequest::MultiDirectional)
        .await?;

    match client_udp_result {
        TcpClientUdpSetupResult::MultiDirectional { stream: mut client_stream } => {
            copy_multidirectional_message(&mut server_stream, &mut client_stream, ...).await?;
            Ok(())
        }
        _ => unreachable!("connect_udp must return matching type"),
    }
}
```

### SessionBasedUdp Case

```rust
TcpServerSetupResult::SessionBasedUdp {
    stream: mut server_stream,
    need_initial_flush: server_need_initial_flush,
    override_proxy_provider,
} => {
    // ... get chain from provider ...

    let client_udp_result = chain
        .connect_udp(&resolver, UdpStreamRequest::SessionBased)
        .await?;

    match client_udp_result {
        TcpClientUdpSetupResult::SessionBased { stream: mut client_stream } => {
            copy_session_messages(&mut server_stream, &mut client_stream, ...).await?;
            Ok(())
        }
        _ => unreachable!("connect_udp must return matching type"),
    }
}
```

## Migration Summary

### Changes Required

1. **`UdpStreamType` → `UdpStreamRequest`**
   - Rename enum
   - Add `target` field to `Bidirectional` variant
   - Add `MultiDirectional` variant back
   - Update all call sites

2. **`TcpClientUdpSetupResult`**
   - Add `MultiDirectional` variant back
   - Keep `Bidirectional` and `SessionBased`

3. **`ClientProxyChain`**
   - Add `udp_final_hop` and `udp_final_hop_next_index` fields
   - Filter UDP-capable connectors at creation time
   - Update `select_connectors(is_udp: bool)`
   - Remove `supports_udp_over_tcp()` method
   - Update `connect_udp()` signature and implementation

4. **`ClientConnector` trait**
   - Rename `supports_udp_over_tcp()` → `supports_udp()`
   - Update UDP method signatures to use `UdpStreamRequest`

5. **`LocalConnector`**
   - Return `true` from `supports_udp()`
   - Implement all three UDP modes in `connect_udp_as_first_hop`

6. **`TcpClientConnector`**
   - Update to use new `UdpStreamRequest`
   - Update handler calls

7. **`tcp_server.rs`**
   - Remove all fallback-to-direct code
   - Update to use new `UdpStreamRequest` API
   - Simplify to just call `chain.connect_udp()` and handle result

8. **Protocol handlers (VLESS, VMess, Trojan)**
   - Update `setup_udp_stream` to accept `UdpStreamRequest`
   - Extract target from `Bidirectional` variant
   - Return `MultiDirectional` for protocols that support it

## Benefits

1. **Type Safety**: Target is only provided when needed (Bidirectional)
2. **No Silent Failures**: UDP errors are explicit, no fallback to direct
3. **Correct Pooling**: UDP connections use only UDP-capable connectors
4. **Full LocalConnector Support**: Direct chains work with all UDP modes
5. **Cleaner API**: No placeholder addresses for SessionBased
