fkm_proxy/utils/
mod.rs

1use anyhow::Result;
2use quinn::VarInt;
3use std::{
4    net::{SocketAddr, ToSocketAddrs},
5    pin::Pin,
6    task::Context,
7    time::Duration,
8};
9use thiserror::Error;
10use tokio::{
11    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
12    net::TcpStream,
13};
14
15pub mod certs;
16pub mod client;
17pub mod http;
18pub mod serve;
19
20#[derive(Debug)]
21pub struct HelloPacket {
22    pub hp_type: HelloPacketType,
23    pub token: u128,
24    pub own_ssl: bool,
25    pub redirect_ssl: bool,
26    pub tunnel_id: u128,
27}
28
29#[derive(Debug, PartialEq)]
30pub enum HelloPacketType {
31    Connector = 0,
32    Tunnel = 1,
33
34    Invalid,
35}
36
37impl HelloPacketType {
38    pub fn to_u8(&self) -> u8 {
39        match self {
40            HelloPacketType::Connector => 0,
41            HelloPacketType::Tunnel => 1,
42            HelloPacketType::Invalid => u8::MAX,
43        }
44    }
45
46    pub fn from_u8(val: u8) -> Self {
47        match val {
48            0 => HelloPacketType::Connector,
49            1 => HelloPacketType::Tunnel,
50            _ => HelloPacketType::Invalid,
51        }
52    }
53}
54
55impl HelloPacket {
56    pub const fn buf_size() -> usize {
57        40
58    }
59
60    pub fn to_buf(&self) -> [u8; Self::buf_size()] {
61        let mut tmp = [0; Self::buf_size()];
62        tmp[0] = self.hp_type.to_u8();
63        tmp[1..17].copy_from_slice(&self.token.to_be_bytes());
64        tmp[17] = self.own_ssl as u8;
65        tmp[18] = self.redirect_ssl as u8;
66        tmp[19..35].copy_from_slice(&self.tunnel_id.to_be_bytes());
67
68        tmp
69    }
70
71    pub fn from_buf(buf: &[u8; Self::buf_size()]) -> Self {
72        Self {
73            hp_type: HelloPacketType::from_u8(buf[0]),
74            token: u128::from_be_bytes(buf[1..17].try_into().unwrap()),
75            own_ssl: buf[17] != 0,
76            redirect_ssl: buf[18] != 0,
77            tunnel_id: u128::from_be_bytes(buf[19..35].try_into().unwrap()),
78        }
79    }
80}
81
82#[derive(Debug, PartialEq)]
83pub enum ConnectorPacketType {
84    Ping = 0,
85    TunnelRequest = 1,
86    Close = 2,
87    ConnectorConnected = 3,
88
89    Invalid,
90}
91
92impl ConnectorPacketType {
93    pub fn to_u8(&self) -> u8 {
94        match self {
95            ConnectorPacketType::Ping => 0,
96            ConnectorPacketType::TunnelRequest => 1,
97            ConnectorPacketType::Close => 2,
98            ConnectorPacketType::ConnectorConnected => 3,
99            ConnectorPacketType::Invalid => u8::MAX,
100        }
101    }
102
103    pub fn from_u8(val: u8) -> Self {
104        match val {
105            0 => ConnectorPacketType::Ping,
106            1 => ConnectorPacketType::TunnelRequest,
107            2 => ConnectorPacketType::Close,
108            3 => ConnectorPacketType::ConnectorConnected,
109            _ => ConnectorPacketType::Invalid,
110        }
111    }
112}
113
114#[derive(Debug)]
115pub struct ConnectorPacket {
116    pub packet_type: ConnectorPacketType,
117    pub tunnel_id: u128,
118    pub ssl: bool,
119}
120
121impl ConnectorPacket {
122    pub const fn buf_size() -> usize {
123        20
124    }
125
126    pub fn to_buf(&self) -> [u8; Self::buf_size()] {
127        let mut tmp = [0; Self::buf_size()];
128        tmp[0] = self.packet_type.to_u8();
129        tmp[1..17].copy_from_slice(&self.tunnel_id.to_be_bytes());
130        tmp[17] = self.ssl as u8;
131
132        tmp
133    }
134
135    pub fn from_buf(buf: &[u8; Self::buf_size()]) -> Self {
136        Self {
137            packet_type: ConnectorPacketType::from_u8(buf[0]),
138            tunnel_id: u128::from_be_bytes(buf[1..17].try_into().unwrap()),
139            ssl: buf[17] != 0,
140        }
141    }
142}
143
144pub fn parse_socketaddr(arg: &str) -> Result<SocketAddr> {
145    for i in 0..10 {
146        let res = arg.to_socket_addrs();
147
148        match res {
149            Ok(addrs) => {
150                for addr in addrs {
151                    if addr.is_ipv4() {
152                        return Ok(addr);
153                    }
154                }
155            }
156            Err(e) => {
157                tracing::warn!("[clap parse_socketaddr] (Try: {}) {e:?}", i + 1);
158                std::thread::sleep(Duration::from_millis(5000));
159            }
160        }
161    }
162
163    Err(anyhow::anyhow!("No ipv4 socketaddr found!"))
164}
165
166pub fn generate_string_packet(string: &str) -> Result<Vec<u8>> {
167    let mut bytes = string.as_bytes().to_vec();
168    bytes.push(0); // null terminator
169
170    Ok(bytes)
171}
172
173pub async fn send_string_to_stream<T>(stream: &mut T, string: &str) -> Result<()>
174where
175    T: AsyncWrite + Unpin,
176{
177    let bytes = generate_string_packet(string)?;
178    stream.write_all(&bytes).await?;
179
180    Ok(())
181}
182
183pub async fn read_string_from_stream<T>(stream: &mut T) -> Result<String>
184where
185    T: AsyncRead + Unpin,
186{
187    let mut buffer = Vec::new();
188    loop {
189        let byte = stream.read_u8().await?;
190        if byte == 0 {
191            break;
192        }
193
194        buffer.push(byte);
195    }
196
197    Ok(String::from_utf8(buffer)?)
198}
199
200#[derive(Error, Debug)]
201pub enum HelloPacketError {
202    #[error("Token mismatch!")]
203    TokenMismatch,
204
205    #[error(transparent)]
206    TryFromSlice(#[from] std::array::TryFromSliceError),
207
208    #[error(transparent)]
209    Anyhow(#[from] anyhow::Error),
210}
211
212pub fn read_http_host(in_buffer: &[u8]) -> Result<String> {
213    let mut lines = in_buffer.split(|&x| x == b'\n');
214    let host = lines
215        .find(|x| x.to_ascii_lowercase().starts_with(b"host:"))
216        .ok_or_else(|| anyhow::anyhow!("No host"))?;
217
218    let host = String::from_utf8_lossy(&host[5..]).trim().to_string();
219    Ok(host)
220}
221
222pub enum ConnectorStream {
223    TcpTlsClient(Box<tokio_rustls::client::TlsStream<TcpStream>>),
224    TcpTlsServer(Box<tokio_rustls::server::TlsStream<TcpStream>>),
225    Quic((quinn::SendStream, quinn::RecvStream)),
226}
227
228impl ConnectorStream {
229    pub async fn shutdown(&mut self) {
230        tokio::time::sleep(Duration::from_millis(100)).await;
231        match self {
232            ConnectorStream::TcpTlsClient(stream) => {
233                _ = stream.flush().await;
234                _ = stream.shutdown().await;
235            }
236            ConnectorStream::TcpTlsServer(stream) => {
237                _ = stream.flush().await;
238                _ = stream.shutdown().await;
239            }
240            ConnectorStream::Quic((send, recv)) => {
241                _ = send.flush().await;
242                _ = send.shutdown().await;
243                _ = recv.stop(VarInt::from_u32(0));
244            }
245        }
246    }
247
248    pub const fn get_name(&self) -> &'static str {
249        match &self {
250            ConnectorStream::TcpTlsClient(_) => "TCP",
251            ConnectorStream::TcpTlsServer(_) => "TCP",
252            ConnectorStream::Quic(_) => "UDP",
253        }
254    }
255}
256
257impl AsyncWrite for ConnectorStream {
258    fn poll_write(
259        self: Pin<&mut Self>,
260        cx: &mut Context<'_>,
261        buf: &[u8],
262    ) -> std::task::Poll<std::result::Result<usize, std::io::Error>> {
263        match self.get_mut() {
264            ConnectorStream::TcpTlsClient(stream) => Pin::new(stream).poll_write(cx, buf),
265            ConnectorStream::TcpTlsServer(stream) => Pin::new(stream).poll_write(cx, buf),
266            ConnectorStream::Quic((stream, _)) => {
267                Pin::new(stream).poll_write(cx, buf).map(|r| match r {
268                    Ok(n) => std::io::Result::Ok(n),
269                    Err(e) => std::io::Result::Err(e.into()),
270                })
271            }
272        }
273    }
274
275    fn poll_flush(
276        self: Pin<&mut Self>,
277        cx: &mut Context<'_>,
278    ) -> std::task::Poll<std::result::Result<(), std::io::Error>> {
279        match self.get_mut() {
280            ConnectorStream::TcpTlsClient(stream) => Pin::new(stream).poll_flush(cx),
281            ConnectorStream::TcpTlsServer(stream) => Pin::new(stream).poll_flush(cx),
282            ConnectorStream::Quic((stream, _)) => Pin::new(stream).poll_flush(cx),
283        }
284    }
285
286    fn poll_shutdown(
287        self: Pin<&mut Self>,
288        cx: &mut Context<'_>,
289    ) -> std::task::Poll<std::result::Result<(), std::io::Error>> {
290        match self.get_mut() {
291            ConnectorStream::TcpTlsClient(stream) => Pin::new(stream).poll_shutdown(cx),
292            ConnectorStream::TcpTlsServer(stream) => Pin::new(stream).poll_shutdown(cx),
293            ConnectorStream::Quic((stream, _)) => Pin::new(stream).poll_shutdown(cx),
294        }
295    }
296}
297
298impl AsyncRead for ConnectorStream {
299    fn poll_read(
300        self: Pin<&mut Self>,
301        cx: &mut Context<'_>,
302        buf: &mut tokio::io::ReadBuf<'_>,
303    ) -> std::task::Poll<std::io::Result<()>> {
304        match self.get_mut() {
305            ConnectorStream::TcpTlsClient(stream) => Pin::new(stream).poll_read(cx, buf),
306            ConnectorStream::TcpTlsServer(stream) => Pin::new(stream).poll_read(cx, buf),
307            ConnectorStream::Quic((_, stream)) => Pin::new(stream).poll_read(cx, buf),
308        }
309    }
310}
fkm_proxy/utils/mod.rs

fkm_proxy/utils/
mod.rs
