2 files changed, 11 insertions, 8 deletions
diff --git a/crates/utils/tcp_connection/src/instance_challenge.rs b/crates/utils/tcp_connection/src/instance_challenge.rs
index c1cf46f..c4ea6a8 100644
--- a/crates/utils/tcp_connection/src/instance_challenge.rs
+++ b/crates/utils/tcp_connection/src/instance_challenge.rs
@@ -35,13 +35,13 @@ impl ConnectionInstance {
     /// * `public_key_dir` - Directory containing public key files for verification
     ///
     /// # Returns
-    /// * `Ok(true)` - Challenge verification successful
-    /// * `Ok(false)` - Challenge verification failed
+    /// * `Ok((true, "KeyId"))` - Challenge verification successful
+    /// * `Ok((false, "KeyId"))` - Challenge verification failed
     /// * `Err(TcpTargetError)` - Error during challenge process
     pub async fn challenge(
         &mut self,
         public_key_dir: impl AsRef<Path>,
-    ) -> Result<bool, TcpTargetError> {
+    ) -> Result<(bool, String), TcpTargetError> {
         // Generate random challenge
         let mut challenge = [0u8; 32];
         rand::rngs::OsRng
@@ -76,7 +76,7 @@ impl ConnectionInstance {
         // Load appropriate public key
         let public_key_path = public_key_dir.as_ref().join(format!("{}.pem", key_id));
         if !public_key_path.exists() {
-            return Ok(false);
+            return Ok((false, key_id));
         }
 
         let public_key_pem = tokio::fs::read_to_string(&public_key_path).await?;
@@ -103,7 +103,7 @@ impl ConnectionInstance {
             false
         };
 
-        Ok(verified)
+        Ok((verified, key_id))
     }
 
     /// Accepts a challenge from the target machine to verify connection security
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 2fc1a87..9327b3e 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
@@ -69,8 +69,9 @@ impl ServerHandle<ExampleChallengeClientHandle> for ExampleChallengeServerHandle
     async fn process(mut instance: ConnectionInstance) {
         // Challenge with correct key
         let key_dir = current_dir().unwrap().join("res").join("key");
-        let result = instance.challenge(key_dir).await.unwrap();
+        let (result, key_id) = instance.challenge(key_dir).await.unwrap();
         assert!(result);
+        assert_eq!(key_id, "test_key");
 
         // Send response
         instance
@@ -80,8 +81,9 @@ impl ServerHandle<ExampleChallengeClientHandle> for ExampleChallengeServerHandle
 
         // Challenge again
         let key_dir = current_dir().unwrap().join("res").join("key");
-        let result = instance.challenge(key_dir).await.unwrap();
+        let (result, key_id) = instance.challenge(key_dir).await.unwrap();
         assert!(!result);
+        assert_eq!(key_id, "test_key");
 
         // Send response
         instance
@@ -91,8 +93,9 @@ impl ServerHandle<ExampleChallengeClientHandle> for ExampleChallengeServerHandle
 
         // Challenge again
         let key_dir = current_dir().unwrap().join("res").join("key");
-        let result = instance.challenge(key_dir).await.unwrap();
+        let (result, key_id) = instance.challenge(key_dir).await.unwrap();
         assert!(!result);
+        assert_eq!(key_id, "test_key__");
 
         // Send response
         instance