1use aes_gcm::aead::{Aead, KeyInit, OsRng as AeadOsRng, generic_array::GenericArray};
2use aes_gcm::{Aes256Gcm, Nonce};
3use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
4use chrono::Utc;
5use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
6use hkdf::Hkdf;
7use rand::RngCore;
8use serde::Serialize;
9use sha2::{Digest, Sha256};
10use thiserror::Error;
11use uuid::Uuid;
12use x25519_dalek::{PublicKey as X25519PublicKey, StaticSecret};
13
14use crate::ids::ThreadId;
15use crate::protocol::{AgentIdentity, AgentKeypairExport, Envelope, MessageKind, ThreadSecret};
16
17#[derive(Debug, Error)]
18pub enum CryptoError {
19 #[error("invalid base64")]
20 Base64(#[from] base64::DecodeError),
21 #[error("invalid key material")]
22 InvalidKey,
23 #[error("signature verification failed")]
24 BadSignature,
25 #[error("encryption failed")]
26 Encrypt,
27 #[error("decryption failed")]
28 Decrypt,
29 #[error("json serialization failed")]
30 Json(#[from] serde_json::Error),
31}
32
33#[derive(Debug, Clone)]
34pub struct AgentKeypair {
35 pub export: AgentKeypairExport,
36}
37
38impl AgentKeypair {
39 pub fn generate(handle: impl Into<String>) -> Self {
40 let signing = SigningKey::generate(&mut rand::rngs::OsRng);
41 let encryption = StaticSecret::random_from_rng(rand::rngs::OsRng);
42 let encryption_public = X25519PublicKey::from(&encryption);
43 let handle = handle.into();
44 let now = Utc::now();
45
46 Self {
47 export: AgentKeypairExport {
48 handle,
49 agent_id: format!("agt_{}", Uuid::new_v4().simple()),
50 signing_secret_key: b64(signing.to_bytes()),
51 signing_public_key: b64(signing.verifying_key().to_bytes()),
52 encryption_secret_key: b64(encryption.to_bytes()),
53 encryption_public_key: b64(encryption_public.to_bytes()),
54 created_at: now,
55 },
56 }
57 }
58
59 pub fn from_export(export: AgentKeypairExport) -> Self {
60 Self { export }
61 }
62
63 pub fn identity(&self) -> AgentIdentity {
64 AgentIdentity {
65 handle: self.export.handle.clone(),
66 agent_id: self.export.agent_id.clone(),
67 signing_public_key: self.export.signing_public_key.clone(),
68 encryption_public_key: self.export.encryption_public_key.clone(),
69 created_at: self.export.created_at,
70 }
71 }
72
73 pub fn sign_json<T: Serialize>(&self, value: &T) -> Result<String, CryptoError> {
74 let bytes = serde_json::to_vec(value)?;
75 self.sign_bytes(&bytes)
76 }
77
78 pub fn sign_bytes(&self, bytes: &[u8]) -> Result<String, CryptoError> {
79 let secret = decode_array::<32>(&self.export.signing_secret_key)?;
80 let key = SigningKey::from_bytes(&secret);
81 Ok(b64(key.sign(bytes).to_bytes()))
82 }
83
84 pub fn encryption_secret(&self) -> Result<StaticSecret, CryptoError> {
85 Ok(StaticSecret::from(decode_array::<32>(
86 &self.export.encryption_secret_key,
87 )?))
88 }
89}
90
91pub fn verify_json_signature<T: Serialize>(
92 public_key: &str,
93 value: &T,
94 signature: &str,
95) -> Result<(), CryptoError> {
96 let bytes = serde_json::to_vec(value)?;
97 verify_bytes_signature(public_key, &bytes, signature)
98}
99
100pub fn verify_bytes_signature(
101 public_key: &str,
102 bytes: &[u8],
103 signature: &str,
104) -> Result<(), CryptoError> {
105 let key = VerifyingKey::from_bytes(&decode_array::<32>(public_key)?)
106 .map_err(|_| CryptoError::InvalidKey)?;
107 let sig = Signature::from_bytes(&decode_array::<64>(signature)?);
108 key.verify(bytes, &sig)
109 .map_err(|_| CryptoError::BadSignature)
110}
111
112pub fn new_thread_secret(thread_id: ThreadId) -> ThreadSecret {
113 let mut key = [0_u8; 32];
114 rand::rngs::OsRng.fill_bytes(&mut key);
115 ThreadSecret {
116 thread_id,
117 thread_key: b64(key),
118 epoch: 1,
119 }
120}
121
122pub fn encrypt_for_thread(
123 thread_key: &str,
124 plaintext: &[u8],
125) -> Result<(String, String), CryptoError> {
126 encrypt_with_key(&decode_array::<32>(thread_key)?, plaintext)
127}
128
129pub fn decrypt_for_thread(
130 thread_key: &str,
131 nonce: &str,
132 ciphertext: &str,
133) -> Result<Vec<u8>, CryptoError> {
134 decrypt_with_key(&decode_array::<32>(thread_key)?, nonce, ciphertext)
135}
136
137pub fn encrypt_for_agent(
138 sender: &AgentKeypair,
139 recipient_encryption_public_key: &str,
140 plaintext: &[u8],
141) -> Result<(String, String), CryptoError> {
142 let sender_secret = sender.encryption_secret()?;
143 let recipient_public =
144 X25519PublicKey::from(decode_array::<32>(recipient_encryption_public_key)?);
145 let shared = sender_secret.diffie_hellman(&recipient_public);
146 let key = derive_welcome_key(shared.as_bytes());
147 encrypt_with_key(&key, plaintext)
148}
149
150pub fn decrypt_from_agent(
151 recipient: &AgentKeypair,
152 sender_encryption_public_key: &str,
153 nonce: &str,
154 ciphertext: &str,
155) -> Result<Vec<u8>, CryptoError> {
156 let recipient_secret = recipient.encryption_secret()?;
157 let sender_public = X25519PublicKey::from(decode_array::<32>(sender_encryption_public_key)?);
158 let shared = recipient_secret.diffie_hellman(&sender_public);
159 let key = derive_welcome_key(shared.as_bytes());
160 decrypt_with_key(&key, nonce, ciphertext)
161}
162
163pub fn build_signed_envelope(
164 keypair: &AgentKeypair,
165 thread_id: ThreadId,
166 epoch: u64,
167 kind: MessageKind,
168 ciphertext: String,
169 nonce: String,
170 artifact: Option<crate::protocol::ArtifactManifest>,
171) -> Result<Envelope, CryptoError> {
172 let mut envelope = Envelope {
173 envelope_id: Uuid::new_v4(),
174 thread_id,
175 seq: None,
176 epoch,
177 sender: keypair.identity(),
178 kind,
179 created_at: Utc::now(),
180 ciphertext,
181 nonce,
182 artifact,
183 signature: String::new(),
184 };
185 envelope.signature = signable_envelope_signature(keypair, &envelope)?;
186 Ok(envelope)
187}
188
189pub fn verify_envelope(envelope: &Envelope) -> Result<(), CryptoError> {
190 let mut signable = envelope.clone();
191 signable.seq = None;
192 signable.signature.clear();
193 let bytes = serde_json::to_vec(&signable)?;
194 verify_bytes_signature(
195 &envelope.sender.signing_public_key,
196 &bytes,
197 &envelope.signature,
198 )
199}
200
201pub fn sha256_b64(bytes: &[u8]) -> String {
202 b64(Sha256::digest(bytes))
203}
204
205fn signable_envelope_signature(
206 keypair: &AgentKeypair,
207 envelope: &Envelope,
208) -> Result<String, CryptoError> {
209 let mut signable = envelope.clone();
210 signable.seq = None;
211 signable.signature.clear();
212 let bytes = serde_json::to_vec(&signable)?;
213 keypair.sign_bytes(&bytes)
214}
215
216fn encrypt_with_key(key: &[u8; 32], plaintext: &[u8]) -> Result<(String, String), CryptoError> {
217 let cipher = Aes256Gcm::new(GenericArray::from_slice(key));
218 let mut nonce = [0_u8; 12];
219 AeadOsRng.fill_bytes(&mut nonce);
220 let ciphertext = cipher
221 .encrypt(Nonce::from_slice(&nonce), plaintext)
222 .map_err(|_| CryptoError::Encrypt)?;
223 Ok((b64(ciphertext), b64(nonce)))
224}
225
226fn decrypt_with_key(key: &[u8; 32], nonce: &str, ciphertext: &str) -> Result<Vec<u8>, CryptoError> {
227 let cipher = Aes256Gcm::new(GenericArray::from_slice(key));
228 cipher
229 .decrypt(
230 Nonce::from_slice(&decode_array::<12>(nonce)?),
231 decode_vec(ciphertext)?.as_ref(),
232 )
233 .map_err(|_| CryptoError::Decrypt)
234}
235
236fn derive_welcome_key(shared: &[u8; 32]) -> [u8; 32] {
237 let hk = Hkdf::<Sha256>::new(Some(b"f8s-welcome-v1"), shared);
238 let mut out = [0_u8; 32];
239 hk.expand(b"agent-thread-key-wrap", &mut out)
240 .expect("fixed output length is valid");
241 out
242}
243
244fn decode_vec(value: &str) -> Result<Vec<u8>, CryptoError> {
245 Ok(URL_SAFE_NO_PAD.decode(value)?)
246}
247
248fn decode_array<const N: usize>(value: &str) -> Result<[u8; N], CryptoError> {
249 let bytes = decode_vec(value)?;
250 bytes.try_into().map_err(|_| CryptoError::InvalidKey)
251}
252
253fn b64(bytes: impl AsRef<[u8]>) -> String {
254 URL_SAFE_NO_PAD.encode(bytes)
255}
256
257#[cfg(test)]
258mod tests {
259 use super::*;
260 use crate::ids::ThreadId;
261
262 #[test]
263 fn signs_and_verifies_json() {
264 let agent = AgentKeypair::generate("codex");
265 let value = serde_json::json!({"hello": "world"});
266 let sig = agent.sign_json(&value).unwrap();
267 verify_json_signature(&agent.identity().signing_public_key, &value, &sig).unwrap();
268 }
269
270 #[test]
271 fn thread_encryption_round_trip() {
272 let secret = new_thread_secret(ThreadId::new());
273 let (ciphertext, nonce) = encrypt_for_thread(&secret.thread_key, b"hello").unwrap();
274 let plaintext = decrypt_for_thread(&secret.thread_key, &nonce, &ciphertext).unwrap();
275 assert_eq!(plaintext, b"hello");
276 }
277
278 #[test]
279 fn agent_encryption_round_trip() {
280 let alice = AgentKeypair::generate("alice");
281 let bob = AgentKeypair::generate("bob");
282 let (ciphertext, nonce) =
283 encrypt_for_agent(&alice, &bob.identity().encryption_public_key, b"welcome").unwrap();
284 let plaintext = decrypt_from_agent(
285 &bob,
286 &alice.identity().encryption_public_key,
287 &nonce,
288 &ciphertext,
289 )
290 .unwrap();
291 assert_eq!(plaintext, b"welcome");
292 }
293
294 #[test]
295 fn server_assigned_sequence_does_not_break_signature() {
296 let agent = AgentKeypair::generate("alice");
297 let secret = new_thread_secret(ThreadId::new());
298 let (ciphertext, nonce) = encrypt_for_thread(&secret.thread_key, b"hello").unwrap();
299 let mut envelope = build_signed_envelope(
300 &agent,
301 secret.thread_id,
302 secret.epoch,
303 MessageKind::Text,
304 ciphertext,
305 nonce,
306 None,
307 )
308 .unwrap();
309 envelope.seq = Some(42);
310 verify_envelope(&envelope).unwrap();
311 }
312}