auths_core/agent/
core.rs

1use crate::crypto::provider_bridge;
2use crate::error::AgentError;
3use auths_crypto::SecureSeed;
4use log::error;
5use std::collections::HashMap;
6use std::fmt;
7use zeroize::Zeroizing;
8
9/// An in-memory registry of SSH keys used by the local agent.
10/// Stores seeds securely using SecureSeed (zeroize-on-drop).
11/// Note: Clone is intentionally NOT derived to prevent accidental copying of key material.
12#[derive(Default)]
13pub struct AgentCore {
14    /// Maps public key bytes (Vec<u8>) to the corresponding SecureSeed.
15    pub keys: HashMap<Vec<u8>, SecureSeed>,
16}
17
18impl fmt::Debug for AgentCore {
19    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
20        f.debug_struct("AgentCore")
21            .field("key_count", &self.keys.len())
22            .finish_non_exhaustive()
23    }
24}
25
26impl AgentCore {
27    /// Create a new `AgentState`.
28    pub fn new() -> Self {
29        Self::default()
30    }
31
32    /// Registers decrypted PKCS#8 key bytes in memory.
33    /// Extracts the seed and public key, validates them, and stores the seed.
34    ///
35    /// Args:
36    /// * `pkcs8_bytes` - The raw, decrypted PKCS#8 bytes for the Ed25519 key, wrapped in `Zeroizing`.
37    pub fn register_key(&mut self, pkcs8_bytes: Zeroizing<Vec<u8>>) -> Result<(), AgentError> {
38        let (seed, pubkey) = crate::crypto::signer::load_seed_and_pubkey(&pkcs8_bytes)?;
39        self.keys.insert(pubkey.to_vec(), seed);
40        Ok(())
41    }
42
43    /// Removes a key by its public key bytes. Returns error if key not found.
44    pub fn unregister_key(&mut self, pubkey: &[u8]) -> Result<(), AgentError> {
45        self.keys
46            .remove(pubkey)
47            .map(|_| ())
48            .ok_or(AgentError::KeyNotFound)
49    }
50
51    /// Signs a message using the key associated with the given public key bytes.
52    /// Routes through CryptoProvider via the sync bridge.
53    pub fn sign(&self, pubkey_to_find: &[u8], data: &[u8]) -> Result<Vec<u8>, AgentError> {
54        let seed = self
55            .keys
56            .get(pubkey_to_find)
57            .ok_or(AgentError::KeyNotFound)?;
58
59        provider_bridge::sign_ed25519_sync(seed, data).map_err(|e| {
60            error!("CryptoProvider signing failed: {}", e);
61            AgentError::CryptoError(format!("Ed25519 signing failed: {}", e))
62        })
63    }
64
65    /// Returns all public key bytes currently registered.
66    pub fn public_keys(&self) -> Vec<Vec<u8>> {
67        self.keys.keys().cloned().collect()
68    }
69
70    /// Returns the number of keys currently loaded in the agent.
71    pub fn key_count(&self) -> usize {
72        self.keys.len()
73    }
74
75    /// Removes all keys from the agent core.
76    pub fn clear_keys(&mut self) {
77        self.keys.clear();
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use ring::rand::SystemRandom;
85    use ring::signature::{ED25519, Ed25519KeyPair, KeyPair, UnparsedPublicKey};
86    use zeroize::Zeroizing;
87
88    fn generate_test_key_bytes() -> (Vec<u8>, Vec<u8>) {
89        let rng = SystemRandom::new();
90        let pkcs8_doc = Ed25519KeyPair::generate_pkcs8(&rng).expect("Failed to generate PKCS#8");
91        let pkcs8_bytes = pkcs8_doc.as_ref().to_vec();
92        let keypair =
93            Ed25519KeyPair::from_pkcs8(&pkcs8_bytes).expect("Failed to parse generated PKCS#8");
94        let pubkey_bytes = keypair.public_key().as_ref().to_vec();
95        (pubkey_bytes, pkcs8_bytes)
96    }
97
98    fn core_with_two_keys() -> (AgentCore, Vec<u8>, Vec<u8>) {
99        let (pubkey1, pkcs8_1) = generate_test_key_bytes();
100        let (pubkey2, pkcs8_2) = generate_test_key_bytes();
101        let mut core = AgentCore::default();
102        core.register_key(Zeroizing::new(pkcs8_1)).unwrap();
103        core.register_key(Zeroizing::new(pkcs8_2)).unwrap();
104        (core, pubkey1, pubkey2)
105    }
106
107    #[test]
108    fn register_keys_updates_count_and_listing() {
109        let (core, pubkey1, pubkey2) = core_with_two_keys();
110        assert_eq!(core.key_count(), 2);
111        let mut keys = core.public_keys();
112        keys.sort();
113        let mut expected = vec![pubkey1, pubkey2];
114        expected.sort();
115        assert_eq!(keys, expected);
116    }
117
118    #[test]
119    fn sign_produces_verifiable_signature() {
120        let (core, pubkey1, _) = core_with_two_keys();
121        let message = b"test message for agent core";
122        let signature = core.sign(&pubkey1, message).unwrap();
123        assert!(!signature.is_empty());
124        let ring_key = UnparsedPublicKey::new(&ED25519, &pubkey1);
125        assert!(ring_key.verify(message, &signature).is_ok());
126    }
127
128    #[test]
129    fn sign_with_different_keys_produces_different_signatures() {
130        let (core, pubkey1, pubkey2) = core_with_two_keys();
131        let message = b"test message";
132        let sig1 = core.sign(&pubkey1, message).unwrap();
133        let sig2 = core.sign(&pubkey2, message).unwrap();
134        assert_ne!(sig1, sig2);
135    }
136
137    #[test]
138    fn sign_with_nonexistent_key_returns_key_not_found() {
139        let (core, _, _) = core_with_two_keys();
140        let err = core.sign(&[99u8; 32], b"msg").unwrap_err();
141        assert!(matches!(err, AgentError::KeyNotFound));
142    }
143
144    #[test]
145    fn unregister_removes_key_and_prevents_signing() {
146        let (mut core, pubkey1, pubkey2) = core_with_two_keys();
147        core.unregister_key(&pubkey1).unwrap();
148        assert_eq!(core.key_count(), 1);
149        assert_eq!(core.public_keys(), vec![pubkey2.clone()]);
150        assert!(matches!(
151            core.sign(&pubkey1, b"msg").unwrap_err(),
152            AgentError::KeyNotFound
153        ));
154        // Remaining key still works
155        assert!(core.sign(&pubkey2, b"msg").is_ok());
156    }
157
158    #[test]
159    fn clear_keys_removes_all() {
160        let (mut core, pubkey1, _) = core_with_two_keys();
161        core.clear_keys();
162        assert_eq!(core.key_count(), 0);
163        assert!(core.public_keys().is_empty());
164        assert!(matches!(
165            core.unregister_key(&pubkey1).unwrap_err(),
166            AgentError::KeyNotFound
167        ));
168    }
169
170    #[test]
171    fn test_register_invalid_key() {
172        let mut core = AgentCore::default();
173        let invalid_bytes = vec![1, 2, 3, 4];
174        let result = core.register_key(Zeroizing::new(invalid_bytes));
175        assert!(result.is_err());
176        assert_eq!(core.key_count(), 0);
177    }
178}
auths_core/agent/core.rs

auths_core/agent/
core.rs
