20. 从零用Rust编写正反向代理，四层反向代理stream(tcp与udp)实现

wmproxy

wmproxy是由Rust编写，已实现http/https代理，socks5代理， 反向代理，静态文件服务器，内网穿透，配置热更新等， 后续将实现websocket代理等，同时会将实现过程分享出来， 感兴趣的可以一起造个轮子法

项目地址

gite: https://gitee.com/tickbh/wmproxy

github: https://github.com/tickbh/wmproxy

四层代理

四层代理，也称为网络层代理，是基于IP地址和端口号的代理方式。它只关心数据包的源IP地址、目的IP地址、源端口号和目的端口号，不关心数据包的具体内容。四层代理主要通过报文中的目标地址和端口，再加上负载均衡设备设置的服务器选择方式，决定最终选择的内部服务器。

因为四层代理不用处理任何相关的包信息，只需将包数据传递给正确的服务器即可，所以实现相对比较简单。

以下是OSI七层模型的示意图，来源于网上

实现方式

双端建立连接，也就是收到客户端的连接的时候，同时建立一条通往服务端的连接，然后做双向绑定即可完成服务。

四层代理还有udp的转发需求，需要同步将udp的数据进行转发，udp的处理方式处理会相对复杂一些，因为当前地址只有绑定一份，但是可能来自各种不同的地址，不同的客户端的(remote_ip, remote_port)我们需要当成一个全新的客户端。

而且有时候无法主动感知是否已经被断开了，所以也必须有超时机制，好在超时的时候能及时释放掉连接，好让系统及时的socket资源。

TCP实现

tcp找到相应的地址，连接，并双向绑定即可

pub async fn process<T>(
    data: Arc<Mutex<StreamConfig>>,
    local_addr: SocketAddr,
    mut inbound: T,
    _addr: SocketAddr,
) -> ProxyResult<()>
where
    T: AsyncRead + AsyncWrite + Unpin + std::marker::Send + 'static,
{
    let value = data.lock().await;
    for (_, s) in value.server.iter().enumerate() {
        if s.bind_addr.port() == local_addr.port() {
            let addr = ReverseHelper::get_upstream_addr(&s.upstream, "")?;
            let mut connect = HealthCheck::connect(&addr).await?;
            copy_bidirectional(&mut inbound, &mut connect).await?;
            break;
        }
    }
    Ok(())
}

UDP实现

UDP相对比较复杂，下面我们先列举内部的流程图
flowchart TD A[绑定反向udp端口] B[客户端] H{是否第一次} I[创建异步协程] D[异步协程中] B <-->|根据地址连接发送数据到| A A --> H H -->|是|I I -->|将Receiver传到以接收数据| D H -->|否,将数据Sender给|D D -->|异步读取数据并发送|A

在stream绑定的时候，要区分出TCP还是UDP的，做分别的绑定

/// stream的绑定，按bind_mode区分出udp或者是tcp，返回相应的列表
pub async fn bind(&mut self) -> ProxyResult<(Vec<TcpListener>, Vec<StreamUdp>)> {
    let mut listeners = vec![];
    let mut udp_listeners = vec![];
    let mut bind_port = HashSet::new();
    for value in &self.server.clone() {
        if bind_port.contains(&value.bind_addr.port()) {
            continue;
        }
        bind_port.insert(value.bind_addr.port());
        if value.bind_mode == "udp" {
            let listener = Helper::bind_upd(value.bind_addr).await?;
            udp_listeners.push(StreamUdp::new(listener, value.clone()));
        } else {
            let listener = Helper::bind(value.bind_addr).await?;
            listeners.push(listener);
        }
    }
    Ok((listeners, udp_listeners))
}

我们会对连接做分别的监听，下面是udp的获取是否有新数据：

async fn multi_udp_listen_work(
    listens: &mut Vec<StreamUdp>,
) -> (io::Result<(Vec<u8>, SocketAddr)>, usize) {
    if !listens.is_empty() {
        let (data, index, _) =
            select_all(listens.iter_mut().map(|listener| {
                listener.next().boxed()
            })).await;
        if data.is_none() {
            return (Err(io::Error::new(io::ErrorKind::InvalidInput, "read none data")), index)
        }
        (data.unwrap(), index)
    } else {
        let pend = std::future::pending();
        let () = pend.await;
        unreachable!()
    }
}

此处我们用next，也就是我们实现了 futures_core::Stream接口，用Poll的方式来注册实现有事件的时候来通知。

在tokio中，在read或者write的时候返回Poll::Pending，将会将socket的可读可写注册到底层，如果一旦系统可读可写就会通知该接口，将会重新执行一遍futures_core::Stream

我们将同时可以处理可读可写可发送事件，如果接口超时我们将关闭相应的接口。

impl Stream for StreamUdp {
    type Item = io::Result<(Vec<u8>, SocketAddr)>;
    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        let _ = self.poll_write(cx)?;
        let _ = self.poll_sender(cx)?;
        self.poll_read(cx)
    }
}

下面是主要的StreamUdp类

/// Udp转发的处理结构，缓存一些数值以做中转
pub struct StreamUdp {
    /// 读的缓冲类，避免每次都释放
    pub buf: BinaryMut,
    /// 核心的udp绑定端口
    pub socket: UdpSocket,
    pub server: ServerConfig,

    /// 如果接收该数据大小为0，那么则代表通知数据关闭
    pub receiver: Receiver<(Vec<u8>, SocketAddr)>,
    /// 将发送器传达给每个子协程
    pub sender: Sender<(Vec<u8>, SocketAddr)>,

    /// 接收的缓存数据，无法保证全部直接进行发送完毕
    pub cache_data: LinkedList<(Vec<u8>, SocketAddr)>,
    /// 发送的缓存数据，无法保证全部直接进行发送完毕
    pub send_cache_data: LinkedList<(Vec<u8>, SocketAddr)>,
    /// 每个地址绑定的对象，包含Sender，最后操作时间，超时时间
    remote_sockets: HashMap<SocketAddr, InnerUdp>,
}

结果测试

我们自己开一个udp服务端，绑定了本地的8089，我们将接收到的数据前面加上from server:并进行返回，代理端我们绑定了84的端口，并将udp数据转发给8089端：

use tokio::net::UdpSocket;
use std::io;

#[tokio::main]
async fn main() -> io::Result<()> {
    let sock = UdpSocket::bind("0.0.0.0:8089").await?;
    let mut buf = [0; 1024];
    loop {
        let (len, addr) = sock.recv_from(&mut buf).await?;
        let mut vec = "from server: ".as_bytes().to_vec();
        vec.extend(&buf[..len]);
        let _ = sock.send_to(&vec, addr).await?;
    }
}

客户端我们用nc运行：

可以看出两个客户端互相独立，彼此返回的数据均符合预期，正常的接收及返回。

TCP我们绑定了83端口并转发到HTTP的本地端口8080，我们用curl进行测试，符合预期，如图：

结语

至此四层的反向代理TCP/UDP均已完成，也符合预期。

点击 [关注] ，[在看] ，[点赞] 是对作者最大的支持