feat: better parsing of HelloPacket

src/client/main.rs

@@ -9,9 +9,8 @@ use tokio::{
 };
 use utils::{
     certs::{cert_from_str, key_from_str},
-    generate_hello_packet,
     http::construct_http_redirect,
-    parse_socketaddr, read_string_from_stream,
+    parse_socketaddr, read_string_from_stream, HelloPacket,
 };
 
 const MAX_REQUEST_TIME: u128 = 1000;
@@ -81,28 +80,33 @@ async fn connector(args: &Args) -> Result<()> {
 
     let stream = TcpStream::connect(&args.proxy_addr).await?;
     let mut stream = acceptor.accept(stream).await?;
-    let mut hello_packet =
-        generate_hello_packet(0, &args.token, &args.hash, args.ssl_addr.is_some())?;
 
-    stream.write_all(&hello_packet).await?;
+    let mut hello_packet = HelloPacket {
+        hp_type: utils::HelloPacketType::Connector,
+        token: args.token,
+        hash: args.hash,
+        own_ssl: args.ssl_addr.is_some(),
+        tunnel_id: 0,
+    };
+
+    stream.write_all(&hello_packet.to_buf()).await?;
     let nonssl_port = stream.read_u16().await?;
     let ssl_port = stream.read_u16().await?;
     let domain = read_string_from_stream(&mut stream).await?;
     tracing::info!(
         "Access through:\n - http://{domain}:{nonssl_port}\n - https://{domain}:{ssl_port}"
     );
 
-    hello_packet[0] = 0x01; // 0x01 - tunnel
+    hello_packet.hp_type = utils::HelloPacketType::Tunnel;
 
-    let mut buf = [0; 16];
     loop {
         let first_byte = stream.read_u8().await?;
         if first_byte == 0x69 {
             stream.write_u8(0x69).await?;
             continue; // ping/pong
         }
 
-        stream.read_exact(&mut buf).await?;
+        let tunnel_id = stream.read_u128().await?;
         let ssl = first_byte == 0x01;
         let domain = domain.to_string();
         let acceptor = acceptor.clone();
@@ -112,10 +116,11 @@ async fn connector(args: &Args) -> Result<()> {
             ssl_addr: args.ssl_addr.unwrap_or(args.addr),
             nonssl_addr: args.addr,
             redirect_ssl: args.redirect_ssl,
-            http3: args.http3
+            http3: args.http3,
         };
 
-        hello_packet[26..42].copy_from_slice(&buf[0..16]);
+        hello_packet.tunnel_id = tunnel_id;
+        let hello_packet = hello_packet.to_buf();
         tokio::task::spawn(async move {
             let res = spawn_tunnel(
                 hello_packet,

src/server/tunnel.rs

@@ -76,7 +76,10 @@ async fn connector_handler(
         .await
         .ok_or_else(|| anyhow!("Cant find token!"))?;
 
-    let own_ssl = ::utils::parse_hello_packet(token, &connection_buff)?;
+    let hello_packet = ::utils::HelloPacket::from_buf(&connection_buff);
+    if hello_packet.token != token {
+        return Err(::utils::HelloPacketError::TokenMismatch.into());
+    }
 
     // im the connector!
     if connection_buff[0] == 0 {
@@ -91,7 +94,9 @@ async fn connector_handler(
         ::utils::send_string_to_stream(&mut stream, &domain).await?;
 
         let (tx, rx) = kanal::unbounded_async::<TunnelRequest>();
-        state.insert_tunnel_connector(token, tx, own_ssl).await;
+        state
+            .insert_tunnel_connector(token, tx, hello_packet.own_ssl)
+            .await;
         let res = connector_loop(&mut stream, rx).await;
         if let Err(e) = res {
             tracing::error!("Connector loop: {e:?}");

utils/src/lib.rs

@@ -7,45 +7,73 @@ pub mod certs;
 pub mod http;
 pub mod udp;
 
-pub fn parse_socketaddr(arg: &str) -> Result<SocketAddr> {
-    let addrs = arg.to_socket_addrs()?;
-    for addr in addrs {
-        if addr.is_ipv4() {
-            return Ok(addr);
+#[derive(Debug)]
+pub struct HelloPacket {
+    pub hp_type: HelloPacketType,
+    pub hash: u64,
+    pub token: u128,
+    pub own_ssl: bool,
+    pub tunnel_id: u128,
+}
+
+#[derive(Debug)]
+pub enum HelloPacketType {
+    Connector = 0,
+    Tunnel = 1,
+
+    Invalid,
+}
+
+impl HelloPacketType {
+    pub fn to_u8(&self) -> u8 {
+        match self {
+            HelloPacketType::Connector => 0,
+            HelloPacketType::Tunnel => 1,
+            HelloPacketType::Invalid => u8::MAX,
         }
     }
 
-    Err(anyhow::anyhow!("No ipv4 socketaddr found!"))
+    pub fn from_u8(val: u8) -> Self {
+        match val {
+            0 => HelloPacketType::Connector,
+            1 => HelloPacketType::Tunnel,
+            _ => HelloPacketType::Invalid,
+        }
+    }
 }
 
-pub fn generate_hello_packet(
-    connector_type: u8,
-    token: &u128,
-    hash: &u64,
-    own_ssl: bool,
-) -> Result<[u8; 80], HelloPacketError> {
-    let mut conn_buff = [0u8; 80];
-    conn_buff[0] = connector_type;
-    conn_buff[1..9].copy_from_slice(&hash.to_be_bytes());
-    conn_buff[10..26].copy_from_slice(&token.to_be_bytes());
-    conn_buff[26] = own_ssl as u8;
-
-    Ok(conn_buff)
-}
+impl HelloPacket {
+    pub fn to_buf(&self) -> [u8; 80] {
+        let mut tmp = [0; 80];
+        tmp[0] = self.hp_type.to_u8();
+        tmp[1..9].copy_from_slice(&self.hash.to_be_bytes());
+        tmp[10..26].copy_from_slice(&self.token.to_be_bytes());
+        tmp[26] = self.own_ssl as u8;
+        tmp[26..42].copy_from_slice(&self.tunnel_id.to_be_bytes());
 
-pub fn parse_hello_packet(
-    token: u128,
-    connection_buff: &[u8; 80],
-) -> Result<bool, HelloPacketError> {
-    //let parsed_connector_type = connection_buff[0];
-    //let parsed_hash = u64::from_be_bytes(connection_buff[1..9].try_into()?);
-    let parsed_token = u128::from_be_bytes(connection_buff[10..26].try_into()?);
+        tmp
+    }
 
-    if parsed_token != token {
-        return Err(HelloPacketError::TokenMismatch);
+    pub fn from_buf(buf: &[u8; 80]) -> Self {
+        Self {
+            hp_type: HelloPacketType::from_u8(buf[0]),
+            hash: u64::from_be_bytes(buf[1..9].try_into().unwrap()),
+            token: u128::from_be_bytes(buf[10..26].try_into().unwrap()),
+            own_ssl: buf[26] != 0,
+            tunnel_id: u128::from_be_bytes(buf[26..42].try_into().unwrap()),
+        }
     }
+}
 
-    Ok(connection_buff[26] != 0x00)
+pub fn parse_socketaddr(arg: &str) -> Result<SocketAddr> {
+    let addrs = arg.to_socket_addrs()?;
+    for addr in addrs {
+        if addr.is_ipv4() {
+            return Ok(addr);
+        }
+    }
+
+    Err(anyhow::anyhow!("No ipv4 socketaddr found!"))
 }
 
 pub fn generate_string_packet(string: &str) -> Result<Vec<u8>> {

utils/src/udp.rs

@@ -10,7 +10,7 @@ use tokio::{
 };
 
 type TunnelMapTest = Arc<RwLock<HashMap<SocketAddr, UnboundedSender<Vec<u8>>>>>;
-async fn server_test(
+pub async fn server_accept(
     listener: &Arc<UdpSocket>,
     recv_buf: &mut [u8],
     tunnel_map: &TunnelMapTest,
@@ -43,7 +43,7 @@ async fn server_test(
     }
 }
 
-struct UdpClient {
+pub struct UdpClient {
     tunnel_map: TunnelMapTest,
     rx: UnboundedReceiver<Vec<u8>>,
     addr: SocketAddr,
@@ -62,7 +62,7 @@ impl UdpClient {
         }
     }
 
-    pub async fn recv(&mut self) -> Option<Vec<u8>> {
+    async fn recv(&mut self) -> Option<Vec<u8>> {
         let timeout = tokio::time::timeout(Duration::from_secs(45), self.rx.recv()).await;
         match timeout {
             Ok(recv) => recv,