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

fkm_proxy/utils/
mod.rs
