From 2e08a3093e5dff1f7e04a55019bace73c8296237 Mon Sep 17 00:00:00 2001
From: PoliEcho <PoliEcho@pupes.org>
Date: Fri, 25 Jul 2025 22:59:26 +0200
Subject: [PATCH] some rebuilding

---
 src/client/main.rs | 27 +++++++++++++++++++++++----
 src/lib.rs         |  7 ++++---
 src/server/net.rs  | 19 +++++++++++++++++++
 3 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/src/client/main.rs b/src/client/main.rs
index 4785f22..4d04bfb 100644
--- a/src/client/main.rs
+++ b/src/client/main.rs
@@ -1,9 +1,10 @@
 use pea_2_pea::SERVER_PORT;
 
 use std::{
-    io::{Read, Write},
+    io::{ErrorKind, Read, Write},
     net::UdpSocket,
     process::exit,
+    time::Duration,
 };
 
 #[derive(clap::Parser)]
@@ -47,6 +48,8 @@ fn main() -> std::io::Result<()> {
         })()
         .expect("Failed to bind to any available port");
 
+        socket.set_read_timeout(Some(Duration::new(10, 0)))?; // set timeout to 10 seconds
+
         // send query request to get server public key
         let server_port: u16 = (|| -> u16 {
             match cli.bind_port {
@@ -59,18 +62,34 @@ fn main() -> std::io::Result<()> {
             .parse()
             .unwrap();
 
+        {
+            let mut query_byte: [u8; 1] = [0; 1];
+            query_byte[0] = pea_2_pea::ServerMethods::QUERY as u8;
+            match socket.send_to(&query_byte, &server_SocketAddr) {
+                Ok(s) => {
+                    #[cfg(debug_assertions)]
+                    eprintln!("send {} bytes", s);
+                }
+                Err(e) => {
+                    eprintln!("Error snding data: {}", e);
+                }
+            }
+        }
+
         let mut buf: [u8; pea_2_pea::BUFFER_SIZE] = [0; pea_2_pea::BUFFER_SIZE];
         loop {
             match socket.recv_from(&mut buf) {
                 Ok((data_length, src)) => {}
+                Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::TimedOut => {
+                    // timedout
+                    continue;
+                }
                 Err(e) => {
                     eprintln!("Error receiving data: {}", e);
                     std::process::exit(-4);
                 }
             }
-            let mut out = std::io::stdout();
-            out.write_all(&buf)?;
-            out.flush()?;
+            break;
         }
     }
     Ok(())
diff --git a/src/lib.rs b/src/lib.rs
index 35f709c..d5c8b7a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,9 +6,10 @@ pub const RSA_SIZE: usize = 2048;
 
 #[repr(u8)]
 pub enum ServerMethods {
-    REGISTER = 0,
-    GET = 1,
-    HEARTBEAT = 2,
+    QUERY = 0,
+    REGISTER = 1,
+    GET = 2,
+    HEARTBEAT = 3,
 }
 
 pub mod shared;
diff --git a/src/server/net.rs b/src/server/net.rs
index 81b0064..e88a8b1 100644
--- a/src/server/net.rs
+++ b/src/server/net.rs
@@ -7,6 +7,25 @@ pub async fn handle_request(
     let mut rng: rand::prelude::ThreadRng = rand::thread_rng();
 
     match buf[0] {
+        x if x == pea_2_pea::ServerMethods::QUERY as u8 => {
+            #[cfg(debug_assertions)]
+            eprintln!("QUERY method");
+
+            let client_sock_addr_str: String = src.to_string();
+            let mut send_vec: Vec<u8> = client_sock_addr_str.into();
+            send_vec.insert(0, pea_2_pea::ServerMethods::QUERY as u8);
+
+            match socket.send_to(&send_vec, &src) {
+                Ok(s) => {
+                    #[cfg(debug_assertions)]
+                    eprintln!("send {} bytes", s);
+                }
+                Err(e) => {
+                    eprintln!("Error snding data: {}", e);
+                }
+            }
+        }
+
         x if x == pea_2_pea::ServerMethods::GET as u8 => {
             #[cfg(debug_assertions)]
             println!("GET method");