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

---
 src/client.rs |  3 +--
 src/main.rs   | 19 ++++++++++++-------
 src/server.rs |  4 +++-
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 1de0909..9b9f00a 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -112,8 +112,7 @@ pub async fn client_mode(client_config: ClientConfiguration) {
     });
 
     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() };
+    let handshake = UDPVpnHandshake{ public_key: pkey, request_ip: client_config.client.address.parse::<Ipv4Addr>().unwrap() };
     sock_snd.send(&handshake.serialize()).await.unwrap();
 
     loop {
diff --git a/src/main.rs b/src/main.rs
index 9af39fc..7335986 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,7 +10,6 @@ use serde_derive::Deserialize;
 use std::str::FromStr;
 use x25519_dalek::{StaticSecret, PublicKey};
 use rand::{rngs::StdRng, SeedableRng};
-use base64::prelude::*;
 
 //mod tcp_client;
 //mod tcp_server;
@@ -49,13 +48,19 @@ impl UDPSerializable for UDPVpnPacket {
 
 struct UDPVpnHandshake {
     public_key: Vec<u8>,
-    request_ip: [u8; 4]
+    request_ip: Ipv4Addr // [u8; 4]
 }
 
 impl UDPSerializable for UDPVpnHandshake {
     fn serialize(&self) -> Vec<u8> {
         let h: &[u8] = &[0];
-        [h, &self.public_key[..], &self.request_ip[..]].concat()
+        [h, &self.public_key[..], &self.request_ip.octets()].concat()
+    }
+}
+
+impl UDPVpnHandshake {
+    fn deserialize(data: &Vec<u8>) -> Self {
+        UDPVpnHandshake { public_key: data[1..=32].to_vec(), request_ip: Ipv4Addr::new(data[33], data[34], data[35], data[36]) }
     }
 }
 
@@ -106,8 +111,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_STANDARD.encode(secret.as_bytes()), 
-                public_key: BASE64_STANDARD.encode(PublicKey::from(&secret).as_bytes()),
+                private_key: base64::encode(secret.as_bytes()), 
+                public_key: base64::encode(PublicKey::from(&secret).as_bytes()),
                 broadcast_mode, 
                 keepalive 
             }, 
@@ -157,8 +162,8 @@ impl ClientConfiguration {
         let secret = StaticSecret::new(&mut csprng);
         ClientConfiguration { 
             client: ClientInterface { 
-                private_key: BASE64_STANDARD.encode(secret.as_bytes()), 
-                public_key: BASE64_STANDARD.encode(PublicKey::from(&secret).as_bytes()),
+                private_key: base64::encode(secret.as_bytes()), 
+                public_key: base64::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 dda2fc7..5e46a7e 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, ServerPeer };
+use crate::{ ServerConfiguration, ServerPeer, UDPSerializable, UDPVpnHandshake, VpnPacket };
 
 pub async fn server_mode(server_config: ServerConfiguration) {
     info!("Starting server...");
@@ -77,6 +77,8 @@ pub async fn server_mode(server_config: ServerConfiguration) {
                     match h {
                         0 => {
                             // (&buf[1..len]).to_vec()
+                            let handshake = UDPVpnHandshake::deserialize(&buf);
+                            info!("Got handshake! ip: {:?}; key: {:?}", handshake.request_ip, base64::encode(handshake.public_key));
                             let internal_ip = IpAddr::V4(Ipv4Addr::new(10,8,0,2));
                             info!("Got handshake");
                             mp.insert(internal_ip, UDPeer { addr });