温锦文欧普厨电及净水器总代理 发表于 2023-11-24 21:52:27

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七层模型的示意图,来源于网上
https://img2023.cnblogs.com/blog/2631821/202310/2631821-20231030071550753-891997292.png
实现方式

双端建立连接,也就是收到客户端的连接的时候,同时建立一条通往服务端的连接,然后做双向绑定即可完成服务。
四层代理还有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;    fn poll_next(      mut self: std::pin::Pin,      cx: &mut std::task::Context
页: [1]
查看完整版本: 20. 从零用Rust编写正反向代理,四层反向代理stream(tcp与udp)实现