use chrono::{DateTime, Utc};
use ed25519_dalek::SigningKey;
use sha2::Digest;
use crate::signer::{verify_hex, DevSigner, SignerBackend, VerifyError};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MockKmsKeyMeta {
pub role: String,
pub version: u32,
pub key_id: String,
pub public_hex: String,
pub created_at: DateTime<Utc>,
pub mock: bool,
}
#[derive(Debug, Clone)]
pub struct MockKmsSigner {
role: String,
root_seed: [u8; 32],
genesis: DateTime<Utc>,
versions: Vec<MockKmsKeyMeta>,
current_index: usize,
}
impl MockKmsSigner {
pub fn new(role: impl Into<String>, root_seed: [u8; 32], genesis: DateTime<Utc>) -> Self {
let role = role.into();
let v1 = Self::build_meta(&role, 1, &root_seed, genesis);
Self {
role,
root_seed,
genesis,
versions: vec![v1],
current_index: 0,
}
}
pub fn from_versions(
role: impl Into<String>,
root_seed: [u8; 32],
genesis: DateTime<Utc>,
current_version: u32,
) -> Self {
assert!(current_version >= 1, "version numbering starts at 1");
let role = role.into();
let versions: Vec<_> = (1..=current_version)
.map(|v| Self::build_meta(&role, v, &root_seed, genesis))
.collect();
let current_index = (current_version as usize) - 1;
Self {
role,
root_seed,
genesis,
versions,
current_index,
}
}
fn build_meta(
role: &str,
version: u32,
root_seed: &[u8; 32],
genesis: DateTime<Utc>,
) -> MockKmsKeyMeta {
let signing_key = derive_signing_key(role, version, root_seed);
let public_hex = hex::encode(signing_key.verifying_key().to_bytes());
let key_id = format!("{role}-v{version}");
let created_at =
genesis + chrono::Duration::nanoseconds(((version as i64).saturating_sub(1)).max(0));
MockKmsKeyMeta {
role: role.to_string(),
version,
key_id,
public_hex,
created_at,
mock: true,
}
}
pub fn role(&self) -> &str {
&self.role
}
pub fn current_version(&self) -> u32 {
self.versions[self.current_index].version
}
pub fn current(&self) -> &MockKmsKeyMeta {
&self.versions[self.current_index]
}
pub fn versions(&self) -> &[MockKmsKeyMeta] {
&self.versions
}
pub fn key_by_id(&self, key_id: &str) -> Option<&MockKmsKeyMeta> {
self.versions.iter().find(|m| m.key_id == key_id)
}
pub fn key_by_version(&self, version: u32) -> Option<&MockKmsKeyMeta> {
self.versions.iter().find(|m| m.version == version)
}
pub fn rotate(&mut self) -> &MockKmsKeyMeta {
let next_version = self.current_version() + 1;
let meta = Self::build_meta(&self.role, next_version, &self.root_seed, self.genesis);
self.versions.push(meta);
self.current_index = self.versions.len() - 1;
self.current()
}
fn current_signing_key(&self) -> SigningKey {
derive_signing_key(&self.role, self.current_version(), &self.root_seed)
}
pub fn verify(
&self,
key_id: &str,
message: &[u8],
signature_hex: &str,
) -> Result<(), VerifyError> {
let meta = self.key_by_id(key_id).ok_or(VerifyError::BadPublicKey)?;
verify_hex(&meta.public_hex, message, signature_hex)
}
pub fn current_as_dev_signer(&self) -> DevSigner {
DevSigner::from_seed(
self.current().key_id.clone(),
seed_for(&self.role, self.current_version(), &self.root_seed),
)
}
}
impl SignerBackend for MockKmsSigner {
fn current_key_id(&self) -> &str {
&self.current().key_id
}
fn sign_hex(&self, message: &[u8]) -> String {
use ed25519_dalek::Signer as _;
let sk = self.current_signing_key();
let sig = sk.sign(message);
hex::encode(sig.to_bytes())
}
fn current_public_hex(&self) -> String {
self.current().public_hex.clone()
}
}
pub fn derive_key_metadata(role: &str, version: u32, root_seed: &[u8; 32]) -> (String, String) {
let signing_key = derive_signing_key(role, version, root_seed);
let public_hex = hex::encode(signing_key.verifying_key().to_bytes());
let key_id = format!("{role}-v{version}");
(key_id, public_hex)
}
fn seed_for(role: &str, version: u32, root_seed: &[u8; 32]) -> [u8; 32] {
let mut hasher = sha2::Sha256::new();
hasher.update(b"sbo3l.mock_kms.v1");
hasher.update(root_seed);
hasher.update((role.len() as u32).to_be_bytes());
hasher.update(role.as_bytes());
hasher.update(version.to_be_bytes());
let digest = hasher.finalize();
let mut out = [0u8; 32];
out.copy_from_slice(&digest);
out
}
fn derive_signing_key(role: &str, version: u32, root_seed: &[u8; 32]) -> SigningKey {
SigningKey::from_bytes(&seed_for(role, version, root_seed))
}
#[cfg(test)]
mod tests {
use super::*;
fn ts(s: &str) -> DateTime<Utc> {
chrono::DateTime::parse_from_rfc3339(s).unwrap().into()
}
#[test]
fn fresh_signer_starts_at_v1() {
let s = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
assert_eq!(s.current_version(), 1);
assert_eq!(s.current_key_id(), "audit-mock-v1");
assert_eq!(s.versions().len(), 1);
assert!(s.current().mock, "metadata must surface mock = true");
}
#[test]
fn key_metadata_is_deterministic() {
let a = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
let b = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
assert_eq!(a.versions(), b.versions());
}
#[test]
fn different_roots_yield_different_keys() {
let a = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
let b = MockKmsSigner::new("audit-mock", [99u8; 32], ts("2026-04-28T00:00:00Z"));
assert_ne!(a.current().public_hex, b.current().public_hex);
}
#[test]
fn signer_backend_round_trip() {
let s = MockKmsSigner::new("decision-mock", [7u8; 32], ts("2026-04-28T00:00:00Z"));
let msg = b"some canonical payload";
let sig = SignerBackend::sign_hex(&s, msg);
s.verify(s.current_key_id(), msg, &sig)
.expect("self-verify must succeed");
}
#[test]
fn signer_backend_reports_consistent_metadata() {
let s = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
assert_eq!(s.current_key_id(), s.current().key_id);
assert_eq!(s.current_public_hex(), s.current().public_hex);
}
#[test]
fn rotate_advances_current_version() {
let mut s = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
let before = s.current().clone();
let after = s.rotate().clone();
assert_eq!(after.version, 2);
assert_eq!(after.key_id, "audit-mock-v2");
assert_ne!(
after.public_hex, before.public_hex,
"rotation must change public key"
);
assert_eq!(s.current_version(), 2);
assert_eq!(s.versions().len(), 2);
}
#[test]
fn old_signature_still_verifies_after_rotation() {
let mut s = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
let v1_key_id = s.current_key_id().to_string();
let msg = b"pre-rotation message";
let sig = SignerBackend::sign_hex(&s, msg);
s.rotate();
assert_eq!(s.current_version(), 2);
s.verify(&v1_key_id, msg, &sig)
.expect("v1 signature must still verify under v1 pubkey after rotation");
let v2_key_id = s.current_key_id().to_string();
let res = s.verify(&v2_key_id, msg, &sig);
assert!(matches!(res, Err(VerifyError::Invalid)));
}
#[test]
fn wrong_key_id_returns_bad_public_key() {
let s = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
let msg = b"x";
let sig = SignerBackend::sign_hex(&s, msg);
let res = s.verify("audit-mock-v999", msg, &sig);
assert!(matches!(res, Err(VerifyError::BadPublicKey)));
}
#[test]
fn from_versions_reconstructs_same_keyring_after_rotations() {
let mut grown = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
for _ in 0..3 {
grown.rotate();
}
let restored =
MockKmsSigner::from_versions("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"), 4);
assert_eq!(grown.versions(), restored.versions());
assert_eq!(grown.current_version(), restored.current_version());
}
#[test]
fn current_as_dev_signer_produces_compatible_signatures() {
let s = MockKmsSigner::new("audit-mock", [42u8; 32], ts("2026-04-28T00:00:00Z"));
let dev = s.current_as_dev_signer();
assert_eq!(dev.current_key_id(), s.current_key_id());
let msg = b"compat";
let sig = dev.sign_hex(msg);
s.verify(s.current_key_id(), msg, &sig)
.expect("DevSigner-produced signature must verify under MockKms keyring");
}
use crate::audit::{AuditEvent, SignedAuditEvent, ZERO_HASH};
use crate::decision_token::{DecisionPayload, TxTemplate};
use crate::receipt::{Decision, UnsignedReceipt};
fn unsigned_receipt(audit_event_id: &str) -> UnsignedReceipt {
UnsignedReceipt {
agent_id: "research-agent-01".to_string(),
decision: Decision::Allow,
deny_code: None,
request_hash: "1111111111111111111111111111111111111111111111111111111111111111"
.to_string(),
policy_hash: "2222222222222222222222222222222222222222222222222222222222222222"
.to_string(),
policy_version: Some(1),
audit_event_id: audit_event_id.to_string(),
execution_ref: None,
issued_at: ts("2026-04-27T12:00:01.500Z"),
expires_at: None,
}
}
fn audit_event(seq: u64) -> AuditEvent {
AuditEvent {
version: 1,
seq,
id: format!("evt-mock-kms-{seq:03}"),
ts: ts("2026-04-27T12:00:01Z"),
event_type: "policy_decided".to_string(),
actor: "policy_engine".to_string(),
subject_id: format!("pr-mock-{seq:03}"),
payload_hash: ZERO_HASH.to_string(),
metadata: serde_json::Map::new(),
policy_version: Some(1),
policy_hash: None,
attestation_ref: None,
prev_event_hash: ZERO_HASH.to_string(),
}
}
fn decision_payload(request_hash: &str) -> DecisionPayload {
DecisionPayload {
version: 1,
request_hash: request_hash.to_string(),
decision: Decision::Allow,
deny_code: None,
policy_version: 1,
policy_hash: "2222222222222222222222222222222222222222222222222222222222222222"
.to_string(),
tx_template: TxTemplate {
chain_id: 8453,
to: "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913".to_string(),
value: "0".to_string(),
data: "0x".to_string(),
gas_limit: 100_000,
max_fee_per_gas: None,
max_priority_fee_per_gas: None,
nonce_hint: None,
},
key_id: "agent-research-01-key".to_string(),
decision_id: "dec-mock-kms-001".to_string(),
issued_at: ts("2026-04-27T12:00:00Z"),
expires_at: ts("2026-04-27T12:05:00Z"),
attestation_ref: None,
}
}
#[test]
fn receipt_round_trip_through_mock_kms() {
let s = MockKmsSigner::new("decision-mock", [7u8; 32], ts("2026-04-28T00:00:00Z"));
let receipt = unsigned_receipt("evt-mock-kms-001").sign(&s).unwrap();
assert_eq!(receipt.signature.key_id, s.current_key_id());
receipt
.verify(&s.current_public_hex())
.expect("receipt must verify under the MockKms keyring's current pubkey");
}
#[test]
fn audit_event_round_trip_through_mock_kms() {
let s = MockKmsSigner::new("audit-mock", [11u8; 32], ts("2026-04-28T00:00:00Z"));
let signed = SignedAuditEvent::sign(audit_event(1), &s).unwrap();
assert_eq!(signed.signature.key_id, s.current_key_id());
signed
.verify_signature(&s.current_public_hex())
.expect("audit event must verify under MockKms keyring");
}
#[test]
fn decision_token_round_trip_through_mock_kms() {
let s = MockKmsSigner::new("decision-mock", [9u8; 32], ts("2026-04-28T00:00:00Z"));
let payload =
decision_payload("c0bd2fab4a7d4686d686edcc9c8356315cd66b820a2072493bf758a1eeb500db");
let token = payload.sign(&s).unwrap();
assert_eq!(token.signing_pubkey_hex, s.current_public_hex());
token.verify().expect("decision token must verify");
}
#[test]
fn pre_rotation_receipt_still_verifies_after_rotation() {
let mut s = MockKmsSigner::new("decision-mock", [7u8; 32], ts("2026-04-28T00:00:00Z"));
let v1_pub = s.current_public_hex();
let receipt_v1 = unsigned_receipt("evt-mock-kms-001").sign(&s).unwrap();
assert_eq!(receipt_v1.signature.key_id, "decision-mock-v1");
s.rotate();
let v2_pub = s.current_public_hex();
assert_ne!(v1_pub, v2_pub);
let resolved = s
.key_by_id(&receipt_v1.signature.key_id)
.expect("keyring still knows about v1");
receipt_v1
.verify(&resolved.public_hex)
.expect("pre-rotation receipt must verify under the resolved v1 pubkey");
let receipt_v2 = unsigned_receipt("evt-mock-kms-002").sign(&s).unwrap();
assert_eq!(receipt_v2.signature.key_id, "decision-mock-v2");
receipt_v2
.verify(&v2_pub)
.expect("v2 receipt must verify under v2 pubkey");
assert!(receipt_v1.verify(&v2_pub).is_err());
}
}