1use crate::agent::AgentHandle;
7use crate::error::AgentError as AuthsAgentError;
8use log::{debug, error, warn};
9use ssh_agent_lib::agent::Session;
10use ssh_agent_lib::error::AgentError as SSHAgentError;
11use ssh_agent_lib::proto::{AddIdentity, Credential, Identity, RemoveIdentity, SignRequest};
12use ssh_key::private::KeypairData;
13use ssh_key::public::{Ed25519PublicKey, KeyData};
14use ssh_key::{Algorithm, Signature};
15use std::convert::TryInto;
16use std::io;
17use std::sync::Arc;
18use zeroize::Zeroizing;
19
20#[derive(Clone)]
25pub struct AgentSession {
26 handle: Arc<AgentHandle>,
28}
29
30impl AgentSession {
31 pub fn new(handle: Arc<AgentHandle>) -> Self {
33 Self { handle }
34 }
35
36 pub fn handle(&self) -> &AgentHandle {
38 &self.handle
39 }
40}
41
42impl std::fmt::Debug for AgentSession {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 f.debug_struct("AgentSession")
45 .field("socket_path", self.handle.socket_path())
46 .field("is_running", &self.handle.is_running())
47 .finish()
48 }
49}
50
51pub use auths_crypto::build_ed25519_pkcs8_v2;
56
57#[ssh_agent_lib::async_trait]
58impl Session for AgentSession {
59 async fn request_identities(&mut self) -> Result<Vec<Identity>, SSHAgentError> {
60 let core = self.handle.lock().map_err(|_| {
61 error!("AgentSession failed to lock agent core (mutex poisoned).");
62 SSHAgentError::Failure
63 })?;
64
65 let pubkey_byte_vectors = core.public_keys();
66 debug!(
67 "request_identities: Agent core has {} keys.",
68 pubkey_byte_vectors.len()
69 );
70
71 let identities = pubkey_byte_vectors
72 .into_iter()
73 .filter_map(|pubkey_bytes| {
74 if pubkey_bytes.len() == 32 {
75 match <&[u8] as TryInto<&[u8; 32]>>::try_into(pubkey_bytes.as_slice()) {
76 Ok(key_bytes_array) => {
77 let ed_pubkey = Ed25519PublicKey(*key_bytes_array);
78 let key_data = KeyData::Ed25519(ed_pubkey);
79 let comment =
80 format!("auths-key-{}", hex::encode(&pubkey_bytes[..4]));
81 debug!("Adding identity with comment: {}", comment);
82 Some(Identity {
83 pubkey: key_data,
84 comment,
85 })
86 }
87 Err(_) => {
88 warn!("request_identities: Key is 32 bytes but failed TryInto<&[u8; 32]>. Skipping.");
89 None
90 }
91 }
92 } else {
93 warn!(
94 "request_identities: Found key with unexpected length ({}) in agent core. Skipping.",
95 pubkey_bytes.len()
96 );
97 None
98 }
99 })
100 .collect();
101
102 Ok(identities)
103 }
104
105 async fn sign(&mut self, request: SignRequest) -> Result<Signature, SSHAgentError> {
106 debug!(
107 "Handling sign request for key type: {:?}",
108 request.pubkey.algorithm()
109 );
110
111 let pubkey_bytes_to_sign_with = match &request.pubkey {
112 KeyData::Ed25519(key) => key.as_ref().to_vec(),
113 other_key_type => {
114 let err_msg = format!(
115 "Unsupported key type requested for signing: {:?}",
116 other_key_type.algorithm()
117 );
118 error!("{}", err_msg);
119 return Err(SSHAgentError::other(io::Error::new(
120 io::ErrorKind::Unsupported,
121 err_msg,
122 )));
123 }
124 };
125
126 let core = self.handle.lock().map_err(|_| {
127 error!("AgentSession failed to lock agent core (mutex poisoned).");
128 SSHAgentError::Failure
129 })?;
130
131 match core.sign(&pubkey_bytes_to_sign_with, &request.data) {
132 Ok(signature_bytes) => {
133 debug!("Successfully signed data using agent core.");
134 Signature::new(Algorithm::Ed25519, signature_bytes).map_err(|e| {
135 let err_msg = format!(
136 "Internal error: Failed to create ssh_key::Signature from core signature: {}",
137 e
138 );
139 error!("{}", err_msg);
140 SSHAgentError::other(io::Error::new(io::ErrorKind::InvalidData, err_msg))
141 })
142 }
143 Err(AuthsAgentError::KeyNotFound) => {
144 warn!("Sign request failed: Key not found in agent core.");
145 Err(SSHAgentError::Failure)
146 }
147 Err(other_core_error) => {
148 let err_msg = format!("Agent core signing error: {}", other_core_error);
149 error!("{}", err_msg);
150 Err(SSHAgentError::other(io::Error::other(err_msg)))
151 }
152 }
153 }
154
155 async fn add_identity(&mut self, identity: AddIdentity) -> Result<(), SSHAgentError> {
156 debug!("Handling add_identity request");
157
158 let (seed, pubkey) = match &identity.credential {
159 Credential::Key { privkey, .. } => match privkey {
160 KeypairData::Ed25519(kp) => (kp.private.to_bytes(), kp.public.0),
161 other => {
162 let err_msg = format!(
163 "Unsupported key type for add_identity: {:?}",
164 other.algorithm()
165 );
166 error!("{}", err_msg);
167 return Err(SSHAgentError::other(io::Error::new(
168 io::ErrorKind::Unsupported,
169 err_msg,
170 )));
171 }
172 },
173 Credential::Cert { .. } => {
174 error!("Certificate credentials are not supported for add_identity");
175 return Err(SSHAgentError::other(io::Error::new(
176 io::ErrorKind::Unsupported,
177 "Certificate credentials are not supported",
178 )));
179 }
180 };
181
182 let pkcs8_bytes = build_ed25519_pkcs8_v2(&seed, &pubkey);
183 self.handle
184 .register_key(Zeroizing::new(pkcs8_bytes))
185 .map_err(|e| {
186 let err_msg = format!("Failed to register key in agent: {}", e);
187 error!("{}", err_msg);
188 SSHAgentError::other(io::Error::other(err_msg))
189 })?;
190
191 debug!("Successfully added identity to agent");
192 Ok(())
193 }
194
195 async fn remove_identity(&mut self, identity: RemoveIdentity) -> Result<(), SSHAgentError> {
196 debug!("Handling remove_identity request");
197
198 let pubkey_bytes = match &identity.pubkey {
199 KeyData::Ed25519(key) => key.as_ref().to_vec(),
200 other => {
201 let err_msg = format!(
202 "Unsupported key type for remove_identity: {:?}",
203 other.algorithm()
204 );
205 error!("{}", err_msg);
206 return Err(SSHAgentError::other(io::Error::new(
207 io::ErrorKind::Unsupported,
208 err_msg,
209 )));
210 }
211 };
212
213 let mut core = self.handle.lock().map_err(|_| {
214 error!("AgentSession failed to lock agent core (mutex poisoned).");
215 SSHAgentError::Failure
216 })?;
217
218 core.unregister_key(&pubkey_bytes).map_err(|e| {
219 let err_msg = format!("Failed to remove key from agent: {}", e);
220 error!("{}", err_msg);
221 SSHAgentError::other(io::Error::new(io::ErrorKind::NotFound, err_msg))
222 })?;
223
224 debug!("Successfully removed identity from agent");
225 Ok(())
226 }
227
228 async fn remove_all_identities(&mut self) -> Result<(), SSHAgentError> {
229 debug!("Handling remove_all_identities request");
230
231 let mut core = self.handle.lock().map_err(|_| {
232 error!("AgentSession failed to lock agent core (mutex poisoned).");
233 SSHAgentError::Failure
234 })?;
235
236 core.clear_keys();
237 debug!("Successfully removed all identities from agent");
238 Ok(())
239 }
240}
241
242#[cfg(test)]
243mod tests {
244 use super::*;
245 use ring::rand::SystemRandom;
246 use ring::signature::Ed25519KeyPair;
247 use ssh_key::private::Ed25519Keypair as SshEd25519Keypair;
248 use std::path::PathBuf;
249 use zeroize::Zeroizing;
250
251 fn generate_test_pkcs8() -> Vec<u8> {
252 let rng = SystemRandom::new();
253 let pkcs8_doc = Ed25519KeyPair::generate_pkcs8(&rng).expect("Failed to generate PKCS#8");
254 pkcs8_doc.as_ref().to_vec()
255 }
256
257 #[test]
258 fn test_agent_session_new() {
259 let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test.sock")));
260 let session = AgentSession::new(handle.clone());
261
262 assert_eq!(
263 session.handle().socket_path(),
264 &PathBuf::from("/tmp/test.sock")
265 );
266 }
267
268 #[test]
269 fn test_agent_session_clone_shares_handle() {
270 let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test.sock")));
271
272 let pkcs8_bytes = generate_test_pkcs8();
273 handle
274 .register_key(Zeroizing::new(pkcs8_bytes))
275 .expect("Failed to register key");
276
277 let session1 = AgentSession::new(handle.clone());
278 let session2 = session1.clone();
279
280 assert_eq!(session1.handle().key_count().unwrap(), 1);
282 assert_eq!(session2.handle().key_count().unwrap(), 1);
283 }
284
285 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
286 async fn test_add_identity_round_trip() {
287 let seed: [u8; 32] = {
289 let pkcs8 = generate_test_pkcs8();
290 let mut s = [0u8; 32];
292 s.copy_from_slice(&pkcs8[16..48]);
293 s
294 };
295
296 let ssh_keypair = SshEd25519Keypair::from_seed(&seed);
297 let pubkey_bytes = ssh_keypair.public.0;
298
299 let identity = AddIdentity {
301 credential: Credential::Key {
302 privkey: KeypairData::Ed25519(ssh_keypair),
303 comment: "test-key".to_string(),
304 },
305 };
306
307 let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test-add.sock")));
309 let mut session = AgentSession::new(handle.clone());
310
311 assert_eq!(handle.key_count().unwrap(), 0);
313
314 session.add_identity(identity).await.unwrap();
316
317 assert_eq!(handle.key_count().unwrap(), 1);
319
320 let sign_request = SignRequest {
322 pubkey: KeyData::Ed25519(Ed25519PublicKey(pubkey_bytes)),
323 data: b"test data for signing".to_vec(),
324 flags: 0,
325 };
326
327 let signature = session.sign(sign_request).await.unwrap();
328 assert_eq!(signature.algorithm(), Algorithm::Ed25519);
329 assert!(!signature.as_bytes().is_empty());
330
331 let ring_pubkey =
333 ring::signature::UnparsedPublicKey::new(&ring::signature::ED25519, &pubkey_bytes);
334 ring_pubkey
335 .verify(b"test data for signing", signature.as_bytes())
336 .expect("Signature verification failed");
337 }
338
339 #[tokio::test]
340 async fn test_remove_identity() {
341 let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test-rm.sock")));
342 let mut session = AgentSession::new(handle.clone());
343
344 let seed: [u8; 32] = {
346 let pkcs8 = generate_test_pkcs8();
347 let mut s = [0u8; 32];
348 s.copy_from_slice(&pkcs8[16..48]);
349 s
350 };
351 let ssh_keypair = SshEd25519Keypair::from_seed(&seed);
352 let pubkey_bytes = ssh_keypair.public.0;
353
354 let identity = AddIdentity {
355 credential: Credential::Key {
356 privkey: KeypairData::Ed25519(ssh_keypair),
357 comment: "test-key".to_string(),
358 },
359 };
360 session.add_identity(identity).await.unwrap();
361 assert_eq!(handle.key_count().unwrap(), 1);
362
363 let remove = RemoveIdentity {
365 pubkey: KeyData::Ed25519(Ed25519PublicKey(pubkey_bytes)),
366 };
367 session.remove_identity(remove).await.unwrap();
368 assert_eq!(handle.key_count().unwrap(), 0);
369 }
370
371 #[tokio::test]
372 async fn test_remove_all_identities() {
373 let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test-rmall.sock")));
374 let mut session = AgentSession::new(handle.clone());
375
376 for _ in 0..2 {
378 let seed: [u8; 32] = {
379 let pkcs8 = generate_test_pkcs8();
380 let mut s = [0u8; 32];
381 s.copy_from_slice(&pkcs8[16..48]);
382 s
383 };
384 let ssh_keypair = SshEd25519Keypair::from_seed(&seed);
385 let identity = AddIdentity {
386 credential: Credential::Key {
387 privkey: KeypairData::Ed25519(ssh_keypair),
388 comment: "test-key".to_string(),
389 },
390 };
391 session.add_identity(identity).await.unwrap();
392 }
393 assert_eq!(handle.key_count().unwrap(), 2);
394
395 session.remove_all_identities().await.unwrap();
397 assert_eq!(handle.key_count().unwrap(), 0);
398 }
399}