From de77445372fe8f6561510349a3d3555425b49823 Mon Sep 17 00:00:00 2001
From: alterdekim <alterwain@protonmail.com>
Date: Mon, 19 Aug 2024 16:37:08 +0300
Subject: [PATCH]  Changes to be committed: 	modified:   src/client.rs 
 modified:   src/main.rs 	modified:   src/server.rs

---
 src/client.rs |  6 ++++--
 src/main.rs   | 19 ++++++++++---------
 src/server.rs | 27 +++++++--------------------
 3 files changed, 21 insertions(+), 31 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 50b1d02..1de0909 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -6,7 +6,7 @@ use std::io::{Read, Write};
 use tun2::BoxError;
 use log::{error, info, LevelFilter};
 use std::sync::Arc;
-use std::net::SocketAddr;
+use std::net::{ SocketAddr, Ipv4Addr };
 use std::collections::HashMap;
 use std::process::Command;
 use tokio::io::AsyncReadExt;
@@ -111,7 +111,9 @@ pub async fn client_mode(client_config: ClientConfiguration) {
         }
     });
 
-    let handshake = UDPVpnHandshake{ public_key: client_config.client.public_key.into_bytes() };
+    let pkey = base64::decode(client_config.client.public_key).unwrap();
+    info!("Handshake public_key: {:?}", pkey.len());
+    let handshake = UDPVpnHandshake{ public_key: pkey, request_ip: client_config.client.address.parse::<Ipv4Addr>().unwrap().octets() };
     sock_snd.send(&handshake.serialize()).await.unwrap();
 
     loop {
diff --git a/src/main.rs b/src/main.rs
index a256323..9af39fc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,7 +10,7 @@ use serde_derive::Deserialize;
 use std::str::FromStr;
 use x25519_dalek::{StaticSecret, PublicKey};
 use rand::{rngs::StdRng, SeedableRng};
-use base64;
+use base64::prelude::*;
 
 //mod tcp_client;
 //mod tcp_server;
@@ -31,7 +31,7 @@ impl VpnPacket {
         u64::from_be_bytes(d)
     }
 
-    fn deserialize(d: Vec<u8>) -> Result<VpnPacket, Error> {
+    fn deserialize(d: Vec<u8>) -> Result<Self, Error> {
         Ok(VpnPacket{ data: d })
     }
 }
@@ -48,13 +48,14 @@ impl UDPSerializable for UDPVpnPacket {
 }
 
 struct UDPVpnHandshake {
-    public_key: Vec<u8>
+    public_key: Vec<u8>,
+    request_ip: [u8; 4]
 }
 
 impl UDPSerializable for UDPVpnHandshake {
     fn serialize(&self) -> Vec<u8> {
         let h: &[u8] = &[0];
-        [h, &self.public_key[..]].concat()
+        [h, &self.public_key[..], &self.request_ip[..]].concat()
     }
 }
 
@@ -73,7 +74,7 @@ struct ServerInterface {
 }
 
 #[derive(Serialize, Deserialize, PartialEq, Debug)]
-struct ServerPeer {
+pub struct ServerPeer {
     public_key: String,
     ip: Ipv4Addr
 }
@@ -105,8 +106,8 @@ impl ServerConfiguration {
         ServerConfiguration { interface: ServerInterface { 
                 bind_address: String::from_str(bind_address).unwrap(), 
                 internal_address: String::from_str(internal_address).unwrap(), 
-                private_key: base64::encode(secret.as_bytes()), 
-                public_key: base64::encode(PublicKey::from(&secret).as_bytes()),
+                private_key: BASE64_STANDARD.encode(secret.as_bytes()), 
+                public_key: BASE64_STANDARD.encode(PublicKey::from(&secret).as_bytes()),
                 broadcast_mode, 
                 keepalive 
             }, 
@@ -156,8 +157,8 @@ impl ClientConfiguration {
         let secret = StaticSecret::new(&mut csprng);
         ClientConfiguration { 
             client: ClientInterface { 
-                private_key: base64::encode(secret.as_bytes()), 
-                public_key: base64::encode(PublicKey::from(&secret).as_bytes()),
+                private_key: BASE64_STANDARD.encode(secret.as_bytes()), 
+                public_key: BASE64_STANDARD.encode(PublicKey::from(&secret).as_bytes()),
                 address: String::from_str(internal_address).unwrap() 
             }, 
             server: EndpointInterface { 
diff --git a/src/server.rs b/src/server.rs
index 2340688..dda2fc7 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -11,7 +11,7 @@ use std::collections::HashMap;
 use tokio::io::AsyncReadExt;
 use std::process::Command;
 
-use crate::{ VpnPacket, ServerConfiguration, UDPSerializable };
+use crate::{ VpnPacket, ServerConfiguration, UDPSerializable, ServerPeer };
 
 pub async fn server_mode(server_config: ServerConfiguration) {
     info!("Starting server...");
@@ -34,6 +34,7 @@ pub async fn server_mode(server_config: ServerConfiguration) {
     let sock_rec = Arc::new(sock);
     let sock_snd = sock_rec.clone();
     let addresses = Arc::new(Mutex::new(HashMap::<IpAddr, UDPeer>::new()));
+    let peers = Arc::new(Mutex::new(Vec::<ServerPeer>::new()));
 
     let (send2tun, recv2tun) = unbounded::<Vec<u8>>();
 
@@ -65,10 +66,12 @@ pub async fn server_mode(server_config: ServerConfiguration) {
     
     let mut buf = vec![0; 2048];
     let addrs_lp = addresses.clone();
-    
+    let peers_lp = peers.clone();
+
     loop {
         if let Ok((len, addr)) = sock_rec.recv_from(&mut buf).await {
             let mut mp = addrs_lp.lock().await;
+            let mut plp = peers_lp.lock().await;
             match buf.first() {
                 Some(h) => {
                     match h {
@@ -88,6 +91,7 @@ pub async fn server_mode(server_config: ServerConfiguration) {
                 },
                 None => error!("There is no header")
             }
+            drop(plp);
             drop(mp);
         }
     }
@@ -95,21 +99,4 @@ pub async fn server_mode(server_config: ServerConfiguration) {
 
 struct UDPeer {
     addr: SocketAddr
-}
-
-/*struct WrappedUDP {
-    sock_rec: Arc<UdpSocket>,
-    sock_snd: Arc<UdpSocket>,
-    addresses: Arc<Mutex<HashMap<IpAddr, UDPeer>>>
-}
-
-impl WrappedUDP {
-    pub async fn new(addr: &str) -> Self {
-        
-        WrappedUDP { sock_rec, sock_snd, addresses }
-    }
-
-    pub async fn init(&self) {
-        
-    }
-}*/
\ No newline at end of file
+}
\ No newline at end of file