use aes_gcm::aead::{Aead, KeyInit, OsRng as AeadOsRng, generic_array::GenericArray};
use aes_gcm::{Aes256Gcm, Nonce};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::Utc;
use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
use hkdf::Hkdf;
use rand::RngCore;
use serde::Serialize;
use sha2::{Digest, Sha256};
use thiserror::Error;
use uuid::Uuid;
use x25519_dalek::{PublicKey as X25519PublicKey, StaticSecret};
use crate::ids::ThreadId;
use crate::protocol::{AgentIdentity, AgentKeypairExport, Envelope, MessageKind, ThreadSecret};
#[derive(Debug, Error)]
pub enum CryptoError {
#[error("invalid base64")]
Base64(#[from] base64::DecodeError),
#[error("invalid key material")]
InvalidKey,
#[error("signature verification failed")]
BadSignature,
#[error("encryption failed")]
Encrypt,
#[error("decryption failed")]
Decrypt,
#[error("json serialization failed")]
Json(#[from] serde_json::Error),
}
#[derive(Debug, Clone)]
pub struct AgentKeypair {
pub export: AgentKeypairExport,
}
impl AgentKeypair {
pub fn generate(handle: impl Into<String>) -> Self {
let signing = SigningKey::generate(&mut rand::rngs::OsRng);
let encryption = StaticSecret::random_from_rng(rand::rngs::OsRng);
let encryption_public = X25519PublicKey::from(&encryption);
let handle = handle.into();
let now = Utc::now();
Self {
export: AgentKeypairExport {
handle,
agent_id: format!("agt_{}", Uuid::new_v4().simple()),
signing_secret_key: b64(signing.to_bytes()),
signing_public_key: b64(signing.verifying_key().to_bytes()),
encryption_secret_key: b64(encryption.to_bytes()),
encryption_public_key: b64(encryption_public.to_bytes()),
created_at: now,
},
}
}
pub fn from_export(export: AgentKeypairExport) -> Self {
Self { export }
}
pub fn identity(&self) -> AgentIdentity {
AgentIdentity {
handle: self.export.handle.clone(),
agent_id: self.export.agent_id.clone(),
signing_public_key: self.export.signing_public_key.clone(),
encryption_public_key: self.export.encryption_public_key.clone(),
created_at: self.export.created_at,
}
}
pub fn sign_json<T: Serialize>(&self, value: &T) -> Result<String, CryptoError> {
let bytes = serde_json::to_vec(value)?;
self.sign_bytes(&bytes)
}
pub fn sign_bytes(&self, bytes: &[u8]) -> Result<String, CryptoError> {
let secret = decode_array::<32>(&self.export.signing_secret_key)?;
let key = SigningKey::from_bytes(&secret);
Ok(b64(key.sign(bytes).to_bytes()))
}
pub fn encryption_secret(&self) -> Result<StaticSecret, CryptoError> {
Ok(StaticSecret::from(decode_array::<32>(
&self.export.encryption_secret_key,
)?))
}
}
pub fn verify_json_signature<T: Serialize>(
public_key: &str,
value: &T,
signature: &str,
) -> Result<(), CryptoError> {
let bytes = serde_json::to_vec(value)?;
verify_bytes_signature(public_key, &bytes, signature)
}
pub fn verify_bytes_signature(
public_key: &str,
bytes: &[u8],
signature: &str,
) -> Result<(), CryptoError> {
let key = VerifyingKey::from_bytes(&decode_array::<32>(public_key)?)
.map_err(|_| CryptoError::InvalidKey)?;
let sig = Signature::from_bytes(&decode_array::<64>(signature)?);
key.verify(bytes, &sig)
.map_err(|_| CryptoError::BadSignature)
}
pub fn new_thread_secret(thread_id: ThreadId) -> ThreadSecret {
let mut key = [0_u8; 32];
rand::rngs::OsRng.fill_bytes(&mut key);
ThreadSecret {
thread_id,
thread_key: b64(key),
epoch: 1,
}
}
pub fn encrypt_for_thread(
thread_key: &str,
plaintext: &[u8],
) -> Result<(String, String), CryptoError> {
encrypt_with_key(&decode_array::<32>(thread_key)?, plaintext)
}
pub fn decrypt_for_thread(
thread_key: &str,
nonce: &str,
ciphertext: &str,
) -> Result<Vec<u8>, CryptoError> {
decrypt_with_key(&decode_array::<32>(thread_key)?, nonce, ciphertext)
}
pub fn encrypt_for_agent(
sender: &AgentKeypair,
recipient_encryption_public_key: &str,
plaintext: &[u8],
) -> Result<(String, String), CryptoError> {
let sender_secret = sender.encryption_secret()?;
let recipient_public =
X25519PublicKey::from(decode_array::<32>(recipient_encryption_public_key)?);
let shared = sender_secret.diffie_hellman(&recipient_public);
let key = derive_welcome_key(shared.as_bytes());
encrypt_with_key(&key, plaintext)
}
pub fn decrypt_from_agent(
recipient: &AgentKeypair,
sender_encryption_public_key: &str,
nonce: &str,
ciphertext: &str,
) -> Result<Vec<u8>, CryptoError> {
let recipient_secret = recipient.encryption_secret()?;
let sender_public = X25519PublicKey::from(decode_array::<32>(sender_encryption_public_key)?);
let shared = recipient_secret.diffie_hellman(&sender_public);
let key = derive_welcome_key(shared.as_bytes());
decrypt_with_key(&key, nonce, ciphertext)
}
pub fn build_signed_envelope(
keypair: &AgentKeypair,
thread_id: ThreadId,
epoch: u64,
kind: MessageKind,
ciphertext: String,
nonce: String,
artifact: Option<crate::protocol::ArtifactManifest>,
) -> Result<Envelope, CryptoError> {
let mut envelope = Envelope {
envelope_id: Uuid::new_v4(),
thread_id,
seq: None,
epoch,
sender: keypair.identity(),
kind,
created_at: Utc::now(),
ciphertext,
nonce,
artifact,
signature: String::new(),
};
envelope.signature = signable_envelope_signature(keypair, &envelope)?;
Ok(envelope)
}
pub fn verify_envelope(envelope: &Envelope) -> Result<(), CryptoError> {
let mut signable = envelope.clone();
signable.seq = None;
signable.signature.clear();
let bytes = serde_json::to_vec(&signable)?;
verify_bytes_signature(
&envelope.sender.signing_public_key,
&bytes,
&envelope.signature,
)
}
pub fn sha256_b64(bytes: &[u8]) -> String {
b64(Sha256::digest(bytes))
}
fn signable_envelope_signature(
keypair: &AgentKeypair,
envelope: &Envelope,
) -> Result<String, CryptoError> {
let mut signable = envelope.clone();
signable.seq = None;
signable.signature.clear();
let bytes = serde_json::to_vec(&signable)?;
keypair.sign_bytes(&bytes)
}
fn encrypt_with_key(key: &[u8; 32], plaintext: &[u8]) -> Result<(String, String), CryptoError> {
let cipher = Aes256Gcm::new(GenericArray::from_slice(key));
let mut nonce = [0_u8; 12];
AeadOsRng.fill_bytes(&mut nonce);
let ciphertext = cipher
.encrypt(Nonce::from_slice(&nonce), plaintext)
.map_err(|_| CryptoError::Encrypt)?;
Ok((b64(ciphertext), b64(nonce)))
}
fn decrypt_with_key(key: &[u8; 32], nonce: &str, ciphertext: &str) -> Result<Vec<u8>, CryptoError> {
let cipher = Aes256Gcm::new(GenericArray::from_slice(key));
cipher
.decrypt(
Nonce::from_slice(&decode_array::<12>(nonce)?),
decode_vec(ciphertext)?.as_ref(),
)
.map_err(|_| CryptoError::Decrypt)
}
fn derive_welcome_key(shared: &[u8; 32]) -> [u8; 32] {
let hk = Hkdf::<Sha256>::new(Some(b"f8s-welcome-v1"), shared);
let mut out = [0_u8; 32];
hk.expand(b"agent-thread-key-wrap", &mut out)
.expect("fixed output length is valid");
out
}
fn decode_vec(value: &str) -> Result<Vec<u8>, CryptoError> {
Ok(URL_SAFE_NO_PAD.decode(value)?)
}
fn decode_array<const N: usize>(value: &str) -> Result<[u8; N], CryptoError> {
let bytes = decode_vec(value)?;
bytes.try_into().map_err(|_| CryptoError::InvalidKey)
}
fn b64(bytes: impl AsRef<[u8]>) -> String {
URL_SAFE_NO_PAD.encode(bytes)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ids::ThreadId;
#[test]
fn signs_and_verifies_json() {
let agent = AgentKeypair::generate("codex");
let value = serde_json::json!({"hello": "world"});
let sig = agent.sign_json(&value).unwrap();
verify_json_signature(&agent.identity().signing_public_key, &value, &sig).unwrap();
}
#[test]
fn thread_encryption_round_trip() {
let secret = new_thread_secret(ThreadId::new());
let (ciphertext, nonce) = encrypt_for_thread(&secret.thread_key, b"hello").unwrap();
let plaintext = decrypt_for_thread(&secret.thread_key, &nonce, &ciphertext).unwrap();
assert_eq!(plaintext, b"hello");
}
#[test]
fn agent_encryption_round_trip() {
let alice = AgentKeypair::generate("alice");
let bob = AgentKeypair::generate("bob");
let (ciphertext, nonce) =
encrypt_for_agent(&alice, &bob.identity().encryption_public_key, b"welcome").unwrap();
let plaintext = decrypt_from_agent(
&bob,
&alice.identity().encryption_public_key,
&nonce,
&ciphertext,
)
.unwrap();
assert_eq!(plaintext, b"welcome");
}
#[test]
fn server_assigned_sequence_does_not_break_signature() {
let agent = AgentKeypair::generate("alice");
let secret = new_thread_secret(ThreadId::new());
let (ciphertext, nonce) = encrypt_for_thread(&secret.thread_key, b"hello").unwrap();
let mut envelope = build_signed_envelope(
&agent,
secret.thread_id,
secret.epoch,
MessageKind::Text,
ciphertext,
nonce,
None,
)
.unwrap();
envelope.seq = Some(42);
verify_envelope(&envelope).unwrap();
}
}