Use raw sockets for TLS, to reduce CPU overhead (we don't need to decrypt the TLS stream for a download speedtest)

Author: 007

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "cf_speedtest"
-version = "0.5.5"
+version = "0.6.1"
 authors = ["https://github.com/12932"]
 repository = "https://github.com/12932/cf_speedtest"
 keywords = ["cross-platform", "speedtest", "cli", "speed", "test"]
@@ -24,6 +24,7 @@ rustls = { version = "0.23.33", default-features = false, features = ["ring"] }
 webpki-roots = "1.0.3"
 ctrlc = "3.5.1"
 phf = { version = "0.13", features = ["macros"] }
+socket2 = "0.6.1"
 
 [profile.release]
 debug = false

Cargo.toml

@@ -1,4 +1,3 @@
-use std::io::Read;
 use std::io::{self, Write};
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
@@ -16,6 +15,8 @@ use args::UserArgs;
 mod agent;
 use crate::agent::create_configured_agent;
 
+mod raw_socket;
+
 mod locations;
 mod table;
 #[cfg(test)]
@@ -259,61 +260,73 @@ fn upload_test(
     }
 }
 
-// download some bytes from cloudflare
+// download some bytes from cloudflare using raw encrypted byte reading
 fn download_test(
     bytes_to_request: usize,
     total_bytes_counter: &Arc<AtomicUsize>,
     current_down_speed: &Arc<AtomicUsize>,
     exit_signal: &Arc<AtomicBool>,
 ) -> Result<()> {
-    let agent: Agent = create_configured_agent();
-
-    let resp = match agent
-        .get(format!("{CLOUDFLARE_SPEEDTEST_DOWNLOAD_URL}&bytes={bytes_to_request}").as_str())
-        .call()
-    {
-        Ok(resp) => resp,
-        Err(err) => {
-            if !CTRL_C_PRESSED.load(Ordering::Relaxed) {
-                eprintln!("Error in download thread: {err}");
-            }
-            return Ok(());
-        }
-    };
-
-    let body = resp.into_body();
-    let mut resp_reader = body.into_reader();
-    let mut total_bytes_sank: usize = 0;
-
+    // Keep making new requests until exit_signal is set
     loop {
         // exit if we have passed deadline
         if exit_signal.load(Ordering::Relaxed) {
             return Ok(());
         }
 
-        // if we are fast, take big chunks
-        // if we are slow, take small chunks
-        let current_recv_buff =
-            get_appropriate_buff_size(current_down_speed.load(Ordering::Relaxed));
+        // Establish connection, perform TLS handshake, send HTTP request
+        let mut conn = match raw_socket::RawDownloadConnection::connect(
+            CLOUDFLARE_SPEEDTEST_DOWNLOAD_URL,
+            bytes_to_request,
+        ) {
+            Ok(conn) => conn,
+            Err(err) => {
+                if !CTRL_C_PRESSED.load(Ordering::Relaxed) {
+                    eprintln!("Error in download thread: {err}");
+                }
+                return Ok(());
+            }
+        };
+
+        let mut total_bytes_sank: usize = 0;
 
-        // copy bytes into the void
-        let bytes_sank = std::io::copy(
-            &mut resp_reader.by_ref().take(current_recv_buff),
-            &mut std::io::sink(),
-        )? as usize;
+        // Read from this connection until it's exhausted
+        loop {
+            // exit if we have passed deadline
+            if exit_signal.load(Ordering::Relaxed) {
+                return Ok(());
+            }
 
-        //println!("Thread {:?} sank {} bytes", std::thread::current().id(), bytes_sank);
+            // if we are fast, take big chunks
+            // if we are slow, take small chunks
+            let current_recv_buff =
+                get_appropriate_buff_size(current_down_speed.load(Ordering::Relaxed)) as usize;
+
+            // Read raw encrypted bytes directly from socket (no TLS decryption!)
+            let mut buf = vec![0u8; current_recv_buff];
+            let bytes_read = match conn.read_encrypted_bytes(&mut buf) {
+                Ok(n) => n,
+                Err(err) => {
+                    if !CTRL_C_PRESSED.load(Ordering::Relaxed) {
+                        eprintln!("Error reading from socket: {err}");
+                    }
+                    // Connection error, break to create a new connection
+                    break;
+                }
+            };
 
-        if bytes_sank == 0 {
-            if total_bytes_sank == 0 {
-                eprintln!("Cloudflare sent an empty response?");
+            if bytes_read == 0 {
+                if total_bytes_sank == 0 {
+                    eprintln!("Cloudflare sent an empty response?");
+                }
+                // Connection exhausted, break inner loop to make a new request
+                break;
             }
 
-            return Ok(());
+            // Count the encrypted bytes we received (wire bytes including TLS overhead)
+            total_bytes_sank += bytes_read;
+            total_bytes_counter.fetch_add(bytes_read, Ordering::SeqCst);
         }
-
-        total_bytes_sank += bytes_sank;
-        total_bytes_counter.fetch_add(bytes_sank, Ordering::SeqCst);
     }
 }