diff options
Diffstat (limited to 'crates')
3 files changed, 38 insertions, 19 deletions
diff --git a/crates/utils/tcp_connection/src/instance.rs b/crates/utils/tcp_connection/src/instance.rs index be7c956..2ef3480 100644 --- a/crates/utils/tcp_connection/src/instance.rs +++ b/crates/utils/tcp_connection/src/instance.rs @@ -1,6 +1,6 @@ use std::{path::Path, time::Duration}; -use rand::{Rng, TryRngCore}; +use rand::TryRngCore; use rsa::{ RsaPrivateKey, RsaPublicKey, pkcs1::{DecodeRsaPrivateKey, DecodeRsaPublicKey}, @@ -12,7 +12,6 @@ use tokio::{ io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter}, net::TcpStream, }; -use uuid::Uuid; use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey}; use ring::rand::SystemRandom; @@ -232,7 +231,6 @@ impl ConnectionInstance { // Open file and get metadata let mut file = File::open(path).await?; let file_size = file.metadata().await?.len(); - // Allow empty files - just send the header with size 0 // Send file header (version + size + crc) self.stream.write_all(&1u64.to_be_bytes()).await?; @@ -443,6 +441,7 @@ impl ConnectionInstance { // Send challenge to target self.stream.write_all(&challenge).await?; + self.stream.flush().await?; // Read signature from target let mut signature = Vec::new(); @@ -453,13 +452,18 @@ impl ConnectionInstance { signature.resize(signature_len, 0); self.stream.read_exact(&mut signature).await?; - // Read UUID from target to identify which public key to use - let mut uuid_buf = [0u8; 16]; - self.stream.read_exact(&mut uuid_buf).await?; - let user_uuid = Uuid::from_bytes(uuid_buf); + // Read key identifier from target to identify which public key to use + let mut key_id_len_buf = [0u8; 4]; + self.stream.read_exact(&mut key_id_len_buf).await?; + let key_id_len = u32::from_be_bytes(key_id_len_buf) as usize; + + let mut key_id_buf = vec![0u8; key_id_len]; + self.stream.read_exact(&mut key_id_buf).await?; + let key_id = String::from_utf8(key_id_buf) + .map_err(|e| TcpTargetError::Crypto(format!("Invalid key identifier: {}", e)))?; // Load appropriate public key - let public_key_path = public_key_dir.as_ref().join(format!("{}.pub", user_uuid)); + let public_key_path = public_key_dir.as_ref().join(format!("{}.pem", key_id)); if !public_key_path.exists() { return Ok(false); } @@ -518,17 +522,22 @@ impl ConnectionInstance { // Send signature length and signature let signature_len = signature.len() as u32; self.stream.write_all(&signature_len.to_be_bytes()).await?; + self.stream.flush().await?; self.stream.write_all(&signature).await?; + self.stream.flush().await?; - // Send UUID for public key identification - self.stream.write_all(verify_public_key.as_bytes()).await?; + // Send key identifier for public key identification + let key_id_bytes = verify_public_key.as_bytes(); + let key_id_len = key_id_bytes.len() as u32; + self.stream.write_all(&key_id_len.to_be_bytes()).await?; + self.stream.flush().await?; + self.stream.write_all(key_id_bytes).await?; + self.stream.flush().await?; Ok(true) } } -// Helper functions for Ed25519 support - /// Parse Ed25519 public key from PEM format fn parse_ed25519_public_key(pem: &str) -> [u8; 32] { // Robust parsing for Ed25519 public key using pem crate @@ -558,8 +567,6 @@ fn parse_ed25519_private_key(pem: &str) -> Result<SigningKey, TcpTargetError> { )) } -// Helper functions for DSA support - /// Parse DSA public key information from PEM fn parse_dsa_public_key( pem: &str, @@ -609,7 +616,7 @@ fn verify_dsa_signature( public_key.verify(message, signature).is_ok() } -/// Sign with DSA (simplified - in practice this would use proper private key operations) +/// Sign with DSA fn sign_with_dsa( algorithm_and_key: &(&'static dyn signature::VerificationAlgorithm, Vec<u8>), message: &[u8], diff --git a/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs b/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs index 57d3819..723bf52 100644 --- a/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs +++ b/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs @@ -6,7 +6,10 @@ use tcp_connection::{ target::TcpServerTarget, target_configure::ServerTargetConfig, }; -use tokio::{join, time::sleep}; +use tokio::{ + join, + time::{sleep, timeout}, +}; pub(crate) struct ExampleChallengeClientHandle; @@ -38,7 +41,7 @@ impl ServerHandle<ExampleChallengeClientHandle> for ExampleChallengeServerHandle #[tokio::test] async fn test_connection_with_challenge_handle() -> Result<(), std::io::Error> { - let host = "localhost"; + let host = "localhost:5011"; // Server setup let Ok(server_target) = TcpServerTarget::< @@ -76,7 +79,16 @@ async fn test_connection_with_challenge_handle() -> Result<(), std::io::Error> { let _ = client_target.connect().await; }; - let _ = async { join!(future_client, future_server) }.await; + let test_timeout = Duration::from_secs(10); + + timeout(test_timeout, async { join!(future_client, future_server) }) + .await + .map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::TimedOut, + format!("Test timed out after {:?}", test_timeout), + ) + })?; Ok(()) } diff --git a/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs b/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs index d9fa123..f35fd80 100644 --- a/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs +++ b/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs @@ -49,7 +49,7 @@ impl ServerHandle<ExampleClientHandle> for ExampleServerHandle { #[tokio::test] async fn test_connection_with_example_handle() { - let host = "localhost"; + let host = "localhost:5012"; // Server setup let Ok(server_target) = |