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/// 信号通道消息类型。
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum SignalKind {
87    /// 有新外部连接到达
88    NewConnection,
89    /// 信号通道已关闭（流关闭或读错误）
90    Disconnected,
91}
92
93/// 从信号通道读取一个信号（1 字节）。
94///
95/// 返回 [`SignalKind`]，不会传播 I/O 错误——任何读失败都映射为 `Disconnected`。
96pub async fn read_signal<R: AsyncRead + Unpin>(reader: &mut R) -> SignalKind {
97    let mut buf = [0u8; 1];
98    match reader.read_exact(&mut buf).await {
99        Ok(_) if buf[0] == crate::SIGNAL_NEW_CONN => SignalKind::NewConnection,
100        Ok(_) => SignalKind::NewConnection,  // treat any non-EOF byte as notification
101        Err(_) => SignalKind::Disconnected,
102    }
103}
104
105/// 向信号通道发送"新外部连接"通知（1 字节 `SIGNAL_NEW_CONN`）。
106pub async fn send_notification<W: AsyncWrite + Unpin>(writer: &mut W) -> std::io::Result<()> {
107    writer.write_all(&[crate::SIGNAL_NEW_CONN]).await
108}
109
110/// 双向 pipe 两个 TCP stream，任一方断开即结束。
111///
112/// 内部使用 `tokio::io::copy` 和 `tokio::select!` 实现双向转发。
113pub async fn pipe_bidirectional(
114    mut a: tokio::net::TcpStream,
115    mut b: tokio::net::TcpStream,
116) {
117    let (mut ar, mut aw) = a.split();
118    let (mut br, mut bw) = b.split();
119    tokio::select! {
120        _ = tokio::io::copy(&mut ar, &mut bw) => {}
121        _ = tokio::io::copy(&mut br, &mut aw) => {}
122    }
123}
gout_api/data_channel.rs

gout_api/
data_channel.rs
