finish register request

src/client/main.rs

@@ -1,4 +1,6 @@
+mod net;
 use pea_2_pea::*;
+use rand::RngCore;
 
 use std::{
     io::{Error, ErrorKind, Read, Write},
@@ -67,88 +69,14 @@ fn main() -> std::io::Result<()> {
             .unwrap();
 
         let mut buf: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE];
-        let mut data_lenght;
+        // query here
-        loop {
+        let mut data_lenght: usize = net::query_request(&mut buf, &server_SocketAddr, socket)?;
-            match socket.send_to(&[ServerMethods::QUERY as u8], &server_SocketAddr) {
-                Ok(s) => {
-                    #[cfg(debug_assertions)]
-                    eprintln!("send {} bytes", s);
-                }
-                Err(e) => {
-                    panic!("Error sending data: {}", e);
-                }
-            }
-
-            match socket.recv_from(&mut buf) {
-                Ok((data_length_recved, _src)) => {
-                    data_lenght = data_length_recved;
-                    if buf[0] != 0 {
-                        continue;
-                    }
-                    break;
-                }
-                Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::TimedOut => {
-                    // timedout
-                    continue;
-                }
-                Err(e) => {
-                    panic!("Error receiving data: {}", e);
-                }
-            }
-        }
 
         let mut public_sock_addr: Vec<u8> = buf[1..data_lenght].to_vec();
 
         // register network
 
-        buf[0] = ServerMethods::REGISTER as u8;
+        let mut salt: Option<[u8; SALT_AND_IV_SIZE as usize]>;
-        buf[RegisterRequestDataPositions::ENCRYPTED as usize] = match cli.password {
-            Some(_) => true as u8,
-            None => false as u8,
-        };
-        buf[RegisterRequestDataPositions::ID_LEN as usize] = cli.network_id.len() as u8;
-        buf[RegisterRequestDataPositions::SOCKADDR_LEN as usize] =
-            server_SocketAddr.to_string().len() as u8;
-
-        buf[RegisterRequestDataPositions::DATA as usize
-            ..RegisterRequestDataPositions::DATA as usize + cli.network_id.len()]
-            .copy_from_slice(cli.network_id.as_bytes());// store network id
-
-
-
-        match cli.password {
-            Some(s) => {},
-            None => {},// do nothig
-        }
-
-        
-
-        buf[RegisterRequestDataPositions::DATA as usize + cli.network_id.len()..RegisterRequestDataPositions::DATA as usize + cli.network_id.len() + ]
-
-        match buf[0] {
-            x if x == ServerResponse::OK as u8 => {
-                eprintln!("network registered");
-            }
-
-            x if x == ServerResponse::GENERAL_ERROR as u8 => {
-                eprintln!(
-                    "{}",
-                    match std::str::from_utf8(&buf[1..data_lenght]) {
-                        Ok(s) => s.to_string(),
-                        Err(e) => {
-                            panic!("id to utf-8 failed: {}", e);
-                        }
-                    }
-                )
-            }
-            x if x == ServerResponse::ID_EXISTS as u8 => {
-                panic!("network ID already exist try differnt one!");
-            }
-
-            _ => {
-                panic!("unknown responce from server code: 0x{:02x}", buf[0])
-            }
-        }
     }
     Ok(())
 }

src/client/net.rs

@@ -0,0 +1,144 @@
+use std::{
+    io::ErrorKind,
+    net::{SocketAddr, UdpSocket},
+};
+
+use pea_2_pea::*;
+use rand::RngCore;
+
+// return data_lenght and number of retryes
+pub fn send_and_recv_with_retry(
+    buf: &mut [u8; BUFFER_SIZE],
+    dst: &SocketAddr,
+    socket: UdpSocket,
+    retry_max: usize,
+) -> Result<(usize, usize), ServerErrorResponses> {
+    let mut send_buf = *buf;
+    let mut retry_count: usize = 0;
+    loop {
+        match socket.send_to(&mut send_buf, dst) {
+            Ok(s) => {
+                #[cfg(debug_assertions)]
+                eprintln!("send {} bytes", s);
+            }
+            Err(e) => {
+                panic!("Error sending data: {}", e);
+            }
+        }
+
+        match socket.recv_from(buf) {
+            Ok((data_lenght, src)) => {
+                if src != *dst {
+                    continue;
+                }
+                match buf[0] {
+                    x if x == send_buf[0] as u8 => {
+                        return Ok((data_lenght, retry_count));
+                    }
+                    x if x == ServerResponse::GENERAL_ERROR as u8 => {
+                        return Err(ServerErrorResponses::IO(std::io::Error::new(
+                            std::io::ErrorKind::InvalidData,
+                            match std::str::from_utf8(&buf[1..data_lenght]) {
+                                // the firts byte is compensated for sice this is len not index
+                                Ok(s) => s.to_string(),
+                                Err(e) => format!("invalid error string: {}", e).to_string(),
+                            },
+                        )));
+                    }
+                    x if x == ServerResponse::ID_DOESNT_EXIST as u8 => {
+                        return Err(ServerErrorResponses::ID_DOESNT_EXIST);
+                    }
+                    x if x == ServerResponse::ID_EXISTS as u8 => {
+                        return Err(ServerErrorResponses::ID_EXISTS);
+                    }
+                    _ => {
+                        continue;
+                    }
+                }
+            }
+            Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::TimedOut => {
+                // timedout
+                if retry_count >= retry_max {
+                    return Err(ServerErrorResponses::IO(std::io::Error::new(
+                        ErrorKind::TimedOut,
+                        "max retry count reached without responce",
+                    )));
+                }
+                retry_count += 1;
+                continue;
+            }
+            Err(e) => {
+                return Err(ServerErrorResponses::IO(e));
+            }
+        }
+    }
+}
+
+pub fn query_request(
+    buf: &mut [u8; BUFFER_SIZE],
+    dst: &SocketAddr,
+    socket: UdpSocket,
+) -> Result<usize, ServerErrorResponses> {
+    match send_and_recv_with_retry(buf, dst, socket, STANDARD_RETRY_MAX) {
+        Ok((data_lenght, _)) => return Ok(data_lenght),
+        Err(e) => return Err(e),
+    }
+}
+
+pub fn register_request(
+    buf: &mut [u8; BUFFER_SIZE],
+    dst: &SocketAddr,
+    socket: UdpSocket,
+    encryption_key: Option<[u8; 32]>,
+    salt_opt: Option<[u8; SALT_AND_IV_SIZE as usize]>,
+    mut public_sock_addr: Vec<u8>,
+    network_id: String,
+) -> Result<usize, ServerErrorResponses> {
+    buf[0] = ServerMethods::REGISTER as u8; // set metod identification byte
+    buf[RegisterRequestDataPositions::ENCRYPTED as usize] = match encryption_key {
+        // stor encryption flag byte
+        Some(_) => true as u8,
+        None => false as u8,
+    };
+    buf[RegisterRequestDataPositions::ID_LEN as usize] = network_id.len() as u8;
+
+    buf[RegisterRequestDataPositions::DATA as usize
+        ..RegisterRequestDataPositions::DATA as usize + network_id.len()]
+        .copy_from_slice(network_id.as_bytes()); // store network id
+
+    let mut iv: [u8; SALT_AND_IV_SIZE as usize] = [0; SALT_AND_IV_SIZE as usize];
+    let salt: [u8; SALT_AND_IV_SIZE as usize];
+    match salt_opt {
+        Some(s) => salt = s,
+        None => salt = [0; SALT_AND_IV_SIZE as usize],
+    }
+    match encryption_key {
+        Some(encryption_key) => {
+            let mut rng = rand::rng();
+            rng.fill_bytes(&mut iv);
+            public_sock_addr =
+                shared::crypto::encrypt(&encryption_key, &iv, public_sock_addr.as_slice()).unwrap();
+        }
+        None => {
+            iv = [0; SALT_AND_IV_SIZE as usize];
+        }
+    };
+
+    buf[RegisterRequestDataPositions::IV as usize
+        ..RegisterRequestDataPositions::IV as usize + SALT_AND_IV_SIZE as usize]
+        .copy_from_slice(&iv); // copy iv ad salt do the request
+    buf[RegisterRequestDataPositions::SALT as usize
+        ..RegisterRequestDataPositions::SALT as usize + SALT_AND_IV_SIZE as usize]
+        .copy_from_slice(&salt);
+
+    buf[RegisterRequestDataPositions::SOCKADDR_LEN as usize] = public_sock_addr.len() as u8;
+
+    buf[RegisterRequestDataPositions::DATA as usize + network_id.len()
+        ..RegisterRequestDataPositions::DATA as usize + network_id.len() + public_sock_addr.len()]
+        .copy_from_slice(&public_sock_addr);
+
+    match send_and_recv_with_retry(buf, dst, socket, STANDARD_RETRY_MAX) {
+        Ok((data_lenght, _)) => return Ok(data_lenght),
+        Err(e) => return Err(e),
+    }
+}

src/lib.rs

@@ -1,8 +1,11 @@
+use core::fmt;
+
 pub const SERVER_PORT: u16 = 3543;
 pub const BUFFER_SIZE: usize = 65535;
 pub const DEFAULT_TIMEOUT: u64 = 30;
 pub const VERSION: &str = "v0.1";
 pub const SALT_AND_IV_SIZE: u8 = 16;
+pub const STANDARD_RETRY_MAX: usize = 10;
 
 #[repr(u8)]
 pub enum ServerMethods {
@@ -11,15 +14,48 @@ pub enum ServerMethods {
     GET = 2,
     HEARTBEAT = 3,
 }
-
+#[allow(non_camel_case_types)]
 pub enum ServerResponse {
-    // avoid 0 from empty buffers
-    OK = 1,
     GENERAL_ERROR = 255,
     ID_EXISTS = 254,
     ID_DOESNT_EXIST = 253, // both error since sometimes it is the problem that the id exist and somethimes problem is that is doesn't
 }
 
+#[allow(non_camel_case_types)]
+#[derive(Debug)]
+pub enum ServerErrorResponses {
+    GENERAL_ERROR(String),
+    ID_EXISTS,
+    ID_DOESNT_EXIST,
+    IO(std::io::Error), // IO errors wraper
+}
+
+impl fmt::Display for ServerErrorResponses {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ServerErrorResponses::GENERAL_ERROR(msg) => write!(f, "General error: {}", msg),
+            ServerErrorResponses::ID_EXISTS => write!(f, "ID is already registered"),
+            ServerErrorResponses::ID_DOESNT_EXIST => write!(f, "ID isn't yet registered"),
+            ServerErrorResponses::IO(err) => write!(f, "IO error: {}", err),
+        }
+    }
+}
+impl std::error::Error for ServerErrorResponses {}
+impl From<std::io::Error> for ServerErrorResponses {
+    fn from(error: std::io::Error) -> Self {
+        ServerErrorResponses::IO(error)
+    }
+}
+impl ServerErrorResponses {
+    pub fn into_io_error(self) -> std::io::Error {
+        match self {
+            ServerErrorResponses::IO(io_err) => io_err,
+            other => std::io::Error::new(std::io::ErrorKind::Other, other),
+        }
+    }
+}
+
+#[allow(non_camel_case_types)]
 pub enum RegisterRequestDataPositions {
     ENCRYPTED = 1, // this feeld should be 0 if not encrypted
     ID_LEN = 2,

src/server/net.rs

@@ -120,7 +120,7 @@ pub async fn handle_request(
                 salt,
                 iv,
             ));
-            match socket.send_to(&[ServerResponse::OK as u8], src) {
+            match socket.send_to(&[ServerMethods::REGISTER as u8], src) {
                 Ok(s) => {
                     #[cfg(debug_assertions)]
                     eprintln!("send {} bytes", s);

src/server/types.rs

@@ -6,13 +6,16 @@ pub struct Client {
     #[readonly]
     pub client_sock_addr: Vec<u8>,
     pub last_heart_beat: i64,
+    #[readonly]
+    pub iv: [u8; SALT_AND_IV_SIZE as usize],
 }
 
 impl Client {
-    pub fn new(client_addr: Vec<u8>, heart_beat: i64) -> Self {
+    pub fn new(client_addr: Vec<u8>, heart_beat: i64, iv: [u8; SALT_AND_IV_SIZE as usize]) -> Self {
         Client {
             client_sock_addr: client_addr,
             last_heart_beat: heart_beat,
+            iv,
         }
     }
 }
@@ -30,8 +33,6 @@ pub struct Registration {
     pub encrypted: bool,
     #[readonly]
     pub salt: [u8; SALT_AND_IV_SIZE as usize],
-    #[readonly]
-    pub iv: [u8; SALT_AND_IV_SIZE as usize],
 }
 
 impl Registration {
@@ -45,11 +46,14 @@ impl Registration {
     ) -> Self {
         Registration {
             net_id,
-            clients: vec![Client::new(client_addr, heart_beat)],
+            clients: vec![Client::new(
+                client_addr,
+                heart_beat,
+                iv.unwrap_or([0; SALT_AND_IV_SIZE as usize]),
+            )],
             encrypted,
             last_heart_beat: heart_beat,
             salt: salt.unwrap_or([0; SALT_AND_IV_SIZE as usize]),
-            iv: iv.unwrap_or([0; SALT_AND_IV_SIZE as usize]),
         }
     }
 }

src/shared/crypto.rs

@@ -18,13 +18,9 @@ pub fn derive_key_from_password(password: &[u8], salt: &[u8]) -> [u8; 32] {
 }
 
 /// Encrypt using AES-256-CBC
-pub fn encrypt(
+pub fn encrypt(key: &[u8], iv: &[u8], data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
-    key: &[u8],
-    iv: &[u8],
-    plaintext: &[u8],
-) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
     let cipher = Aes256CbcEnc::new_from_slices(key, iv)?;
-    Ok(cipher.encrypt_padded_vec_mut::<Pkcs7>(plaintext))
+    Ok(cipher.encrypt_padded_vec_mut::<Pkcs7>(data))
 }
 
 /// Decrypt using AES-256-CBC

src/shared/mod.rs

@@ -1 +1 @@
-mod crypto;
+pub mod crypto;