6. 用Rust手把手編寫一個wmproxy(代理，內網穿透等), 通訊協議原始碼解讀篇

問蒙服務框架發表於2023-09-30
Rust內網穿透協議原始碼

用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()))
    }

}

至此基本幾個大類已設定完畢，接下來僅需簡單的擴充就能實現內網穿透功能。

相關文章