Merge pull request #11 from NordSecurity/xray_bug_fixes

XRay bug fixes

Author: mathiaspeters

.gitignore

@@ -6,3 +6,4 @@
 **/*.pcap
 **/*.sk
 **/*.pk
+**/__pycache__

xray/README.md

@@ -18,6 +18,29 @@ Execution consists of these steps:
 5. The wireguard interface is destroyed
 6. The `.csv` file and pcap are analyzed
 
+The structure and exeuction of the XRay binary can be represented with the following image:
+
+```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.
+
 ## Running it
 
 X-Ray currently only works on linux
@@ -48,8 +71,6 @@ The application is executed with the `run.py` script. I takes some arguments, al
 
 ## Known issues
 
--  The analysis of pcaps is quite limited right now because it doesn't decrypt the packets (this is being worked on)
+- The analysis of pcaps is quite limited right now because it doesn't decrypt the packets (this is being worked on)
 
 - There are multiple inefficiencies that could potentially impact the test results, the main one being not reusing buffers when creating and sending packets. Each packet that gets constructed allocates a new buffer when they could all reuse the same one
-
-- Sometimes the wireguard handshake times out. When that happens, just rerun and it should be fine. Having to rerun multiple times is rare, but it happens

xray/analyze.py

@@ -22,7 +22,8 @@ def __init__(self, csv_path):
             reader = csv.reader(csvfile, delimiter=",", quotechar="|")
             next(reader)
             for row in reader:
-                self.indices.append(int(row[0]))
+                if len(row[0]) > 0:
+                    self.indices.append(int(row[0]))
                 send_ts = int(row[1])
                 if len(row[2]) > 0:
                     recv_ts = int(row[2])
@@ -147,7 +148,12 @@ def packet_latency(self, ax):
         ax.hist(buckets, color="blue", bins=num_buckets)
 
     def dropped_packets(self, ax):
-        num_buckets = 100
+        if self.count >= 100:
+            num_buckets = 100
+        elif self.count >= 10:
+            num_buckets = 10
+        else:
+            num_buckets = self.count
         bucket_size = int(self.count / (num_buckets - 1))
         buckets = []
         for iter, index in enumerate(self.csv_data.indices):

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,6 +39,26 @@ impl Client {
 
     pub async fn do_handshake(&mut self, wg_addr: SocketAddrV4) -> XRayResult<()> {
         println!("Handshake: starting");
+        for _ in 0..HANDSHAKE_MAX_TRIES {
+            match self.try_handshake(wg_addr).await {
+                Ok(_) => {
+                    println!("Handshake: done");
+                    return Ok(());
+                }
+                Err(XRayError::HandshakeTimedOut) => {
+                    println!("Handshake timed out")
+                }
+                Err(err) => {
+                    println!("Handshake failed");
+                    return Err(err);
+                }
+            }
+        }
+        println!("Could not establish handshake");
+        Err(XRayError::HandshakeTimedOut)
+    }
+
+    async fn try_handshake(&mut self, wg_addr: SocketAddrV4) -> XRayResult<()> {
         let tunn = self
             .tunn
             .as_mut()
@@ -46,18 +72,17 @@ impl Client {
             }
         }
         let mut handshake_buf = vec![0; 512];
-        let handshake_timeout = tokio::time::sleep(tokio::time::Duration::from_secs(3));
-        tokio::pin!(handshake_timeout);
+        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 {
-            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 => {
-                    return Err(XRayError::HandshakeTimedOut);
+            if let Ok(recv_type) = self.recv_encrypted(buf).await {
+                if matches!(recv_type, RecvType::HandshakeResponse) {
+                    return Ok(());
                 }
             }
         }

xray/src/event_loop.rs

@@ -42,7 +42,7 @@ impl EventLoop {
             is_done: false,
             crypto_buf: vec![0; 1024],
             plaintext_buf: vec![0; 1024],
-            recv_counter: 1,
+            recv_counter: 0,
         }
     }
 
@@ -55,7 +55,7 @@ impl EventLoop {
         loop {
             tokio::select! {
                 _ = &mut finish_timeout, if self.is_done => {
-                    self.on_finished(self.packets.len()).await?;
+                    self.on_finished(self.recv_counter).await?;
                     break;
                 },
                 _ = wg_tick_interval.tick() => {
@@ -150,8 +150,9 @@ impl EventLoop {
             RecvType::Data { length: bytes_read } => {
                 if bytes_read == Packet::send_size() {
                     if self.recv_counter % (self.cli_args.packet_count / 10) == 0 {
-                        println!("[Crypto] Received {} packets", self.recv_counter);
+                        println!("[Crypto] Received {} packets", self.recv_counter + 1);
                     }
+                    self.recv_counter += 1;
                     let send_index = u64::from_le_bytes(
                         self.crypto_buf[0..8]
                             .try_into()
@@ -162,7 +163,6 @@ impl EventLoop {
                         .as_micros();
                     self.packets[send_index].recv_index = Some(self.recv_counter as u64);
                     self.packets[send_index].recv_ts = Some(recv_ts);
-                    self.recv_counter += 1;
                     self.on_maybe_recv_all(finish_timeout);
                 }
             }
@@ -178,8 +178,9 @@ impl EventLoop {
         if let RecvType::Data { length: bytes_read } = rt {
             if bytes_read == Packet::send_size() {
                 if self.recv_counter % (self.cli_args.packet_count / 10) == 0 {
-                    println!("[Plaintext] Received {} packets", self.recv_counter);
+                    println!("[Plaintext] Received {} packets", self.recv_counter + 1);
                 }
+                self.recv_counter += 1;
                 let send_index = u64::from_le_bytes(
                     self.plaintext_buf[0..8]
                         .try_into()
@@ -190,19 +191,18 @@ impl EventLoop {
                     .as_micros();
                 self.packets[send_index].recv_index = Some(self.recv_counter as u64);
                 self.packets[send_index].recv_ts = Some(recv_ts);
-                self.recv_counter += 1;
                 self.on_maybe_recv_all(finish_timeout);
             }
         }
         Ok(())
     }
 
     fn on_maybe_recv_all(&self, finish_timeout: &mut Pin<&mut tokio::time::Sleep>) {
-        if self.recv_counter > self.cli_args.packet_count {
-            println!("All packets were received. Waiting 3 seconds to make sure pcap is properly populated");
+        if self.recv_counter >= self.cli_args.packet_count {
+            println!("All packets were received. Finishing test...");
             finish_timeout
                 .as_mut()
-                .reset(Instant::now() + Duration::from_secs(3));
+                .reset(Instant::now() + Duration::from_secs(1));
         }
     }