# ClientConnector Refactor Part 4: Handler-Driven UDP Stream Pairing

**Status: IMPLEMENTED**

## Problem Statement

Currently, when the server receives a UDP connection request (e.g., `SessionBasedUdp`), it extracts the server stream and then calls `chain_group.connect_udp()` with a `UdpStreamRequest` enum. The handler must return a `TcpClientUdpSetupResult` that **exactly matches** the requested type, or the server code hits an `unreachable!()`.

This is inflexible for several reasons:

1. **Type Mismatch Handling**: If the server speaks `SessionBased` (XUDP) but the best available client is `MultiDirectional` (native UDP), we cannot connect them efficiently. The handler must either fail or the server must wrap/unwrap.

2. **Optimal Wrapping Unknown to Server**: The protocol handler knows which wrapping direction is optimal:
   - Bidirectional → MultiDirectional/SessionBased: easy (add per-packet/session metadata)
   - MultiDirectional/SessionBased → Bidirectional: hard/impossible (lose per-packet addressing)
   - Which side to wrap (server or client) depends on protocol-specific optimizations

3. **Server Code Complexity**: The server has three near-identical UDP handling blocks that:
   - Request a specific UDP type
   - Match the response
   - Hit unreachable if mismatched
   - Call the appropriate copy function

## Design Goals

1. **Handler-driven pairing**: Protocol handler receives server stream, decides optimal client setup, returns matched pair
2. **Simplified server code**: Server receives ready-to-copy stream pair, calls unified copy function
3. **Type safety**: Result type clearly indicates both streams have compatible types
4. **Future extensibility**: Support mismatched pairs with dedicated copy functions later
5. **Method renaming**: Clearer naming for TCP vs UDP setup methods

## Proposed Changes

### 1. Rename Trait Methods for Clarity

In `TcpClientHandler` trait (`src/tcp/tcp_handler.rs`):

```rust
// Current names
async fn setup_client_stream(...) -> std::io::Result<TcpClientSetupResult>;
async fn setup_udp_stream(...) -> std::io::Result<TcpClientUdpSetupResult>;

// New names - clearer distinction
async fn setup_client_tcp_stream(...) -> std::io::Result<TcpClientSetupResult>;
async fn setup_client_udp_stream(...) -> std::io::Result<TcpClientUdpSetupResult>;
```

This affects all handlers implementing `TcpClientHandler`:
- `vmess_handler.rs`
- `vless_client_handler.rs`
- `trojan_handler.rs`
- `tls_client_handler.rs`
- `socks_handler.rs`
- `http_handler.rs`
- `websocket_handler.rs`
- `snell_handler.rs`
- `shadowsocks_tcp_handler.rs`
- `reality_client_handler.rs`
- `shadow_tls_client_handler.rs`
- `port_forward_handler.rs`

### 2. New `ServerUdpStream` Enum

Pass the server stream into `connect_udp()` so the handler can decide wrapping:

```rust
/// Server-side UDP stream types that need to be matched with client streams.
pub enum ServerUdpStream {
    /// Single fixed destination
    Bidirectional {
        stream: Box<dyn AsyncMessageStream>,
        target: NetLocation,
    },
    /// Multiple destinations with address per packet
    MultiDirectional {
        stream: Box<dyn AsyncTargetedMessageStream>,
    },
    /// Session-based multiplexing (XUDP)
    SessionBased {
        stream: Box<dyn AsyncSessionMessageStream>,
    },
}

impl ServerUdpStream {
    /// Returns the target for Bidirectional streams
    pub fn target(&self) -> Option<&NetLocation> {
        match self {
            ServerUdpStream::Bidirectional { target, .. } => Some(target),
            _ => None,
        }
    }
}
```

### 3. Updated `UdpStreamRequest`

Now carries the server stream for handler inspection:

```rust
/// Request type for UDP connections through the chain.
/// Contains the server stream so the handler can decide optimal pairing.
pub enum UdpStreamRequest {
    /// Single fixed destination - server provides AsyncMessageStream
    Bidirectional {
        server_stream: Box<dyn AsyncMessageStream>,
        target: NetLocation,
    },

    /// Multiple destinations - server provides AsyncTargetedMessageStream
    MultiDirectional {
        server_stream: Box<dyn AsyncTargetedMessageStream>,
    },

    /// Session-based multiplexing - server provides AsyncSessionMessageStream
    SessionBased {
        server_stream: Box<dyn AsyncSessionMessageStream>,
    },
}
```

### 4. Updated `TcpClientUdpSetupResult`

Returns both server and client streams as a matched pair:

```rust
/// Result type for UDP client setup.
/// Contains both server and client streams, ready for copying.
/// The handler has already performed any necessary wrapping.
pub enum TcpClientUdpSetupResult {
    /// Both sides are Bidirectional (AsyncMessageStream)
    Bidirectional {
        server_stream: Box<dyn AsyncMessageStream>,
        client_stream: Box<dyn AsyncMessageStream>,
    },

    /// Both sides are MultiDirectional
    /// Server: AsyncTargetedMessageStream (read gives source, write takes target)
    /// Client: AsyncSourcedMessageStream (read gives source, write takes target)
    MultiDirectional {
        server_stream: Box<dyn AsyncTargetedMessageStream>,
        client_stream: Box<dyn AsyncSourcedMessageStream>,
    },

    /// Both sides are SessionBased (AsyncSessionMessageStream)
    SessionBased {
        server_stream: Box<dyn AsyncSessionMessageStream>,
        client_stream: Box<dyn AsyncSessionMessageStream>,
    },

    // Future: Mixed pairs with dedicated copy implementations
    // BidirectionalToMulti { ... },
    // MultiToSession { ... },
}
```

### 5. Changes to `ClientConnector` Trait

```rust
#[async_trait]
pub trait ClientConnector: Send + Sync + Debug {
    // ... existing methods ...

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

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

### 6. Changes to `TcpClientHandler` Trait

```rust
#[async_trait]
pub trait TcpClientHandler: Send + Sync + Debug {
    /// Setup a TCP client connection through this proxy.
    /// RENAMED from setup_client_stream
    async fn setup_client_tcp_stream(
        &self,
        client_stream: Box<dyn AsyncStream>,
        remote_location: NetLocation,
    ) -> std::io::Result<TcpClientSetupResult>;

    fn supports_udp_over_tcp(&self) -> bool { false }

    /// Setup a UDP stream over TCP.
    /// RENAMED from setup_udp_stream
    ///
    /// Handler receives the server stream, decides optimal client setup,
    /// and returns a matched pair ready for copying.
    async fn setup_client_udp_stream(
        &self,
        tcp_stream: Box<dyn AsyncStream>,
        request: UdpStreamRequest,
    ) -> std::io::Result<TcpClientUdpSetupResult> {
        Err(std::io::Error::new(
            std::io::ErrorKind::Unsupported,
            "UDP-over-TCP not supported by this protocol",
        ))
    }
}
```

### 7. Changes to `ClientProxyChain`

```rust
impl ClientProxyChain {
    /// Connect for UDP traffic through the chain.
    /// Takes server stream, returns matched pair ready for copying.
    pub async fn connect_udp(
        &self,
        resolver: &Arc<dyn Resolver>,
        request: UdpStreamRequest,
    ) -> std::io::Result<TcpClientUdpSetupResult> {
        // ... existing selection logic ...

        // Single hop case
        if connectors.len() == 1 {
            return connectors[0]
                .connect_udp_as_first_hop(resolver, request)
                .await;
        }

        // Multi-hop: build TCP tunnel, then setup UDP on last hop
        // ... existing TCP tunnel building ...

        let result = last_connector
            .setup_udp_on_existing_stream(tcp_result.client_stream, request)
            .await?;

        Ok(result)
    }
}
```

### 8. Handler Implementation Example: LocalConnector

```rust
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 { server_stream, target } => {
                // Create connected UDP socket for single target
                let remote_addr = resolve_single_address(resolver, &target).await?;
                let socket = new_udp_socket(remote_addr.is_ipv6(), self.bind_interface.clone())?;
                socket.connect(remote_addr).await?;

                Ok(TcpClientUdpSetupResult::Bidirectional {
                    server_stream,
                    client_stream: Box::new(socket),
                })
            }

            UdpStreamRequest::MultiDirectional { server_stream } => {
                // Unconnected socket for multi-target
                let socket = new_udp_socket(true, self.bind_interface.clone())?;

                Ok(TcpClientUdpSetupResult::MultiDirectional {
                    server_stream,
                    client_stream: Box::new(UdpMessageStream::new(socket, resolver.clone())),
                })
            }

            UdpStreamRequest::SessionBased { server_stream } => {
                // Session manager for XUDP-like operation
                let config = UdpSocketConfig {
                    bind_interface: self.bind_interface.clone(),
                };

                Ok(TcpClientUdpSetupResult::SessionBased {
                    server_stream,
                    client_stream: Box::new(UdpSessionMessageStream::new(config)),
                })
            }
        }
    }
}
```

### 9. Handler Implementation Example: VMess (XUDP)

```rust
impl TcpClientHandler for VmessTcpClientHandler {
    async fn setup_client_udp_stream(
        &self,
        tcp_stream: Box<dyn AsyncStream>,
        request: UdpStreamRequest,
    ) -> std::io::Result<TcpClientUdpSetupResult> {
        match request {
            UdpStreamRequest::SessionBased { server_stream } => {
                // VMess native XUDP - perfect match
                let client_stream = self.setup_vmess_xudp(tcp_stream).await?;

                Ok(TcpClientUdpSetupResult::SessionBased {
                    server_stream,
                    client_stream: Box::new(client_stream),
                })
            }

            UdpStreamRequest::Bidirectional { server_stream, target } => {
                // VMess supports single-target UDP (COMMAND_UDP)
                let client_stream = self.setup_vmess_udp(tcp_stream, target).await?;

                Ok(TcpClientUdpSetupResult::Bidirectional {
                    server_stream,
                    client_stream: Box::new(client_stream),
                })
            }

            UdpStreamRequest::MultiDirectional { .. } => {
                // VMess doesn't support MultiDirectional natively
                // Could potentially wrap server stream to SessionBased and use XUDP
                Err(std::io::Error::new(
                    std::io::ErrorKind::Unsupported,
                    "VMess does not support MultiDirectional UDP. Use SessionBased.",
                ))
            }
        }
    }
}
```

### 10. Simplified Server Code

The server code becomes much simpler:

```rust
// In tcp_server.rs / quic_server.rs
async fn process_stream(...) -> std::io::Result<()> {
    let setup_result = server_handler.setup_server_stream(stream).await?;

    match setup_result {
        TcpServerSetupResult::TcpForward { ... } => {
            // ... TCP handling unchanged ...
        }

        TcpServerSetupResult::BidirectionalUdp {
            remote_location,
            stream: server_stream,
            need_initial_flush: server_need_initial_flush,
            override_proxy_provider,
        } => {
            let selected_provider = select_provider(override_proxy_provider, client_proxy_selector);
            let action = selected_provider.judge(remote_location.clone(), &resolver).await?;

            match action {
                ConnectDecision::Allow { chain_group, remote_location } => {
                    // Pass server stream to handler, get matched pair back
                    let result = chain_group.connect_udp(
                        &resolver,
                        UdpStreamRequest::Bidirectional {
                            server_stream,
                            target: remote_location,
                        },
                    ).await?;

                    // No match needed - dispatch to appropriate copy function
                    copy_udp_streams(result, server_need_initial_flush).await?;
                }
                ConnectDecision::Block => {}
            }
        }

        TcpServerSetupResult::MultiDirectionalUdp { ... } => {
            // Similar pattern
        }

        TcpServerSetupResult::SessionBasedUdp { ... } => {
            // Similar pattern
        }
    }
}

/// Dispatch to appropriate copy function based on result type
async fn copy_udp_streams(
    result: TcpClientUdpSetupResult,
    server_need_initial_flush: bool,
) -> std::io::Result<()> {
    match result {
        TcpClientUdpSetupResult::Bidirectional { mut server_stream, mut client_stream } => {
            copy_bidirectional_message(
                &mut server_stream,
                &mut client_stream,
                server_need_initial_flush,
                false,
            ).await
        }
        TcpClientUdpSetupResult::MultiDirectional { mut server_stream, mut client_stream } => {
            copy_multidirectional_message(
                &mut server_stream,
                &mut client_stream,
                server_need_initial_flush,
                false,
            ).await
        }
        TcpClientUdpSetupResult::SessionBased { mut server_stream, mut client_stream } => {
            copy_session_messages(
                &mut server_stream,
                &mut client_stream,
                server_need_initial_flush,
                false,
            ).await
        }
    }
}
```

## Future Extensibility: Mixed Stream Pairs

In the future, we can support mismatched pairs:

```rust
pub enum TcpClientUdpSetupResult {
    // ... existing matched pairs ...

    /// Server is Bidirectional, client is MultiDirectional
    /// Requires wrapping client to present single-target interface
    BidirectionalToMulti {
        server_stream: Box<dyn AsyncMessageStream>,
        client_stream: Box<dyn AsyncSourcedMessageStream>,
        target: NetLocation, // For wrapping client writes
    },

    /// Server is SessionBased, client is MultiDirectional
    /// Session manager wraps multi-target to session-based
    SessionToMulti {
        server_stream: Box<dyn AsyncSessionMessageStream>,
        client_stream: Box<dyn AsyncSourcedMessageStream>,
    },
}
```

Each variant would need a dedicated copy function that handles the conversion.

## Wrapping Strategy Rationale

### Easy wrappings (handler can do these):
- **Bidirectional → MultiDirectional**: Add fixed target to all writes, ignore source on reads
- **Bidirectional → SessionBased**: Create single session, add session metadata

### Hard/impossible wrappings:
- **MultiDirectional → Bidirectional**: Can't reduce multiple targets to one
- **SessionBased → Bidirectional**: Can't collapse sessions to single stream

### Handler's advantage:
The handler knows its native mode and can:
1. Choose not to wrap if native mode matches request
2. Wrap the server side if that's more efficient
3. Wrap the client side if that's more efficient
4. Return an error if wrapping is impossible

## Implementation Order

1. **Phase 1: Rename methods** (non-breaking with deprecation)
   - Add new method names with `#[deprecated]` on old ones
   - Update all handler implementations
   - Remove deprecated methods

2. **Phase 2: Update types**
   - Add `ServerUdpStream` enum
   - Update `UdpStreamRequest` to carry server stream
   - Update `TcpClientUdpSetupResult` to return pair

3. **Phase 3: Update handlers**
   - Update `ClientConnector` implementations
   - Update `TcpClientHandler` implementations
   - Add `copy_udp_streams` helper

4. **Phase 4: Simplify servers**
   - Update `tcp_server.rs`
   - Update `quic_server.rs`
   - Update `hysteria2_server.rs`
   - Update `tuic_server.rs`

## Files to Modify

### Core types (Phase 1-2)
- `src/tcp/tcp_handler.rs` - `UdpStreamRequest`, `TcpClientUdpSetupResult`, `TcpClientHandler`
- `src/tcp/client_connector.rs` - `ClientConnector` trait

### Connectors (Phase 3)
- `src/tcp/local_connector.rs` - `LocalConnector`
- `src/tcp/tcp_client_connector.rs` - `TcpClientConnector`

### Handlers (Phase 1, 3)
- `src/vmess/vmess_handler.rs`
- `src/vless/vless_client_handler.rs`
- `src/trojan_handler.rs`
- `src/tls_client_handler.rs`
- `src/socks_handler.rs`
- `src/http_handler.rs`
- `src/websocket/websocket_handler.rs`
- `src/snell/snell_handler.rs`
- `src/shadowsocks/shadowsocks_tcp_handler.rs`
- `src/reality_client_handler.rs`
- `src/shadow_tls/shadow_tls_client_handler.rs`
- `src/port_forward_handler.rs`

### Chain logic (Phase 3)
- `src/client_proxy_chain.rs` - `ClientProxyChain`, `ClientChainGroup`

### Servers (Phase 4)
- `src/tcp/tcp_server.rs`
- `src/quic_server.rs`
- `src/hysteria2_server.rs`
- `src/tuic_server.rs`

## Summary

This refactoring:
1. **Gives handlers control** over optimal stream pairing
2. **Simplifies server code** by removing match-and-dispatch boilerplate
3. **Enables future flexibility** for mismatched stream pairs
4. **Improves naming** with clearer TCP vs UDP method names
5. **Maintains type safety** with explicit result variants

The key insight is that the handler (not the server) should decide how to pair server and client streams, because the handler knows:
- Its native UDP mode
- Which wrapping directions are possible
- Which side is more efficient to wrap
