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

fkm_proxy/utils/
mod.rs
