auths_core/agent/
core.rs

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

auths_core/agent/
core.rs
