diff options
| -rw-r--r-- | crates/utils/tcp_connection/src/instance.rs | 488 | ||||
| -rw-r--r-- | crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs | 16 |
2 files changed, 296 insertions, 208 deletions
diff --git a/crates/utils/tcp_connection/src/instance.rs b/crates/utils/tcp_connection/src/instance.rs index db36b54..be7c956 100644 --- a/crates/utils/tcp_connection/src/instance.rs +++ b/crates/utils/tcp_connection/src/instance.rs @@ -1,7 +1,6 @@ use std::{path::Path, time::Duration}; -use base64::{engine::general_purpose::STANDARD, prelude::*}; -use rand::Rng; +use rand::{Rng, TryRngCore}; use rsa::{ RsaPrivateKey, RsaPublicKey, pkcs1::{DecodeRsaPrivateKey, DecodeRsaPublicKey}, @@ -16,46 +15,77 @@ use tokio::{ use uuid::Uuid; use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey}; +use ring::rand::SystemRandom; use ring::signature::{ - self, ECDSA_P256_SHA256_ASN1, ECDSA_P384_SHA384_ASN1, RSA_PKCS1_2048_8192_SHA256, + self, ECDSA_P256_SHA256_ASN1, ECDSA_P384_SHA384_ASN1, EcdsaKeyPair, RSA_PKCS1_2048_8192_SHA256, UnparsedPublicKey, }; use crate::error::TcpTargetError; -const CHUNK_SIZE: usize = 8 * 1024; +const DEFAULT_CHUNK_SIZE: usize = 4096; +const DEFAULT_TIMEOUT_SECS: u64 = 10; -pub struct ConnectionInstance { - stream: TcpStream, +const ECDSA_P256_SHA256_ASN1_SIGNING: &signature::EcdsaSigningAlgorithm = + &signature::ECDSA_P256_SHA256_ASN1_SIGNING; +const ECDSA_P384_SHA384_ASN1_SIGNING: &signature::EcdsaSigningAlgorithm = + &signature::ECDSA_P384_SHA384_ASN1_SIGNING; + +#[derive(Debug, Clone)] +pub struct ConnectionConfig { + pub chunk_size: usize, + pub timeout_secs: u64, + pub enable_crc_validation: bool, } -impl From<TcpStream> for ConnectionInstance { - fn from(value: TcpStream) -> Self { - Self { stream: value } +impl Default for ConnectionConfig { + fn default() -> Self { + Self { + chunk_size: DEFAULT_CHUNK_SIZE, + timeout_secs: DEFAULT_TIMEOUT_SECS, + enable_crc_validation: false, + } } } -// Helper trait for reading u64 from TcpStream -trait ReadU64Ext { - async fn read_u64(&mut self) -> Result<u64, std::io::Error>; +pub struct ConnectionInstance { + stream: TcpStream, + config: ConnectionConfig, } -impl ReadU64Ext for TcpStream { - async fn read_u64(&mut self) -> Result<u64, std::io::Error> { - let mut buf = [0u8; 8]; - self.read_exact(&mut buf).await?; - Ok(u64::from_be_bytes(buf)) +impl From<TcpStream> for ConnectionInstance { + fn from(stream: TcpStream) -> Self { + Self { + stream, + config: ConnectionConfig::default(), + } } } impl ConnectionInstance { + /// Create a new ConnectionInstance with custom configuration + pub fn with_config(stream: TcpStream, config: ConnectionConfig) -> Self { + Self { stream, config } + } + + /// Get a reference to the current configuration + pub fn config(&self) -> &ConnectionConfig { + &self.config + } + + /// Get a mutable reference to the current configuration + pub fn config_mut(&mut self) -> &mut ConnectionConfig { + &mut self.config + } /// Serialize data and write to the target machine pub async fn write<Data>(&mut self, data: Data) -> Result<(), TcpTargetError> where Data: Default + Serialize, { let Ok(json_text) = serde_json::to_string(&data) else { - return Err(TcpTargetError::from("Serialize failed.")); + return Err(TcpTargetError::Serialization( + "Serialize failed.".to_string(), + )); }; Self::write_text(self, json_text).await?; Ok(()) @@ -67,10 +97,12 @@ impl ConnectionInstance { Data: Default + serde::de::DeserializeOwned, { let Ok(json_text) = Self::read_text(self, buffer_size).await else { - return Err(TcpTargetError::from("Read failed.")); + return Err(TcpTargetError::Io("Read failed.".to_string())); }; let Ok(deser_obj) = serde_json::from_str::<Data>(&json_text) else { - return Err(TcpTargetError::from("Deserialize failed.")); + return Err(TcpTargetError::Serialization( + "Deserialize failed.".to_string(), + )); }; Ok(deser_obj) } @@ -81,7 +113,9 @@ impl ConnectionInstance { Data: Default + Serialize, { let Ok(json_text) = serde_json::to_string(&data) else { - return Err(TcpTargetError::from("Serialize failed.")); + return Err(TcpTargetError::Serialization( + "Serialize failed.".to_string(), + )); }; Self::write_large_text(self, json_text).await?; Ok(()) @@ -96,10 +130,12 @@ impl ConnectionInstance { Data: Default + serde::de::DeserializeOwned, { let Ok(json_text) = Self::read_large_text(self, buffer_size).await else { - return Err(TcpTargetError::from("Read failed.")); + return Err(TcpTargetError::Io("Read failed.".to_string())); }; let Ok(deser_obj) = serde_json::from_str::<Data>(&json_text) else { - return Err(TcpTargetError::from("Deserialize failed.")); + return Err(TcpTargetError::Serialization( + "Deserialize failed.".to_string(), + )); }; Ok(deser_obj) } @@ -111,7 +147,7 @@ impl ConnectionInstance { // Write match self.stream.write_all(text.as_bytes()).await { Ok(_) => Ok(()), - Err(err) => Err(TcpTargetError::from(err.to_string())), + Err(err) => Err(TcpTargetError::Io(err.to_string())), } } @@ -128,7 +164,7 @@ impl ConnectionInstance { let text = String::from_utf8_lossy(&buffer[..n]).to_string(); Ok(text) } - Err(err) => Err(TcpTargetError::from(err.to_string())), + Err(err) => Err(TcpTargetError::Io(err.to_string())), } } @@ -145,7 +181,7 @@ impl ConnectionInstance { let chunk = &bytes[offset..]; let written = match self.stream.write(chunk).await { Ok(n) => n, - Err(err) => return Err(TcpTargetError::from(err.to_string())), + Err(err) => return Err(TcpTargetError::Io(err.to_string())), }; offset += written; } @@ -168,7 +204,7 @@ impl ConnectionInstance { Ok(n) => { buffer.extend_from_slice(&chunk_buf[..n]); } - Err(err) => return Err(TcpTargetError::from(err.to_string())), + Err(err) => return Err(TcpTargetError::Io(err.to_string())), } } @@ -181,59 +217,66 @@ impl ConnectionInstance { // Validate file if !path.exists() { - return Err(TcpTargetError::from(format!( + return Err(TcpTargetError::File(format!( "File not found: {}", path.display() ))); } if path.is_dir() { - return Err(TcpTargetError::from(format!( + return Err(TcpTargetError::File(format!( "Path is directory: {}", path.display() ))); } // Open file and get metadata - let mut file = File::open(path) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - let file_size = file - .metadata() - .await - .map_err(|e| TcpTargetError::from(e.to_string()))? - .len(); - if file_size == 0 { - return Err(TcpTargetError::from("Cannot send empty file")); - } + 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?; + self.stream.write_all(&file_size.to_be_bytes()).await?; + + // Calculate and send CRC32 if enabled + let file_crc = if self.config.enable_crc_validation { + let crc32 = crc::Crc::<u32>::new(&crc::CRC_32_ISO_HDLC); + let mut crc_calculator = crc32.digest(); + + let mut temp_reader = + BufReader::with_capacity(self.config.chunk_size, File::open(path).await?); + let mut temp_buffer = vec![0u8; self.config.chunk_size]; + let mut temp_bytes_read = 0; + + while temp_bytes_read < file_size { + let bytes_to_read = + (file_size - temp_bytes_read).min(self.config.chunk_size as u64) as usize; + temp_reader + .read_exact(&mut temp_buffer[..bytes_to_read]) + .await?; + crc_calculator.update(&temp_buffer[..bytes_to_read]); + temp_bytes_read += bytes_to_read as u64; + } + + crc_calculator.finalize() + } else { + 0 + }; - // Send file header (version + size) - self.stream - .write_all(&1u64.to_be_bytes()) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - self.stream - .write_all(&file_size.to_be_bytes()) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.write_all(&file_crc.to_be_bytes()).await?; // Transfer file content - let mut reader = BufReader::with_capacity(CHUNK_SIZE, &mut file); + let mut reader = BufReader::with_capacity(self.config.chunk_size, &mut file); let mut bytes_sent = 0; while bytes_sent < file_size { - let buffer = reader - .fill_buf() - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + let buffer = reader.fill_buf().await?; if buffer.is_empty() { break; } let chunk_size = buffer.len().min((file_size - bytes_sent) as usize); - self.stream - .write_all(&buffer[..chunk_size]) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.write_all(&buffer[..chunk_size]).await?; reader.consume(chunk_size); bytes_sent += chunk_size as u64; @@ -241,26 +284,27 @@ impl ConnectionInstance { // Verify transfer completion if bytes_sent != file_size { - return Err(TcpTargetError::from(format!( + return Err(TcpTargetError::File(format!( "Transfer incomplete: expected {} bytes, sent {} bytes", file_size, bytes_sent ))); } - self.stream - .flush() - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.flush().await?; // Wait for receiver confirmation let mut ack = [0u8; 1]; - tokio::time::timeout(Duration::from_secs(10), self.stream.read_exact(&mut ack)) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))? - .map_err(|e| TcpTargetError::from(e.to_string()))?; + tokio::time::timeout( + Duration::from_secs(self.config.timeout_secs), + self.stream.read_exact(&mut ack), + ) + .await + .map_err(|_| TcpTargetError::Timeout("Ack timeout".to_string()))??; if ack[0] != 1 { - return Err(TcpTargetError::from("Receiver verification failed")); + return Err(TcpTargetError::Protocol( + "Receiver verification failed".to_string(), + )); } Ok(()) @@ -269,35 +313,42 @@ impl ConnectionInstance { /// Read file from target machine pub async fn read_file(&mut self, save_path: impl AsRef<Path>) -> Result<(), TcpTargetError> { let path = save_path.as_ref(); + // Create CRC instance at function scope to ensure proper lifetime + let crc_instance = crc::Crc::<u32>::new(&crc::CRC_32_ISO_HDLC); // Make sure parent directory exists if let Some(parent) = path.parent() { if !parent.exists() { - tokio::fs::create_dir_all(parent) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + tokio::fs::create_dir_all(parent).await?; } } - // Read file header (version + size) + // Read file header (version + size + crc) let mut version_buf = [0u8; 8]; - self.stream - .read_exact(&mut version_buf) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.read_exact(&mut version_buf).await?; let version = u64::from_be_bytes(version_buf); if version != 1 { - return Err(TcpTargetError::from("Unsupported transfer version")); + return Err(TcpTargetError::Protocol( + "Unsupported transfer version".to_string(), + )); } let mut size_buf = [0u8; 8]; - self.stream - .read_exact(&mut size_buf) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.read_exact(&mut size_buf).await?; let file_size = u64::from_be_bytes(size_buf); + + let mut expected_crc_buf = [0u8; 4]; + self.stream.read_exact(&mut expected_crc_buf).await?; + let expected_crc = u32::from_be_bytes(expected_crc_buf); if file_size == 0 { - return Err(TcpTargetError::from("Cannot receive zero-length file")); + // Create empty file and return early + let _file = OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(path) + .await?; + return Ok(()); } // Prepare output file @@ -306,57 +357,74 @@ impl ConnectionInstance { .create(true) .truncate(true) .open(path) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - let mut writer = BufWriter::with_capacity(CHUNK_SIZE, file); + .await?; + let mut writer = BufWriter::with_capacity(self.config.chunk_size, file); - // Receive file content - let mut buffer = vec![0u8; CHUNK_SIZE]; + // Receive file content with CRC calculation if enabled let mut bytes_received = 0; + let mut buffer = vec![0u8; self.config.chunk_size]; + let mut crc_calculator = if self.config.enable_crc_validation { + Some(crc_instance.digest()) + } else { + None + }; while bytes_received < file_size { - let read_size = buffer.len().min((file_size - bytes_received) as usize); - self.stream - .read_exact(&mut buffer[..read_size]) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - - writer - .write_all(&buffer[..read_size]) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - bytes_received += read_size as u64; + let bytes_to_read = + (file_size - bytes_received).min(self.config.chunk_size as u64) as usize; + let chunk = &mut buffer[..bytes_to_read]; + + self.stream.read_exact(chunk).await?; + + writer.write_all(chunk).await?; + + // Update CRC if validation is enabled + if let Some(ref mut crc) = crc_calculator { + crc.update(chunk); + } + + bytes_received += bytes_to_read as u64; + } + + // Verify transfer completion + if bytes_received != file_size { + return Err(TcpTargetError::File(format!( + "Transfer incomplete: expected {} bytes, received {} bytes", + file_size, bytes_received + ))); + } + + writer.flush().await?; + + // Validate CRC if enabled + if self.config.enable_crc_validation { + if let Some(crc_calculator) = crc_calculator { + let actual_crc = crc_calculator.finalize(); + if actual_crc != expected_crc && expected_crc != 0 { + return Err(TcpTargetError::File(format!( + "CRC validation failed: expected {:08x}, got {:08x}", + expected_crc, actual_crc + ))); + } + } } // Final flush and sync - writer - .flush() - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - writer - .into_inner() - .sync_all() - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + writer.flush().await?; + writer.into_inner().sync_all().await?; // Verify completion if bytes_received != file_size { let _ = tokio::fs::remove_file(path).await; - return Err(TcpTargetError::from(format!( + return Err(TcpTargetError::File(format!( "Transfer incomplete: expected {} bytes, received {} bytes", file_size, bytes_received ))); } // Send confirmation - self.stream - .write_all(&[1]) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - self.stream - .flush() - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.write_all(&[1u8]).await?; + self.stream.flush().await?; Ok(()) } @@ -366,36 +434,28 @@ impl ConnectionInstance { public_key_dir: impl AsRef<Path>, ) -> Result<bool, TcpTargetError> { // Generate random challenge - let mut rng = rand::rng(); - let challenge: [u8; 32] = rng.random(); + let mut challenge = [0u8; 32]; + rand::rngs::OsRng + .try_fill_bytes(&mut challenge) + .map_err(|e| { + TcpTargetError::Crypto(format!("Failed to generate random challenge: {}", e)) + })?; // Send challenge to target - self.stream - .write_all(&challenge) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.write_all(&challenge).await?; // Read signature from target let mut signature = Vec::new(); let mut signature_len_buf = [0u8; 4]; - self.stream - .read_exact(&mut signature_len_buf) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.read_exact(&mut signature_len_buf).await?; let signature_len = u32::from_be_bytes(signature_len_buf) as usize; signature.resize(signature_len, 0); - self.stream - .read_exact(&mut signature) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + 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 - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.read_exact(&mut uuid_buf).await?; let user_uuid = Uuid::from_bytes(uuid_buf); // Load appropriate public key @@ -404,9 +464,7 @@ impl ConnectionInstance { return Ok(false); } - let public_key_pem = tokio::fs::read_to_string(&public_key_path) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + let public_key_pem = tokio::fs::read_to_string(&public_key_path).await?; // Try to verify with different key types let verified = if let Ok(rsa_key) = RsaPublicKey::from_pkcs1_pem(&public_key_pem) { @@ -415,9 +473,15 @@ impl ConnectionInstance { } else if let Ok(ed25519_key) = VerifyingKey::from_bytes(&parse_ed25519_public_key(&public_key_pem)) { - let sig_bytes: [u8; 64] = signature.as_slice().try_into().unwrap_or([0u8; 64]); - let sig = Signature::from_bytes(&sig_bytes); - ed25519_key.verify(&challenge, &sig).is_ok() + if signature.len() == 64 { + let sig_bytes: [u8; 64] = signature.as_slice().try_into().map_err(|_| { + TcpTargetError::Crypto("Invalid signature length for Ed25519".to_string()) + })?; + let sig = Signature::from_bytes(&sig_bytes); + ed25519_key.verify(&challenge, &sig).is_ok() + } else { + false + } } else if let Ok(dsa_key_info) = parse_dsa_public_key(&public_key_pem) { verify_dsa_signature(&dsa_key_info, &challenge, &signature) } else { @@ -430,50 +494,34 @@ impl ConnectionInstance { pub async fn accept_challenge( &mut self, private_key_file: impl AsRef<Path>, - verify_user_uuid: Uuid, + verify_public_key: &str, ) -> Result<bool, TcpTargetError> { // Read challenge from initiator let mut challenge = [0u8; 32]; - self.stream - .read_exact(&mut challenge) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.read_exact(&mut challenge).await?; // Load private key - let private_key_pem = tokio::fs::read_to_string(&private_key_file) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + let private_key_pem = tokio::fs::read_to_string(&private_key_file).await?; // Sign the challenge with supported key types let signature = if let Ok(rsa_key) = RsaPrivateKey::from_pkcs1_pem(&private_key_pem) { let padding = rsa::pkcs1v15::Pkcs1v15Sign::new::<sha2::Sha256>(); - rsa_key - .sign(padding, &challenge) - .map_err(|e| TcpTargetError::from(e.to_string()))? + rsa_key.sign(padding, &challenge)? } else if let Ok(ed25519_key) = parse_ed25519_private_key(&private_key_pem) { ed25519_key.sign(&challenge).to_bytes().to_vec() } else if let Ok(dsa_key_info) = parse_dsa_private_key(&private_key_pem) { - sign_with_dsa(&dsa_key_info, &challenge) + sign_with_dsa(&dsa_key_info, &challenge)? } else { return Ok(false); }; // Send signature length and signature let signature_len = signature.len() as u32; - self.stream - .write_all(&signature_len.to_be_bytes()) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; - self.stream - .write_all(&signature) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.write_all(&signature_len.to_be_bytes()).await?; + self.stream.write_all(&signature).await?; // Send UUID for public key identification - self.stream - .write_all(verify_user_uuid.as_bytes()) - .await - .map_err(|e| TcpTargetError::from(e.to_string()))?; + self.stream.write_all(verify_public_key.as_bytes()).await?; Ok(true) } @@ -483,15 +531,13 @@ impl ConnectionInstance { /// Parse Ed25519 public key from PEM format fn parse_ed25519_public_key(pem: &str) -> [u8; 32] { - // Simple parsing for Ed25519 public key (assuming raw 32-byte key) - let lines: Vec<&str> = pem.lines().collect(); + // Robust parsing for Ed25519 public key using pem crate let mut key_bytes = [0u8; 32]; - if lines.len() >= 2 && lines[0].contains("PUBLIC KEY") { - if let Ok(decoded) = STANDARD.decode(lines[1].trim()) { - if decoded.len() >= 32 { - key_bytes.copy_from_slice(&decoded[decoded.len() - 32..]); - } + if let Ok(pem_data) = pem::parse(pem) { + if pem_data.tag() == "PUBLIC KEY" && pem_data.contents().len() >= 32 { + let contents = pem_data.contents(); + key_bytes.copy_from_slice(&contents[contents.len() - 32..]); } } key_bytes @@ -499,18 +545,17 @@ fn parse_ed25519_public_key(pem: &str) -> [u8; 32] { /// Parse Ed25519 private key from PEM format fn parse_ed25519_private_key(pem: &str) -> Result<SigningKey, TcpTargetError> { - let lines: Vec<&str> = pem.lines().collect(); - - if lines.len() >= 2 && lines[0].contains("PRIVATE KEY") { - if let Ok(decoded) = STANDARD.decode(lines[1].trim()) { - if decoded.len() >= 32 { - let mut seed = [0u8; 32]; - seed.copy_from_slice(&decoded[decoded.len() - 32..]); - return Ok(SigningKey::from_bytes(&seed)); - } + if let Ok(pem_data) = pem::parse(pem) { + if pem_data.tag() == "PRIVATE KEY" && pem_data.contents().len() >= 32 { + let contents = pem_data.contents(); + let mut seed = [0u8; 32]; + seed.copy_from_slice(&contents[contents.len() - 32..]); + return Ok(SigningKey::from_bytes(&seed)); } } - Err(TcpTargetError::from("Invalid Ed25519 private key format")) + Err(TcpTargetError::Crypto( + "Invalid Ed25519 private key format".to_string(), + )) } // Helper functions for DSA support @@ -519,23 +564,30 @@ fn parse_ed25519_private_key(pem: &str) -> Result<SigningKey, TcpTargetError> { fn parse_dsa_public_key( pem: &str, ) -> Result<(&'static dyn signature::VerificationAlgorithm, Vec<u8>), TcpTargetError> { - let lines: Vec<&str> = pem.lines().collect(); + if let Ok(pem_data) = pem::parse(pem) { + let contents = pem_data.contents().to_vec(); - if lines.len() >= 2 { - if let Ok(decoded) = STANDARD.decode(lines[1].trim()) { - // Try different DSA algorithms - if pem.contains("ECDSA") || pem.contains("ecdsa") { + // Try different DSA algorithms based on PEM tag + match pem_data.tag() { + "EC PUBLIC KEY" | "PUBLIC KEY" if pem.contains("ECDSA") || pem.contains("ecdsa") => { if pem.contains("P-256") { - return Ok((&ECDSA_P256_SHA256_ASN1, decoded)); + return Ok((&ECDSA_P256_SHA256_ASN1, contents)); } else if pem.contains("P-384") { - return Ok((&ECDSA_P384_SHA384_ASN1, decoded)); + return Ok((&ECDSA_P384_SHA384_ASN1, contents)); } } - // Default to RSA if no specific algorithm detected - return Ok((&RSA_PKCS1_2048_8192_SHA256, decoded)); + "RSA PUBLIC KEY" | "PUBLIC KEY" => { + return Ok((&RSA_PKCS1_2048_8192_SHA256, contents)); + } + _ => {} } + + // Default to RSA for unknown types + return Ok((&RSA_PKCS1_2048_8192_SHA256, contents)); } - Err(TcpTargetError::from("Invalid DSA public key format")) + Err(TcpTargetError::Crypto( + "Invalid DSA public key format".to_string(), + )) } /// Parse DSA private key information from PEM @@ -559,13 +611,51 @@ fn verify_dsa_signature( /// Sign with DSA (simplified - in practice this would use proper private key operations) fn sign_with_dsa( - _algorithm_and_key: &(&'static dyn signature::VerificationAlgorithm, Vec<u8>), + algorithm_and_key: &(&'static dyn signature::VerificationAlgorithm, Vec<u8>), message: &[u8], -) -> Vec<u8> { - // Note: This is a simplified implementation. In a real scenario, - // you would use proper private key signing operations with ring or other crypto library. - // For now, we'll return a dummy signature for demonstration. - let mut signature = vec![0u8; 64]; - signature[..32].copy_from_slice(&message[..32.min(message.len())]); - signature +) -> Result<Vec<u8>, TcpTargetError> { + let (algorithm, key_bytes) = algorithm_and_key; + + // Handle different DSA/ECDSA algorithms by comparing algorithm identifiers + // Since we can't directly compare trait objects, we use pointer comparison + let algorithm_ptr = algorithm as *const _ as *const (); + let ecdsa_p256_ptr = &ECDSA_P256_SHA256_ASN1 as *const _ as *const (); + let ecdsa_p384_ptr = &ECDSA_P384_SHA384_ASN1 as *const _ as *const (); + + if algorithm_ptr == ecdsa_p256_ptr { + let key_pair = EcdsaKeyPair::from_pkcs8( + ECDSA_P256_SHA256_ASN1_SIGNING, + &key_bytes, + &SystemRandom::new(), + ) + .map_err(|e| { + TcpTargetError::Crypto(format!("Failed to create ECDSA P-256 key pair: {}", e)) + })?; + + let signature = key_pair + .sign(&SystemRandom::new(), message) + .map_err(|e| TcpTargetError::Crypto(format!("ECDSA P-256 signing failed: {}", e)))?; + + Ok(signature.as_ref().to_vec()) + } else if algorithm_ptr == ecdsa_p384_ptr { + let key_pair = EcdsaKeyPair::from_pkcs8( + ECDSA_P384_SHA384_ASN1_SIGNING, + &key_bytes, + &SystemRandom::new(), + ) + .map_err(|e| { + TcpTargetError::Crypto(format!("Failed to create ECDSA P-384 key pair: {}", e)) + })?; + + let signature = key_pair + .sign(&SystemRandom::new(), message) + .map_err(|e| TcpTargetError::Crypto(format!("ECDSA P-384 signing failed: {}", e)))?; + + Ok(signature.as_ref().to_vec()) + } else { + // RSA or unsupported algorithm + Err(TcpTargetError::Unsupported( + "DSA/ECDSA signing not supported for this algorithm type".to_string(), + )) + } } 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 ac57451..57d3819 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 @@ -1,7 +1,4 @@ -use std::{ - env::{current_dir, set_current_dir}, - time::Duration, -}; +use std::{env::current_dir, time::Duration}; use tcp_connection::{ handle::{ClientHandle, ServerHandle}, @@ -18,7 +15,9 @@ impl ClientHandle<ExampleChallengeServerHandle> for ExampleChallengeClientHandle mut instance: ConnectionInstance, ) -> impl std::future::Future<Output = ()> + Send + Sync { async move { - // TODO :: Complete the implementation + let key = current_dir().unwrap().join("res").join("test_key"); + let result = instance.accept_challenge(key, "test_key").await.unwrap(); + assert_eq!(true, result); } } } @@ -30,7 +29,9 @@ impl ServerHandle<ExampleChallengeClientHandle> for ExampleChallengeServerHandle mut instance: ConnectionInstance, ) -> impl std::future::Future<Output = ()> + Send + Sync { async move { - // TODO :: Complete the implementation + let key_dir = current_dir().unwrap().join("res"); + let result = instance.challenge(key_dir).await.unwrap(); + assert_eq!(true, result); } } } @@ -39,9 +40,6 @@ impl ServerHandle<ExampleChallengeClientHandle> for ExampleChallengeServerHandle async fn test_connection_with_challenge_handle() -> Result<(), std::io::Error> { let host = "localhost"; - // Enter temp directory - set_current_dir(current_dir().unwrap().join(".temp/"))?; - // Server setup let Ok(server_target) = TcpServerTarget::< ExampleChallengeClientHandle, |