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                let key_data = match pubkey_bytes.len() {
75                    32 => {
76                        let key_bytes_array: &[u8; 32] = pubkey_bytes.as_slice().try_into().ok()?;
77                        KeyData::Ed25519(Ed25519PublicKey(*key_bytes_array))
78                    }
79                    33 | 65 => {
80                        // P-256: try to parse as ECDSA SEC1 point
81                        let ecdsa_pk =
82                            ssh_key::public::EcdsaPublicKey::from_sec1_bytes(&pubkey_bytes).ok()?;
83                        KeyData::Ecdsa(ecdsa_pk)
84                    }
85                    n => {
86                        warn!("request_identities: unsupported key length ({n}). Skipping.");
87                        return None;
88                    }
89                };
90                let comment = format!(
91                    "auths-key-{}",
92                    hex::encode(&pubkey_bytes[..4.min(pubkey_bytes.len())])
93                );
94                debug!("Adding identity with comment: {}", comment);
95                Some(Identity {
96                    pubkey: key_data,
97                    comment,
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, algorithm) = match &request.pubkey {
112            KeyData::Ed25519(key) => (key.as_ref().to_vec(), Algorithm::Ed25519),
113            KeyData::Ecdsa(ecdsa_key) => {
114                let point_bytes = ecdsa_key.as_ref();
115                (
116                    point_bytes.to_vec(),
117                    Algorithm::Ecdsa {
118                        curve: ssh_key::EcdsaCurve::NistP256,
119                    },
120                )
121            }
122            other_key_type => {
123                let err_msg = format!(
124                    "Unsupported key type requested for signing: {:?}",
125                    other_key_type.algorithm()
126                );
127                error!("{}", err_msg);
128                return Err(SSHAgentError::other(io::Error::new(
129                    io::ErrorKind::Unsupported,
130                    err_msg,
131                )));
132            }
133        };
134
135        let core = self.handle.lock().map_err(|_| {
136            error!("AgentSession failed to lock agent core (mutex poisoned).");
137            SSHAgentError::Failure
138        })?;
139
140        match core.sign(&pubkey_bytes_to_sign_with, &request.data) {
141            Ok(signature_bytes) => {
142                debug!("Successfully signed data using agent core.");
143                Signature::new(algorithm, signature_bytes).map_err(|e| {
144                    let err_msg = format!(
145                        "Internal error: Failed to create ssh_key::Signature from core signature: {}",
146                        e
147                    );
148                    error!("{}", err_msg);
149                    SSHAgentError::other(io::Error::new(io::ErrorKind::InvalidData, err_msg))
150                })
151            }
152            Err(AuthsAgentError::KeyNotFound) => {
153                warn!("Sign request failed: Key not found in agent core.");
154                Err(SSHAgentError::Failure)
155            }
156            Err(other_core_error) => {
157                let err_msg = format!("Agent core signing error: {}", other_core_error);
158                error!("{}", err_msg);
159                Err(SSHAgentError::other(io::Error::other(err_msg)))
160            }
161        }
162    }
163
164    async fn add_identity(&mut self, identity: AddIdentity) -> Result<(), SSHAgentError> {
165        debug!("Handling add_identity request");
166
167        let pkcs8_bytes: Vec<u8> = match &identity.credential {
168            Credential::Key { privkey, .. } => match privkey {
169                KeypairData::Ed25519(kp) => {
170                    let seed = kp.private.to_bytes();
171                    let pubkey = kp.public.0;
172                    build_ed25519_pkcs8_v2(&seed, &pubkey)
173                }
174                KeypairData::Ecdsa(ssh_key::private::EcdsaKeypair::NistP256 {
175                    private, ..
176                }) => {
177                    use auths_crypto::{TypedSeed, TypedSignerKey};
178                    let scalar_bytes = private.as_slice();
179                    if scalar_bytes.len() != 32 {
180                        let err_msg = format!(
181                            "Invalid P-256 scalar length: expected 32, got {}",
182                            scalar_bytes.len()
183                        );
184                        error!("{}", err_msg);
185                        return Err(SSHAgentError::other(io::Error::new(
186                            io::ErrorKind::InvalidData,
187                            err_msg,
188                        )));
189                    }
190                    #[allow(clippy::expect_used)] // INVARIANT: length checked above
191                    let mut scalar = [0u8; 32];
192                    scalar.copy_from_slice(scalar_bytes);
193                    let signer =
194                        TypedSignerKey::from_seed(TypedSeed::P256(scalar)).map_err(|e| {
195                            let err_msg = format!("P-256 signer construction failed: {}", e);
196                            error!("{}", err_msg);
197                            SSHAgentError::other(io::Error::other(err_msg))
198                        })?;
199                    let pkcs8 = signer.to_pkcs8().map_err(|e| {
200                        let err_msg = format!("P-256 PKCS8 encoding failed: {}", e);
201                        error!("{}", err_msg);
202                        SSHAgentError::other(io::Error::other(err_msg))
203                    })?;
204                    pkcs8.as_ref().to_vec()
205                }
206                other => {
207                    let err_msg = format!(
208                        "Unsupported key type for add_identity: {:?}",
209                        other.algorithm()
210                    );
211                    error!("{}", err_msg);
212                    return Err(SSHAgentError::other(io::Error::new(
213                        io::ErrorKind::Unsupported,
214                        err_msg,
215                    )));
216                }
217            },
218            Credential::Cert { .. } => {
219                error!("Certificate credentials are not supported for add_identity");
220                return Err(SSHAgentError::other(io::Error::new(
221                    io::ErrorKind::Unsupported,
222                    "Certificate credentials are not supported",
223                )));
224            }
225        };
226
227        self.handle
228            .register_key(Zeroizing::new(pkcs8_bytes))
229            .map_err(|e| {
230                let err_msg = format!("Failed to register key in agent: {}", e);
231                error!("{}", err_msg);
232                SSHAgentError::other(io::Error::other(err_msg))
233            })?;
234
235        debug!("Successfully added identity to agent");
236        Ok(())
237    }
238
239    async fn remove_identity(&mut self, identity: RemoveIdentity) -> Result<(), SSHAgentError> {
240        debug!("Handling remove_identity request");
241
242        let pubkey_bytes = match &identity.pubkey {
243            KeyData::Ed25519(key) => key.as_ref().to_vec(),
244            KeyData::Ecdsa(key) => key.as_ref().to_vec(),
245            other => {
246                let err_msg = format!(
247                    "Unsupported key type for remove_identity: {:?}",
248                    other.algorithm()
249                );
250                error!("{}", err_msg);
251                return Err(SSHAgentError::other(io::Error::new(
252                    io::ErrorKind::Unsupported,
253                    err_msg,
254                )));
255            }
256        };
257
258        let mut core = self.handle.lock().map_err(|_| {
259            error!("AgentSession failed to lock agent core (mutex poisoned).");
260            SSHAgentError::Failure
261        })?;
262
263        core.unregister_key(&pubkey_bytes).map_err(|e| {
264            let err_msg = format!("Failed to remove key from agent: {}", e);
265            error!("{}", err_msg);
266            SSHAgentError::other(io::Error::new(io::ErrorKind::NotFound, err_msg))
267        })?;
268
269        debug!("Successfully removed identity from agent");
270        Ok(())
271    }
272
273    async fn remove_all_identities(&mut self) -> Result<(), SSHAgentError> {
274        debug!("Handling remove_all_identities request");
275
276        let mut core = self.handle.lock().map_err(|_| {
277            error!("AgentSession failed to lock agent core (mutex poisoned).");
278            SSHAgentError::Failure
279        })?;
280
281        core.clear_keys();
282        debug!("Successfully removed all identities from agent");
283        Ok(())
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use ring::rand::SystemRandom;
291    use ring::signature::Ed25519KeyPair;
292    use ssh_key::private::Ed25519Keypair as SshEd25519Keypair;
293    use std::path::PathBuf;
294    use zeroize::Zeroizing;
295
296    fn generate_test_pkcs8() -> Vec<u8> {
297        let rng = SystemRandom::new();
298        let pkcs8_doc = Ed25519KeyPair::generate_pkcs8(&rng).expect("Failed to generate PKCS#8");
299        pkcs8_doc.as_ref().to_vec()
300    }
301
302    #[test]
303    fn test_agent_session_new() {
304        let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test.sock")));
305        let session = AgentSession::new(handle.clone());
306
307        assert_eq!(
308            session.handle().socket_path(),
309            &PathBuf::from("/tmp/test.sock")
310        );
311    }
312
313    #[test]
314    fn test_agent_session_clone_shares_handle() {
315        let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test.sock")));
316
317        let pkcs8_bytes = generate_test_pkcs8();
318        handle
319            .register_key(Zeroizing::new(pkcs8_bytes))
320            .expect("Failed to register key");
321
322        let session1 = AgentSession::new(handle.clone());
323        let session2 = session1.clone();
324
325        // Both sessions share the same handle
326        assert_eq!(session1.handle().key_count().unwrap(), 1);
327        assert_eq!(session2.handle().key_count().unwrap(), 1);
328    }
329
330    #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
331    async fn test_add_identity_round_trip() {
332        // Generate a test seed and create an SSH keypair from it
333        let seed: [u8; 32] = {
334            let pkcs8 = generate_test_pkcs8();
335            // Extract seed from the PKCS#8 bytes (bytes 16..48 in ring's format)
336            let mut s = [0u8; 32];
337            s.copy_from_slice(&pkcs8[16..48]);
338            s
339        };
340
341        let ssh_keypair = SshEd25519Keypair::from_seed(&seed);
342        let pubkey_bytes = ssh_keypair.public.0;
343
344        // Build an AddIdentity request (what the client sends over the wire)
345        let identity = AddIdentity {
346            credential: Credential::Key {
347                privkey: KeypairData::Ed25519(ssh_keypair),
348                comment: "test-key".to_string(),
349            },
350        };
351
352        // Create session with empty handle
353        let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test-add.sock")));
354        let mut session = AgentSession::new(handle.clone());
355
356        // Verify no keys initially
357        assert_eq!(handle.key_count().unwrap(), 0);
358
359        // Add the identity via the session (this is what was broken before)
360        session.add_identity(identity).await.unwrap();
361
362        // Verify the key is now registered
363        assert_eq!(handle.key_count().unwrap(), 1);
364
365        // Sign data via the session and verify the signature
366        let sign_request = SignRequest {
367            pubkey: KeyData::Ed25519(Ed25519PublicKey(pubkey_bytes)),
368            data: b"test data for signing".to_vec(),
369            flags: 0,
370        };
371
372        let signature = session.sign(sign_request).await.unwrap();
373        assert_eq!(signature.algorithm(), Algorithm::Ed25519);
374        assert!(!signature.as_bytes().is_empty());
375
376        // Verify the signature using ring
377        let ring_pubkey =
378            ring::signature::UnparsedPublicKey::new(&ring::signature::ED25519, &pubkey_bytes);
379        ring_pubkey
380            .verify(b"test data for signing", signature.as_bytes())
381            .expect("Signature verification failed");
382    }
383
384    #[tokio::test]
385    async fn test_remove_identity() {
386        let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test-rm.sock")));
387        let mut session = AgentSession::new(handle.clone());
388
389        // Add a key via add_identity
390        let seed: [u8; 32] = {
391            let pkcs8 = generate_test_pkcs8();
392            let mut s = [0u8; 32];
393            s.copy_from_slice(&pkcs8[16..48]);
394            s
395        };
396        let ssh_keypair = SshEd25519Keypair::from_seed(&seed);
397        let pubkey_bytes = ssh_keypair.public.0;
398
399        let identity = AddIdentity {
400            credential: Credential::Key {
401                privkey: KeypairData::Ed25519(ssh_keypair),
402                comment: "test-key".to_string(),
403            },
404        };
405        session.add_identity(identity).await.unwrap();
406        assert_eq!(handle.key_count().unwrap(), 1);
407
408        // Remove it
409        let remove = RemoveIdentity {
410            pubkey: KeyData::Ed25519(Ed25519PublicKey(pubkey_bytes)),
411        };
412        session.remove_identity(remove).await.unwrap();
413        assert_eq!(handle.key_count().unwrap(), 0);
414    }
415
416    #[tokio::test]
417    async fn test_remove_all_identities() {
418        let handle = Arc::new(AgentHandle::new(PathBuf::from("/tmp/test-rmall.sock")));
419        let mut session = AgentSession::new(handle.clone());
420
421        // Add two keys
422        for _ in 0..2 {
423            let seed: [u8; 32] = {
424                let pkcs8 = generate_test_pkcs8();
425                let mut s = [0u8; 32];
426                s.copy_from_slice(&pkcs8[16..48]);
427                s
428            };
429            let ssh_keypair = SshEd25519Keypair::from_seed(&seed);
430            let identity = AddIdentity {
431                credential: Credential::Key {
432                    privkey: KeypairData::Ed25519(ssh_keypair),
433                    comment: "test-key".to_string(),
434                },
435            };
436            session.add_identity(identity).await.unwrap();
437        }
438        assert_eq!(handle.key_count().unwrap(), 2);
439
440        // Remove all
441        session.remove_all_identities().await.unwrap();
442        assert_eq!(handle.key_count().unwrap(), 0);
443    }
444}