From 53d7ca0d696fc7cee1952dc7453153b2ef8dc401 Mon Sep 17 00:00:00 2001
From: Mathias Peters <mathias.peters@nordsec.com>
Date: Tue, 17 Dec 2024 20:53:11 +0100
Subject: [PATCH] fixes

---
 xray/README.md     | 19 +++++++++++-
 xray/src/client.rs | 74 +++++++++++++++++++++++++++++-----------------
 2 files changed, 65 insertions(+), 28 deletions(-)

diff --git a/xray/README.md b/xray/README.md
index 2bb93ad..3a7ba99 100644
--- a/xray/README.md
+++ b/xray/README.md
@@ -20,7 +20,24 @@ Execution consists of these steps:
 
 The structure and exeuction of the XRay binary can be represented with the following image:
 
-![XRay diagram](xray_setup.png)
+```mermaid
+flowchart LR
+    subgraph Encryption boundary
+        subgraph Crypto socket
+            tunn[noise::Tunn]<-->cs[UdpSocket]
+        end
+        subgraph Wireguard interface
+            wg[NepTUN/WireguardGo/LinuxNative/etc]
+        end
+    end
+    subgraph Plaintext socket
+        ps[UdpSocket]
+    end
+    cp["Packet(index: u64)"]-->tunn
+    cs<--Encrypted packet-->wg
+    wg<--Plaintext packet-->ps[Plaintext Socket]
+    ps-->pp["Packet(index: u64)"]
+```
 
 The red arrows represent crypto packets and the blue arrows represent plaintext packets.
 
diff --git a/xray/src/client.rs b/xray/src/client.rs
index d6a942f..a471843 100644
--- a/xray/src/client.rs
+++ b/xray/src/client.rs
@@ -1,4 +1,7 @@
-use std::net::{SocketAddr, SocketAddrV4};
+use std::{
+    net::{SocketAddr, SocketAddrV4},
+    time::Duration,
+};
 
 use neptun::noise::{Tunn, TunnResult};
 use pnet::packet::{
@@ -14,6 +17,9 @@ use crate::{
     XRayError, XRayResult,
 };
 
+const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(2);
+const HANDSHAKE_MAX_TRIES: usize = 5;
+
 pub struct Client {
     pub addr: SocketAddrV4,
     pub tunn: Option<Tunn>,
@@ -33,34 +39,18 @@ impl Client {
 
     pub async fn do_handshake(&mut self, wg_addr: SocketAddrV4) -> XRayResult<()> {
         println!("Handshake: starting");
-        'new_handshake: for _ in 0..3 {
-            let tunn = self
-                .tunn
-                .as_mut()
-                .expect("This function should only be called on clients with a Tunn object");
-            match tunn.format_handshake_initiation(&mut self.buf, true) {
-                TunnResult::WriteToNetwork(packet) => {
-                    self.sock.send_to(packet, wg_addr).await?;
+        for _ in 0..HANDSHAKE_MAX_TRIES {
+            match self.try_handshake(wg_addr).await {
+                Ok(_) => {
+                    println!("Handshake: done");
+                    return Ok(());
                 }
-                unexpected => {
-                    return Err(XRayError::UnexpectedTunnResult(format!("{unexpected:?}")));
+                Err(XRayError::HandshakeTimedOut) => {
+                    println!("Handshake timed out")
                 }
-            }
-            let mut handshake_buf = vec![0; 512];
-            let handshake_timeout = tokio::time::sleep(tokio::time::Duration::from_secs(2));
-            tokio::pin!(handshake_timeout);
-            loop {
-                tokio::select! {
-                    Ok(recv_type) = self.recv_encrypted(&mut handshake_buf) => {
-                        if matches!(recv_type, RecvType::HandshakeResponse) {
-                            println!("Handshake: done");
-                            return Ok(());
-                        }
-                    }
-                    _ = &mut handshake_timeout => {
-                        println!("Handshake timed out trying again...");
-                        continue 'new_handshake;
-                    }
+                Err(err) => {
+                    println!("Handshake failed");
+                    return Err(err);
                 }
             }
         }
@@ -68,6 +58,36 @@ impl Client {
         Err(XRayError::HandshakeTimedOut)
     }
 
+    async fn try_handshake(&mut self, wg_addr: SocketAddrV4) -> XRayResult<()> {
+        let tunn = self
+            .tunn
+            .as_mut()
+            .expect("This function should only be called on clients with a Tunn object");
+        match tunn.format_handshake_initiation(&mut self.buf, true) {
+            TunnResult::WriteToNetwork(packet) => {
+                self.sock.send_to(packet, wg_addr).await?;
+            }
+            unexpected => {
+                return Err(XRayError::UnexpectedTunnResult(format!("{unexpected:?}")));
+            }
+        }
+        let mut handshake_buf = vec![0; 512];
+        tokio::time::timeout(HANDSHAKE_TIMEOUT, self.recv_handshake(&mut handshake_buf))
+            .await
+            .map_err(|_| XRayError::HandshakeTimedOut)??;
+        Ok(())
+    }
+
+    async fn recv_handshake(&mut self, buf: &mut [u8]) -> XRayResult<()> {
+        loop {
+            if let Ok(recv_type) = self.recv_encrypted(buf).await {
+                if matches!(recv_type, RecvType::HandshakeResponse) {
+                    return Ok(());
+                }
+            }
+        }
+    }
+
     pub async fn send_packet(
         &mut self,
         sock_dst: SocketAddrV4,