飞蛾头条

6. 用Rust手把手编写一个wmproxy(代理，内网穿透等), 通讯协议源码解读篇

2023-09-30 08:12 由 问蒙服务框架 发表于 #后端开发

用Rust手把手编写一个wmproxy(代理，内网穿透等), 通讯协议源码解读篇

项目 ++wmproxy++

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

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

事件模型的选取

  • OS线程, 简单的一个IO对应一个系统级别的线程，通常单进程创建的线程数是有限的，在线程与线程间同步数据会相当困难，线程间的调度争用会相当损耗效率，不适合IO密集的场景。
  • 事件驱动(Event driven), 事件驱动基本上是最早的高并发的IO密集型的编程模式了，如C++的libevent，RUST的MIO，通过监听IO的可读可写从而进行编程设计，缺点通常跟回调( Callback )一起使用，如果使用不好，回调层级过多会有回调地狱的风险。
  • 协程(Coroutines) 可能是目前比较火的并发模型，火遍全球的Go语言的协程设计就非常优秀。协程跟线程类似，无需改变编程模型，同时它也跟async类似，可以支持大量的任务并发运行。
  • actor模型 是erlang的杀手锏之一，它将所有并发计算分割成一个一个单元，这些单元被称为actor,单元之间通过消息传递的方式进行通信和数据传递，跟分布式系统的设计理念非常相像。由于actor模型跟现实很贴近，因此它相对来说更容易实现，但是一旦遇到流控制、失败重试等场景时，就会变得不太好用
  • async/await， 该模型为异步编辑模型，async模型的问题就是内部实现机制过于复杂，对于用户来说，理解和使用起来也没有线程和协程简单。主要是等待完成状态await，就比如读socket数据，等待系统将数据送达再继续触发读操作的执行，从而答到无损耗的运行。

这里我们选择的是async/await的模式

Rust中的async

  • Future 在 Rust 中是惰性的，只有在被轮询(poll)时才会运行， 因此丢弃一个 future 会阻止它未来再被运行, 你可以将Future理解为一个在未来某个时间点被调度执行的任务。在Rust中调用异步函数没有用await会被编辑器警告，因为这不符合预期。
  • Async 在 Rust 中使用开销是零， 意味着只有你能看到的代码(自己的代码)才有性能损耗，你看不到的(async 内部实现)都没有性能损耗，例如，你可以无需分配任何堆内存、也无需任何动态分发来使用 async，这对于热点路径的性能有非常大的好处，正是得益于此，Rust 的异步编程性能才会这么高。
  • Rust 异步运行时，Rust社区生态中已经提供了非常优异的运行时实现例如tokio，官方版本的async目前的生态相对tokio会差许多
  • 运行时同时支持单线程和多线程

流代码的封装

跟数据通讯相关的代码均放在streams目录下面。

  1. center_client.rs中的CenterClient表示中心客户端，提供主动连接服务端的能力并可选择为加密（TLS）或者普通模式，并且将该客户端收发的消息转给服务端
  2. center_server.rs中的CenterServer表示中心服务端，接受中心客户端的连接，并且将信息处理或者转发
  3. trans_stream.rs中的TransStream表示转发流量端，提供与中心端绑定的读出写入功能，在代理服务器中客户端接收的连接因为无需处理任何数据，直接绑定为TransStream将数据完整的转发给服务端
  4. virtual_stream.rs中的VirtualStream表示虚拟端，虚拟出一个流连接，并实现AsyncRead及AsyncRead，可以和流一样正常操作，在代理服务器中服务端接收到新连接，把他虚拟成一个VirtualStream，就可以直接和他连接的服务器上做双向绑定。

几种流式在代码中的转化

HTTP代理

下面展示的是http代理，通过加密TLS中的转化

flowchart TD A[TcpStream请求到代理]<-->|建立连接/明文|B[代理转化成TransStream] B<-->|转发到/内部|C[中心客户端] C<-->|建立加密连接/加密|D[TlsStream< TcpStream>绑定中心服务端] D<-->|收到Create/内部|E[虚拟出VirtualStream] E<-->|解析到host并连接/明文|F[TcpStream连接到http服务器]

上述过程实现了程序中实现了http的代理转发

HTTP内网穿透

以下是http内网穿透在代理中的转化

flowchart TD A[服务端绑定http对外端口]<-->|接收连接/明文|B[外部的TcpStream] B<-->|转发到/内部|C[中心服务端并绑定TransStream] C<-->|通过客户的加密连接推送/加密|D[TlsStream< TcpStream>绑定中心客户端] D<-->|收到Create/内部|E[虚拟出VirtualStream] E<-->|解析对应的连接信息/明文|F[TcpStream连接到内网的http服务器]

上述过程可以主动把公网的请求连接转发到内网，由内网提供完服务后再转发到公网的请求，从而实现内网穿透。

流代码的介绍

CenterClient中心客端

下面是代码类的定义

/// 中心客户端
/// 负责与服务端建立连接，断开后自动再重连
pub struct CenterClient {
    /// tls的客户端连接信息
    tls_client: Option<Arc<rustls::ClientConfig>>,
    /// tls的客户端连接域名
    domain: Option<String>,
    /// 连接中心服务器的地址
    server_addr: SocketAddr,
    /// 内网映射的相关消息
    mappings: Vec<MappingConfig>,
    /// 存在普通连接和加密连接，此处不为None则表示普通连接
    stream: Option<TcpStream>,
    /// 存在普通连接和加密连接，此处不为None则表示加密连接
    tls_stream: Option<TlsStream<TcpStream>>,
    /// 绑定的下一个sock_map映射
    next_id: u32,

    /// 发送Create，并将绑定的Sender发到做绑定
    sender_work: Sender<(ProtCreate, Sender<ProtFrame>)>,
    /// 接收的Sender绑定，开始服务时这值move到工作协程中，所以不能二次调用服务
    receiver_work: Option<Receiver<(ProtCreate, Sender<ProtFrame>)>>,

    /// 发送协议数据，接收到服务端的流数据，转发给相应的Stream
    sender: Sender<ProtFrame>,
    /// 接收协议数据，并转发到服务端。
    receiver: Option<Receiver<ProtFrame>>,
}

主要的逻辑流程，循环监听数据流的到达，同时等待多个异步的到达，这里用的是tokio::select!宏

loop {
    let _ = tokio::select! {
        // 严格的顺序流
        biased;
        // 新的流建立，这里接收Create并进行绑定
        r = receiver_work.recv() => {
            if let Some((create, sender)) = r {
                map.insert(create.sock_map(), sender);
                let _ = create.encode(&mut write_buf);
            }
        }
        // 数据的接收，并将数据写入给远程端
        r = receiver.recv() => {
            if let Some(p) = r {
                let _ = p.encode(&mut write_buf);
            }
        }
        // 数据的等待读取，一旦流可读则触发，读到0则关闭主动关闭所有连接
        r = reader.read(&mut vec) => {
            match r {
                Ok(0)=>{
                    is_closed=true;
                    break;
                }
                Ok(n) => {
                    read_buf.put_slice(&vec[..n]);
                }
                Err(_err) => {
                    is_closed = true;
                    break;
                },
            }
        }
        // 一旦有写数据，则尝试写入数据，写入成功后扣除相应的数据
        r = writer.write(write_buf.chunk()), if write_buf.has_remaining() => {
            match r {
                Ok(n) => {
                    write_buf.advance(n);
                    if !write_buf.has_remaining() {
                        write_buf.clear();
                    }
                }
                Err(e) => {
                    println!("center_client errrrr = {:?}", e);
                },
            }
        }
    };

    loop {
        // 将读出来的数据全部解析成ProtFrame并进行相应的处理，如果是0则是自身消息，其它进行转发
        match Helper::decode_frame(&mut read_buf)? {
            Some(p) => {
                match p {
                    ProtFrame::Create(p) => {
                    }
                    ProtFrame::Close(_) | ProtFrame::Data(_) => {
                    },
                }
            }
            None => {
                break;
            }
        }
    }
}

CenterServer中心服务端

下面是代码类的定义

/// 中心服务端
/// 接受中心客户端的连接，并且将信息处理或者转发
pub struct CenterServer {
    /// 代理的详情信息，如用户密码这类
    option: ProxyOption,
    
    /// 发送协议数据，接收到服务端的流数据，转发给相应的Stream
    sender: Sender<ProtFrame>,
    /// 接收协议数据，并转发到服务端。
    receiver: Option<Receiver<ProtFrame>>,

    /// 发送Create，并将绑定的Sender发到做绑定
    sender_work: Sender<(ProtCreate, Sender<ProtFrame>)>,
    /// 接收的Sender绑定，开始服务时这值move到工作协程中，所以不能二次调用服务
    receiver_work: Option<Receiver<(ProtCreate, Sender<ProtFrame>)>>,
    /// 绑定的下一个sock_map映射，为双数
    next_id: u32,
}

主要的逻辑流程，循环监听数据流的到达，同时等待多个异步的到达，这里用的是tokio::select!宏，select处理方法与Client相同，均处理相同逻辑，不同的是接收数据包后数据端是处理的proxy的请求，而Client处理的是内网穿透的逻辑

loop {
    // 将读出来的数据全部解析成ProtFrame并进行相应的处理，如果是0则是自身消息，其它进行转发
    match Helper::decode_frame(&mut read_buf)? {
        Some(p) => {
            match p {
                ProtFrame::Create(p) => {
                    tokio::spawn(async move {
                        let _ = Proxy::deal_proxy(stream, flag, username, password, udp_bind).await;
                    });
                }
                ProtFrame::Close(_) | ProtFrame::Data(_) => {
                },
            }
        }
        None => {
            break;
        }
    }
}

TransStream转发流量端

下面是代码类的定义

/// 转发流量端
/// 提供与中心端绑定的读出写入功能
pub struct TransStream<T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    // 流有相应的AsyncRead + AsyncWrite + Unpin均可
    stream: T,
    // sock绑定的句柄
    id: u32,
    // 读取的数据缓存，将转发成ProtFrame
    read: BinaryMut,
    // 写的数据缓存，直接写入到stream下，从ProtFrame转化而来
    write: BinaryMut,
    // 收到数据通过sender发送给中心端
    in_sender: Sender<ProtFrame>,
    // 收到中心端的写入请求，转成write
    out_receiver: Receiver<ProtFrame>,
}

主要的逻辑流程，循环监听数据流的到达，同时等待多个异步的到达，这里用的是tokio::select!宏，监听的对象有stream可读，可写，sender的写发送及receiver的可接收

loop {
    // 有剩余数据，优先转化成Prot，因为数据可能从外部直接带入
    if self.read.has_remaining() {
        link.push_back(ProtFrame::new_data(self.id, self.read.copy_to_binary()));
        self.read.clear();
    }

    tokio::select! {
        n = reader.read(&mut buf) => {
            let n = n?;
            if n == 0 {
                return Ok(())
            } else {
                self.read.put_slice(&buf[..n]);
            }
        },
        r = writer.write(self.write.chunk()), if self.write.has_remaining() => {
            match r {
                Ok(n) => {
                    self.write.advance(n);
                    if !self.write.has_remaining() {
                        self.write.clear();
                    }
                }
                Err(_) => todo!(),
            }
        }
        r = self.out_receiver.recv() => {
            if let Some(v) = r {
                if v.is_close() || v.is_create() {
                    return Ok(())
                } else if v.is_data() {
                    match v {
                        ProtFrame::Data(d) => {
                            self.write.put_slice(&d.data().chunk());
                        }
                        _ => unreachable!(),
                    }
                }
            } else {
                return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid frame"))
            }
        }
        p = self.in_sender.reserve(), if link.len() > 0 => {
            match p {
                Err(_)=>{
                    return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid frame"))
                }
                Ok(p) => {
                    p.send(link.pop_front().unwrap())
                }, 
            }
        }
    }

VirtualStream虚拟端

下面是代码类的定义，我们并未有真实的socket，通过虚拟出的端方便后续的操作

/// 虚拟端
/// 虚拟出一个流连接，并实现AsyncRead及AsyncRead，可以和流一样正常操作
pub struct VirtualStream
{
    // sock绑定的句柄
    id: u32,
    // 收到数据通过sender发送给中心端
    sender: PollSender<ProtFrame>,
    // 收到中心端的写入请求，转成write
    receiver: Receiver<ProtFrame>,
    // 读取的数据缓存，将转发成ProtFrame
    read: BinaryMut,
    // 写的数据缓存，直接写入到stream下，从ProtFrame转化而来
    write: BinaryMut,
}

虚拟的流主要通过实现AsyncRead及AsyncWrite


impl AsyncRead for VirtualStream
{
    // 有读取出数据，则返回数据，返回数据0的Ready状态则表示已关闭
    fn poll_read(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut [std](https://note.youdao.com/)[link](https://note.youdao.com/)::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        loop {
            match self.receiver.poll_recv(cx) {
                Poll::Ready(value) => {
                    if let Some(v) = value {
                        if v.is_close() || v.is_create() {
                            return Poll::Ready(Ok(()))
                        } else if v.is_data() {
                            match v {
                                ProtFrame::Data(d) => {
                                    self.read.put_slice(&d.data().chunk());
                                }
                                _ => unreachable!(),
                            }
                        }
                    } else {
                        return Poll::Ready(Ok(()))
                    }
                },
                Poll::Pending => {
                    if !self.read.has_remaining() {
                        return Poll::Pending;
                    }
                },
            }


            if self.read.has_remaining() {
                let copy = std::cmp::min(self.read.remaining(), buf.remaining());
                buf.put_slice(&self.read.chunk()[..copy]);
                self.read.advance(copy);
                return Poll::Ready(Ok(()));
            }
        }
        
    }
}


impl AsyncWrite for VirtualStream
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<Result<usize, std::io::Error>> {
        self.write.put_slice(buf);
        if let Err(_) = ready!(self.sender.poll_reserve(cx)) {
            return Poll::Pending;
        }
        let binary = Binary::from(self.write.chunk().to_vec());
        let id = self.id;
        if let Ok(_) = self.sender.send_item(ProtFrame::Data(ProtData::new(id, binary))) {
            self.write.clear();
        }
        Poll::Ready(Ok(buf.len()))
    }

}

至此基本几个大类已设置完毕，接下来仅需简单的拓展就能实现内网穿透功能。

热门相关：峡谷正能量   名门盛婚：首席，别来无恙！   重生之嫡女祸妃   重生之嫡女祸妃   魔神狂后