Bump rust edition to 2024 and refactor code

Author: TeddyHuang-00

Cargo.lock

@@ -1,7 +1,7 @@
 [package]
 name = "sshping"
 version = "0.2.2"
-edition = "2021"
+edition = "2024"
 description = "SSH-based ping that measures interactive character echo latency and file transfer throughput. Pronounced \"shipping\"."
 authors = ["Nan Huang <teddyhuangnan@gmail.com>"]
 license = "MIT"
@@ -28,12 +28,12 @@ rand = "0.9.0"
 regex = "1.11.1"
 russh = "0.54.6"
 russh-sftp = "2.1.1"
-ssh2-config = "0.6.0"
 serde = { version = "1.0.217", features = ["derive"] }
 serde_json = "1.0.138"
 shellexpand = "3.1.0"
 simple_logger = "5.0.0"
 size = "0.5.0"
+ssh2-config = "0.6.0"
 tabled = "0.20.0"
 tokio = { version = "1.42.0", features = ["full"] }
 whoami = "1.5.1"

Cargo.toml

@@ -5,8 +5,57 @@ use std::{
 };
 
 use log::{info, warn};
-use russh::client;
-use russh::keys::{decode_secret_key, PrivateKeyWithHashAlg};
+use russh::{
+    client,
+    keys::{decode_secret_key, PrivateKeyWithHashAlg},
+};
+
+async fn authenticate_publickey<H: client::Handler>(
+    session: &mut client::Handle<H>,
+    user: &str,
+    identity: &PathBuf,
+    timeout: f64,
+) -> Result<(), String> {
+    let key_content = std::fs::read_to_string(identity)
+        .map_err(|e| format!("Failed to read identity file: {e}"))?;
+    let key = decode_secret_key(&key_content, None)
+        .map_err(|e| format!("Failed to decode secret key: {e}"))?;
+    let timeout_result = tokio::time::timeout(
+        Duration::from_secs_f64(timeout),
+        session.authenticate_publickey(user, PrivateKeyWithHashAlg::new(Arc::new(key), None)),
+    )
+    .await
+    .map_err(|_| format!("Public key authentication timed out after {timeout} seconds"))?;
+    let auth_result =
+        timeout_result.map_err(|e| format!("Public key authentication failed: {e}"))?;
+    if !auth_result.success() {
+        return Err("Public key authentication returned false".to_string());
+    }
+
+    info!("Public key authentication succeeded");
+    Ok(())
+}
+
+async fn authenticate_password<H: client::Handler>(
+    session: &mut client::Handle<H>,
+    user: &str,
+    password: &str,
+    timeout: f64,
+) -> Result<(), String> {
+    let timeout_result = tokio::time::timeout(
+        Duration::from_secs_f64(timeout),
+        session.authenticate_password(user, password),
+    )
+    .await
+    .map_err(|_| format!("Password authentication timed out after {timeout} seconds"))?;
+    let auth_result = timeout_result.map_err(|e| format!("Password authentication failed: {e}"))?;
+    if !auth_result.success() {
+        return Err("Password authentication returned false".to_string());
+    }
+
+    info!("Password authentication succeeded");
+    Ok(())
+}
 
 pub async fn authenticate_all<H: client::Handler>(
     session: &mut client::Handle<H>,
@@ -18,58 +67,23 @@ pub async fn authenticate_all<H: client::Handler>(
     let start = Instant::now();
 
     // Try public key authentication if identity file is provided
-    if let Some(identity_path) = identity {
-        match std::fs::read_to_string(identity_path) {
-            Ok(key_content) => {
-                match decode_secret_key(&key_content, password) {
-                    Ok(key) => {
-                        match tokio::time::timeout(
-                            Duration::from_secs_f64(timeout),
-                            session.authenticate_publickey(
-                                user,
-                                PrivateKeyWithHashAlg::new(Arc::new(key), None),
-                            ),
-                        )
-                        .await
-                        {
-                            Ok(Ok(auth_result)) => {
-                                if auth_result.success() {
-                                    info!("Public key authentication succeeded");
-                                    return Ok(start.elapsed());
-                                } else {
-                                    warn!("Public key authentication returned false");
-                                }
-                            }
-                            Ok(Err(e)) => warn!("Public key authentication failed: {e}"),
-                            Err(_) => warn!("Public key authentication timed out"),
-                        }
-                    }
-                    Err(e) => warn!("Failed to decode secret key: {e}"),
-                }
-            }
-            Err(e) => warn!("Failed to read identity file: {e}"),
-        }
+    if let Some(identity_path) = identity
+        && authenticate_publickey(session, user, identity_path, timeout)
+            .await
+            .inspect_err(|e| warn!("{e}"))
+            .is_ok()
+    {
+        return Ok(start.elapsed());
     }
 
     // Try password authentication
-    if let Some(pwd) = password {
-        match tokio::time::timeout(
-            Duration::from_secs_f64(timeout),
-            session.authenticate_password(user, pwd),
-        )
-        .await
-        {
-            Ok(Ok(auth_result)) => {
-                if auth_result.success() {
-                    info!("Password authentication succeeded");
-                    return Ok(start.elapsed());
-                } else {
-                    warn!("Password authentication returned false");
-                }
-            }
-            Ok(Err(e)) => warn!("Password authentication failed: {e}"),
-            Err(_) => warn!("Password authentication timed out"),
-        }
+    if let Some(pwd) = password
+        && authenticate_password(session, user, pwd, timeout)
+            .await
+            .inspect_err(|e| warn!("{e}"))
+            .is_ok()
+    {
+        return Ok(start.elapsed());
     }
 
     // Fails if all authentication methods fail

src/auth.rs

@@ -99,7 +99,7 @@ async fn main() -> ExitCode {
         ..Default::default()
     });
     let handler = SshHandler;
-    
+
     let addr = (opts.target.host.as_str(), opts.target.port);
     let mut session = match tokio::time::timeout(
         std::time::Duration::from_secs_f64(opts.ssh_timeout),

src/main.rs

@@ -9,8 +9,7 @@ use rand::{
     distr::{Distribution, Uniform},
     rng,
 };
-use russh::client;
-use russh::ChannelMsg;
+use russh::{client, ChannelMsg};
 use russh_sftp::client::SftpSession;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
@@ -45,33 +44,33 @@ pub async fn run_echo_test<H: client::Handler>(
     debug!("Running echo test with command: {echo_cmd:?}");
     debug!("Number of characters to echo: {char_count:?}");
     debug!("Time limit for echo: {time_limit:?} seconds");
-    
+
     // Start the channel server
     trace!("Preparing channel session");
     let mut channel = session
         .channel_open_session()
         .await
         .map_err(|e| e.to_string())?;
-    
+
     // Request a pseudo-terminal for the interactive shell
     channel
         .request_pty(true, "sshping", 10, 5, 0, 0, &[])
         .await
         .map_err(|e| e.to_string())?;
-    
+
     channel
         .request_shell(false)
         .await
         .map_err(|e| e.to_string())?;
-    
+
     // Send the echo command to accept input
     trace!("Starting echo command");
     let echo_cmd_bytes = format!("{echo_cmd}\n").into_bytes();
     channel
         .data(&echo_cmd_bytes[..])
         .await
         .map_err(|e| e.to_string())?;
-    
+
     // Read the initial buffer to clear the echo command
     tokio::time::sleep(Duration::from_millis(100)).await;
     while let Some(msg) = channel.wait().await {
@@ -93,14 +92,11 @@ pub async fn run_echo_test<H: client::Handler>(
 
     for (n, idx) in (0..char_count).zip((0..write_buffer.len()).cycle()) {
         let start = Instant::now();
-        
+
         // Send one character
         let byte_slice = &write_buffer[idx..idx + 1];
-        channel
-            .data(byte_slice)
-            .await
-            .map_err(|e| e.to_string())?;
-        
+        channel.data(byte_slice).await.map_err(|e| e.to_string())?;
+
         // Wait for echo back
         loop {
             if let Some(msg) = channel.wait().await {
@@ -117,14 +113,14 @@ pub async fn run_echo_test<H: client::Handler>(
                 }
             }
         }
-        
+
         let latency = start.elapsed().as_nanos();
         latencies.push(latency);
-        
-        if let Some(timeout) = timeout {
-            if start_time.elapsed() > timeout {
-                break;
-            }
+
+        if let Some(timeout) = timeout
+            && start_time.elapsed() > timeout
+        {
+            break;
         }
         progress_bar.set_position((n as u64) + 1);
     }
@@ -189,7 +185,7 @@ async fn run_upload_test<H: client::Handler>(
     formatter: &Formatter,
 ) -> Result<SpeedTestResult, String> {
     info!("Running upload speed test");
-    
+
     // Establish SFTP channel
     trace!("Establishing SFTP channel");
     let channel = session
@@ -203,7 +199,7 @@ async fn run_upload_test<H: client::Handler>(
     let sftp = SftpSession::new(channel.into_stream())
         .await
         .map_err(|e| e.to_string())?;
-    
+
     // Generate random data to upload
     trace!("Generating random data");
     let dist = Uniform::try_from(0..128_u8).unwrap();
@@ -212,14 +208,11 @@ async fn run_upload_test<H: client::Handler>(
         .take(size as usize)
         .map(|v| (v & 0x3f) + 32)
         .collect();
-    
+
     // Open remote file for writing
     let remote_path = remote_file.to_str().ok_or("Invalid remote file path")?;
-    let mut file = sftp
-        .create(remote_path)
-        .await
-        .map_err(|e| e.to_string())?;
-    
+    let mut file = sftp.create(remote_path).await.map_err(|e| e.to_string())?;
+
     // Preparing logging variables
     let mut total_bytes_sent = 0;
     let start_time: Instant = Instant::now();
@@ -234,7 +227,7 @@ async fn run_upload_test<H: client::Handler>(
         progress_bar.set_position(total_bytes_sent as u64);
     }
     progress_bar.finish_and_clear();
-    
+
     // Close the file
     file.shutdown().await.map_err(|e| e.to_string())?;
 
@@ -254,7 +247,7 @@ async fn run_download_test<H: client::Handler>(
     formatter: &Formatter,
 ) -> Result<SpeedTestResult, String> {
     info!("Running download speed test");
-    
+
     // Establish SFTP channel
     trace!("Establishing SFTP channel");
     let channel = session
@@ -268,25 +261,22 @@ async fn run_download_test<H: client::Handler>(
     let sftp = SftpSession::new(channel.into_stream())
         .await
         .map_err(|e| e.to_string())?;
-    
+
     // Get file size
     let remote_path = remote_file.to_str().ok_or("Invalid remote file path")?;
     let metadata = sftp
         .metadata(remote_path)
         .await
         .map_err(|e| e.to_string())?;
     let size = metadata.len();
-    
+
     if size == 0 {
         return Err("Remote file is empty".to_string());
     }
-    
+
     // Open remote file for reading
-    let mut file = sftp
-        .open(remote_path)
-        .await
-        .map_err(|e| e.to_string())?;
-    
+    let mut file = sftp.open(remote_path).await.map_err(|e| e.to_string())?;
+
     // Prepare buffer for downloading
     trace!("Preparing buffer for downloading");
     let mut buffer = vec![0; chunk_size as usize];
@@ -299,13 +289,17 @@ async fn run_download_test<H: client::Handler>(
     // Starting downloading file
     trace!("Receiving file in chunks");
     while size - total_bytes_recv > chunk_size {
-        file.read_exact(&mut buffer).await.map_err(|e| e.to_string())?;
+        file.read_exact(&mut buffer)
+            .await
+            .map_err(|e| e.to_string())?;
         total_bytes_recv += chunk_size;
         progress_bar.set_position(total_bytes_recv);
     }
     if size - total_bytes_recv > 0 {
         let mut remaining = vec![0; (size - total_bytes_recv) as usize];
-        file.read_exact(&mut remaining).await.map_err(|e| e.to_string())?;
+        file.read_exact(&mut remaining)
+            .await
+            .map_err(|e| e.to_string())?;
         total_bytes_recv += remaining.len() as u64;
         progress_bar.set_position(total_bytes_recv);
     }