Skip to main content

auths_core/agent/
session.rs

1//! SSH agent session handler.
2//!
3//! This module provides `AgentSession`, which implements the `ssh_agent_lib::agent::Session`
4//! trait to handle SSH agent protocol requests.
5
6use 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/// Wraps an `AgentHandle` to implement the `ssh_agent_lib::agent::Session` trait.
21///
22/// Each `AgentSession` holds a reference to an `AgentHandle`, enabling multiple
23/// independent agent instances to coexist.
24#[derive(Clone)]
25pub struct AgentSession {
26    /// Reference to the agent handle
27    handle: Arc<AgentHandle>,
28}
29
30impl AgentSession {
31    /// Creates a new AgentSession wrapping the given AgentHandle.
32    pub fn new(handle: Arc<AgentHandle>) -> Self {
33        Self { handle }
34    }
35
36    /// Returns a reference to the underlying agent handle.
37    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
51/// Build a PKCS#8 v2 DER encoding for an Ed25519 key from seed and public key.
52///
53/// Ring's `Ed25519KeyPair::from_pkcs8()` requires v2 format (RFC 8410) which
54/// includes the public key. Produces the fixed 85-byte structure:
55pub 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        // Both sessions share the same handle
281        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        // Generate a test seed and create an SSH keypair from it
288        let seed: [u8; 32] = {
289            let pkcs8 = generate_test_pkcs8();
290            // Extract seed from the PKCS#8 bytes (bytes 16..48 in ring's format)
291            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        // Build an AddIdentity request (what the client sends over the wire)
300        let identity = AddIdentity {
301            credential: Credential::Key {
302                privkey: KeypairData::Ed25519(ssh_keypair),
303                comment: "test-key".to_string(),
304            },
305        };
306
307        // Create session with empty handle
308        let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test-add.sock")));
309        let mut session = AgentSession::new(handle.clone());
310
311        // Verify no keys initially
312        assert_eq!(handle.key_count().unwrap(), 0);
313
314        // Add the identity via the session (this is what was broken before)
315        session.add_identity(identity).await.unwrap();
316
317        // Verify the key is now registered
318        assert_eq!(handle.key_count().unwrap(), 1);
319
320        // Sign data via the session and verify the signature
321        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        // Verify the signature using ring
332        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        // Add a key via add_identity
345        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        // Remove it
364        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        // Add two keys
377        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        // Remove all
396        session.remove_all_identities().await.unwrap();
397        assert_eq!(handle.key_count().unwrap(), 0);
398    }
399}