fkm_proxy/utils/
mod.rs

1use anyhow::Result;
2use std::{
3    net::{SocketAddr, ToSocketAddrs},
4    time::Duration,
5};
6use thiserror::Error;
7use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
8
9pub mod certs;
10pub mod http;
11pub mod udp;
12
13#[derive(Debug)]
14pub struct HelloPacket {
15    pub hp_type: HelloPacketType,
16    //pub hash: u64,
17    pub token: u128,
18    pub own_ssl: bool,
19    pub tunnel_id: u128,
20}
21
22#[derive(Debug, PartialEq)]
23pub enum HelloPacketType {
24    Connector = 0,
25    Tunnel = 1,
26
27    Invalid,
28}
29
30impl HelloPacketType {
31    pub fn to_u8(&self) -> u8 {
32        match self {
33            HelloPacketType::Connector => 0,
34            HelloPacketType::Tunnel => 1,
35            HelloPacketType::Invalid => u8::MAX,
36        }
37    }
38
39    pub fn from_u8(val: u8) -> Self {
40        match val {
41            0 => HelloPacketType::Connector,
42            1 => HelloPacketType::Tunnel,
43            _ => HelloPacketType::Invalid,
44        }
45    }
46}
47
48impl HelloPacket {
49    pub const fn buf_size() -> usize {
50        80
51    }
52
53    pub fn to_buf(&self) -> [u8; 80] {
54        let mut tmp = [0; 80];
55        tmp[0] = self.hp_type.to_u8();
56        //tmp[1..9].copy_from_slice(&self.hash.to_be_bytes());
57        tmp[10..26].copy_from_slice(&self.token.to_be_bytes());
58        tmp[26] = self.own_ssl as u8;
59        tmp[27..43].copy_from_slice(&self.tunnel_id.to_be_bytes());
60
61        tmp
62    }
63
64    pub fn from_buf(buf: &[u8; 80]) -> Self {
65        Self {
66            hp_type: HelloPacketType::from_u8(buf[0]),
67            //hash: u64::from_be_bytes(buf[1..9].try_into().unwrap()),
68            token: u128::from_be_bytes(buf[10..26].try_into().unwrap()),
69            own_ssl: buf[26] != 0,
70            tunnel_id: u128::from_be_bytes(buf[27..43].try_into().unwrap()),
71        }
72    }
73}
74
75#[derive(Debug, PartialEq)]
76pub enum ConnectorPacketType {
77    Ping = 0,
78    TunnelRequest = 1,
79    Close = 2,
80
81    Invalid,
82}
83
84impl ConnectorPacketType {
85    pub fn to_u8(&self) -> u8 {
86        match self {
87            ConnectorPacketType::Ping => 0,
88            ConnectorPacketType::TunnelRequest => 1,
89            ConnectorPacketType::Close => 2,
90            ConnectorPacketType::Invalid => u8::MAX,
91        }
92    }
93
94    pub fn from_u8(val: u8) -> Self {
95        match val {
96            0 => ConnectorPacketType::Ping,
97            1 => ConnectorPacketType::TunnelRequest,
98            2 => ConnectorPacketType::Close,
99            _ => ConnectorPacketType::Invalid,
100        }
101    }
102}
103
104#[derive(Debug)]
105pub struct ConnectorPacket {
106    pub packet_type: ConnectorPacketType,
107    pub tunnel_id: u128,
108    pub ssl: bool,
109}
110
111impl ConnectorPacket {
112    pub const fn buf_size() -> usize {
113        20
114    }
115
116    pub fn to_buf(&self) -> [u8; 20] {
117        let mut tmp = [0; 20];
118        tmp[0] = self.packet_type.to_u8();
119        tmp[1..17].copy_from_slice(&self.tunnel_id.to_be_bytes());
120        tmp[17] = self.ssl as u8;
121
122        tmp
123    }
124
125    pub fn from_buf(buf: &[u8; 20]) -> Self {
126        Self {
127            packet_type: ConnectorPacketType::from_u8(buf[0]),
128            tunnel_id: u128::from_be_bytes(buf[1..17].try_into().unwrap()),
129            ssl: buf[17] != 0,
130        }
131    }
132}
133
134pub fn parse_socketaddr(arg: &str) -> Result<SocketAddr> {
135    for i in 0..10 {
136        let res = arg.to_socket_addrs();
137
138        match res {
139            Ok(addrs) => {
140                for addr in addrs {
141                    if addr.is_ipv4() {
142                        return Ok(addr);
143                    }
144                }
145            }
146            Err(e) => {
147                println!("[clap parse_socketaddr] (Try: {}) {e:?}", i + 1);
148                std::thread::sleep(Duration::from_millis(5000));
149            }
150        }
151    }
152
153    Err(anyhow::anyhow!("No ipv4 socketaddr found!"))
154}
155
156pub fn generate_string_packet(string: &str) -> Result<Vec<u8>> {
157    let mut bytes = string.as_bytes().to_vec();
158    bytes.push(0); // null terminator
159
160    Ok(bytes)
161}
162
163pub async fn send_string_to_stream<T>(stream: &mut T, string: &str) -> Result<()>
164where
165    T: AsyncWrite + Unpin,
166{
167    let bytes = generate_string_packet(string)?;
168    stream.write_all(&bytes).await?;
169
170    Ok(())
171}
172
173pub async fn read_string_from_stream<T>(stream: &mut T) -> Result<String>
174where
175    T: AsyncRead + Unpin,
176{
177    let mut buffer = Vec::new();
178    loop {
179        let byte = stream.read_u8().await?;
180        if byte == 0 {
181            break;
182        }
183
184        buffer.push(byte);
185    }
186
187    Ok(String::from_utf8(buffer)?)
188}
189
190#[derive(Error, Debug)]
191pub enum HelloPacketError {
192    #[error("Token mismatch!")]
193    TokenMismatch,
194
195    #[error(transparent)]
196    TryFromSlice(#[from] std::array::TryFromSliceError),
197
198    #[error(transparent)]
199    Anyhow(#[from] anyhow::Error),
200}
201
202pub fn read_http_host(in_buffer: &[u8]) -> Result<String> {
203    let mut lines = in_buffer.split(|&x| x == b'\n');
204    let host = lines
205        .find(|x| x.to_ascii_lowercase().starts_with(b"host:"))
206        .ok_or_else(|| anyhow::anyhow!("No host"))?;
207
208    let host = String::from_utf8_lossy(&host[5..]).trim().to_string();
209    Ok(host)
210}