gout_api/
data_channel.rs

1//! 数据通道协议 — 握手、确认、双向 pipe。
2//!
3//! 提供客户端和服务端两端的数据通道握手函数，
4//! 以及双向 TCP 数据转发（pipe）功能。
5
6use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
7
8use crate::{decode_handshake, encode_handshake, TunnelType, STATUS_OK};
9
10/// 数据通道握手错误。
11#[derive(Debug)]
12pub enum HandshakeError {
13    /// I/O 错误（连接断开、超时等）
14    Io(std::io::Error),
15    /// 服务端拒绝了握手（token 无效或隧道已关闭）
16    Rejected,
17    /// token 不存在或隧道已关闭
18    NotFound,
19}
20
21impl std::fmt::Display for HandshakeError {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        match self {
24            Self::Io(e) => write!(f, "I/O error: {e}"),
25            Self::Rejected => write!(f, "server rejected handshake"),
26            Self::NotFound => write!(f, "unknown token or tunnel closed"),
27        }
28    }
29}
30
31impl std::error::Error for HandshakeError {}
32
33impl From<std::io::Error> for HandshakeError {
34    fn from(e: std::io::Error) -> Self { Self::Io(e) }
35}
36
37/// 客户端发起握手：发送 `[token: u64 BE][tunnel_type: u8]` 并等待服务端确认。
38///
39/// 成功返回 `Ok(())`，失败返回 [`HandshakeError`]。
40pub async fn client_handshake(
41    stream: &mut (impl AsyncWrite + AsyncRead + Unpin),
42    token: u64,
43    tunnel_type: TunnelType,
44) -> Result<(), HandshakeError> {
45    let buf = encode_handshake(token, tunnel_type);
46    stream.write_all(&buf).await?;
47    let mut status = [0u8; 1];
48    stream.read_exact(&mut status).await?;
49    if status[0] == STATUS_OK {
50        Ok(())
51    } else {
52        Err(HandshakeError::Rejected)
53    }
54}
55
56/// 服务端接收并解析客户端握手，返回 `(token, tunnel_type)`。
57pub async fn server_receive_handshake(
58    stream: &mut (impl AsyncRead + Unpin),
59) -> Result<(u64, TunnelType), HandshakeError> {
60    let mut buf = [0u8; crate::HANDSHAKE_SIZE];
61    stream.read_exact(&mut buf).await?;
62    Ok(decode_handshake(&buf))
63}
64
65/// 服务端发送握手成功响应（1 字节 `STATUS_OK`）。
66pub async fn server_accept(
67    stream: &mut (impl AsyncWrite + Unpin),
68) -> Result<(), std::io::Error> {
69    stream.write_all(&[STATUS_OK]).await
70}
71
72/// 服务端发送握手拒绝响应：`[STATUS_ERR][err_len: u16 BE][reason]`。
73pub async fn server_reject(
74    stream: &mut (impl AsyncWrite + Unpin),
75    reason: &str,
76) -> Result<(), std::io::Error> {
77    let mut resp = vec![crate::STATUS_ERR];
78    let msg_bytes = reason.as_bytes();
79    resp.extend_from_slice(&(msg_bytes.len() as u16).to_be_bytes());
80    resp.extend_from_slice(msg_bytes);
81    stream.write_all(&resp).await
82}
83
84/// 双向 pipe 两个 TCP stream，任一方断开即结束。
85///
86/// 内部使用 `tokio::io::copy` 和 `tokio::select!` 实现双向转发。
87pub async fn pipe_bidirectional(
88    mut a: tokio::net::TcpStream,
89    mut b: tokio::net::TcpStream,
90) {
91    let (mut ar, mut aw) = a.split();
92    let (mut br, mut bw) = b.split();
93    tokio::select! {
94        _ = tokio::io::copy(&mut ar, &mut bw) => {}
95        _ = tokio::io::copy(&mut br, &mut aw) => {}
96    }
97}
gout_api/data_channel.rs

gout_api/
data_channel.rs
