From 4b42c23b24a0fd7658cb9d0b485d51ee2d4b505a Mon Sep 17 00:00:00 2001
From: alterdekim <alterwain@protonmail.com>
Date: Mon, 12 Aug 2024 18:43:15 +0300
Subject: [PATCH]  Changes to be committed: 	modified:   src/main.rs 
 modified:   src/tcp_client.rs 	modified:   src/tcp_server.rs

---
 src/main.rs       |  9 ++++++---
 src/tcp_client.rs | 11 ++++++-----
 src/tcp_server.rs | 11 ++++++-----
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index d6e997e..1d98fd5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -13,16 +13,19 @@ use serde_derive::Deserialize;
 mod tcp_client;
 mod tcp_server;
 
+const HEADER: [u8;3] = [0x56, 0x66, 0x76];
+const TAIL: [u8;3] = [0x76, 0x66, 0x56];
+
 #[derive(Serialize, Deserialize)]
 struct VpnPacket {
-    start: u8,
+    start: Vec<u8>,
     data: Vec<u8>,
-    end: u8
+    end: Vec<u8>
 }
 
 impl VpnPacket {
     fn init(d: Vec<u8>) -> Self {
-        VpnPacket{start: 0x56, data: d, end: 0x66}
+        VpnPacket{start: (&HEADER).to_vec(), data: d, end: (&TAIL).to_vec()}
     }
 }
 
diff --git a/src/tcp_client.rs b/src/tcp_client.rs
index 10fdba9..f0d0c29 100644
--- a/src/tcp_client.rs
+++ b/src/tcp_client.rs
@@ -11,7 +11,7 @@ use std::collections::HashMap;
 use std::process::Command;
 use tokio::io::AsyncReadExt;
 
-use crate::VpnPacket;
+use crate::{VpnPacket, HEADER, TAIL};
 
 fn configure_routes() {
     let ip_output = Command::new("ip")
@@ -92,18 +92,19 @@ pub async fn client_mode(remote_addr: String) {
     });
 
     tokio::spawn(async move {
-        let mut buf = Vec::<u8>::new();
-        while let Ok(n) = dev_reader.read_to_end(&mut buf) {
+        let mut buf = vec![0; 1024];
+        while let Ok(n) = dev_reader.read(&mut buf) {
             dx.send(buf[..n].to_vec()).unwrap();
         }
     });
 
     tokio::spawn(async move {
-        let mut buf = Vec::<u8>::new();
+        let mut buf = vec![0; 1024];
         loop {
-            if let Ok(n) = sock_reader.read_to_end(&mut buf).await {
+            if let Ok(n) = sock_reader.read(&mut buf).await {
                 //info!("Catch from socket: {:?}", &buf[..n]);
                 let vpn_packet: VpnPacket = bincode::deserialize(&buf[..n]).unwrap();
+                if vpn_packet.start != &HEADER || vpn_packet.end != &TAIL { error!("Bad packet"); continue; }
                 tx.send(vpn_packet.data).unwrap();
             }
         }
diff --git a/src/tcp_server.rs b/src/tcp_server.rs
index 3e227f3..a1ef514 100644
--- a/src/tcp_server.rs
+++ b/src/tcp_server.rs
@@ -10,7 +10,7 @@ use std::net::SocketAddr;
 use std::collections::HashMap;
 use tokio::io::AsyncReadExt;
 
-use crate::VpnPacket;
+use crate::{VpnPacket, HEADER, TAIL};
 
 pub async fn server_mode() {
     info!("Starting server...");
@@ -38,8 +38,8 @@ pub async fn server_mode() {
     });
 
     tokio::spawn(async move {
-        let mut buf = Vec::<u8>::new();
-        while let Ok(n) = dev_reader.read_to_end(&mut buf) {
+        let mut buf = vec![0; 1024];
+        while let Ok(n) = dev_reader.read(&mut buf) {
             dx.send(buf[..n].to_vec()).unwrap();
         }
     });
@@ -64,11 +64,12 @@ pub async fn server_mode() {
         });
 
         tokio::spawn(async move {
-            let mut buf = Vec::<u8>::new();
+            let mut buf = vec![0; 1024];
             loop {
-                if let Ok(n) = sock_reader.read_to_end(&mut buf).await {
+                if let Ok(n) = sock_reader.read(&mut buf).await {
                     info!("Catched from sock: {:?}", &buf[..n]);
                     let vpn_packet: VpnPacket = bincode::deserialize(&buf[..n]).unwrap();
+                    if vpn_packet.start != &HEADER || vpn_packet.end != &TAIL { error!("Bad packet"); continue; }
                     thread_tx.send(vpn_packet.data).unwrap();
                 }
             }