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

fkm_proxy/utils/
mod.rs
