1use std::collections::HashMap;
11
12use chrono::Utc;
13use serde::{Deserialize, Serialize};
14
15use zlayer_types::api::internal::SecretsRaftOp;
16use zlayer_types::storage::{NodeIdentity, ReplicatedSecret, WrappedDek};
17
18use crate::SecretsError;
19
20#[derive(Debug, Default, Clone, Serialize, Deserialize)]
25pub struct SecretsState {
26 pub nodes: HashMap<String, NodeIdentity>,
31
32 pub wrapped_dek: Option<WrappedDek>,
35
36 pub secrets: HashMap<String, ReplicatedSecret>,
38
39 pub revoked_tokens: HashMap<String, chrono::DateTime<chrono::Utc>>,
43
44 #[serde(default)]
53 pub trusted_bundles: HashMap<String, zlayer_types::api::cluster::TrustBundle>,
54
55 #[serde(default)]
62 pub jwt_algorithm: zlayer_types::api::cluster::JwtAlgorithm,
63
64 #[serde(default)]
79 pub join_secret_wiped_at: Option<chrono::DateTime<chrono::Utc>>,
80}
81
82impl SecretsState {
83 pub fn apply(&mut self, op: SecretsRaftOp) -> Result<(), SecretsError> {
95 match op {
96 SecretsRaftOp::RegisterNode { identity } => {
97 self.nodes.insert(identity.node_id.clone(), identity);
99 Ok(())
100 }
101 SecretsRaftOp::RevokeNode { node_id } => {
102 let entry = self.nodes.get_mut(&node_id).ok_or_else(|| {
103 SecretsError::Provider(format!("RevokeNode for unknown node_id: {node_id}"))
104 })?;
105 if entry.revoked_at.is_none() {
106 entry.revoked_at = Some(Utc::now());
107 }
108 Ok(())
109 }
110 SecretsRaftOp::RotateDek { new_wraps } => {
111 self.wrapped_dek = Some(new_wraps);
116 Ok(())
117 }
118 SecretsRaftOp::PutSecret { secret } => {
119 self.secrets.insert(secret.storage_key.clone(), secret);
120 Ok(())
121 }
122 SecretsRaftOp::DeleteSecret { storage_key } => {
123 self.secrets.remove(&storage_key).ok_or_else(|| {
124 SecretsError::Provider(format!(
125 "DeleteSecret for unknown storage_key: {storage_key}"
126 ))
127 })?;
128 Ok(())
129 }
130 SecretsRaftOp::RevokeToken {
131 token_hash,
132 expires_at,
133 } => {
134 let now = Utc::now();
138 if expires_at > now {
139 self.revoked_tokens.insert(token_hash, expires_at);
140 }
141 self.revoked_tokens.retain(|_, exp| *exp > now);
143 Ok(())
144 }
145 SecretsRaftOp::ImportTrustBundle { bundle } => {
146 self.trusted_bundles
152 .insert(bundle.cluster_domain.clone(), bundle);
153 Ok(())
154 }
155 SecretsRaftOp::RemoveTrustBundle { cluster_domain } => {
156 self.trusted_bundles.remove(&cluster_domain);
159 Ok(())
160 }
161 SecretsRaftOp::SetJwtAlgorithm { algorithm } => {
162 self.jwt_algorithm = algorithm;
163 Ok(())
164 }
165 SecretsRaftOp::WipeJoinSecret => {
166 if self.join_secret_wiped_at.is_none() {
176 self.join_secret_wiped_at = Some(Utc::now());
177 }
178 Ok(())
179 }
180 }
181 }
182
183 pub fn snapshot(&self) -> Result<Vec<u8>, SecretsError> {
190 serde_json::to_vec(self).map_err(|e| SecretsError::Storage(format!("snapshot: {e}")))
191 }
192
193 pub fn restore(bytes: &[u8]) -> Result<Self, SecretsError> {
198 serde_json::from_slice(bytes).map_err(|e| SecretsError::Storage(format!("restore: {e}")))
199 }
200
201 #[must_use]
204 pub fn node_can_decrypt(&self, node_id: &str) -> bool {
205 self.wrapped_dek
206 .as_ref()
207 .is_some_and(|w| w.wraps.contains_key(node_id))
208 }
209
210 #[must_use]
215 pub fn token_revoked(&self, token_hash: &str) -> bool {
216 self.revoked_tokens.contains_key(token_hash)
217 }
218
219 #[must_use]
221 pub fn trust_bundle_for(
222 &self,
223 cluster_domain: &str,
224 ) -> Option<&zlayer_types::api::cluster::TrustBundle> {
225 self.trusted_bundles.get(cluster_domain)
226 }
227
228 #[must_use]
230 pub fn jwt_algorithm(&self) -> zlayer_types::api::cluster::JwtAlgorithm {
231 self.jwt_algorithm
232 }
233
234 #[must_use]
236 pub fn join_secret_wiped(&self) -> bool {
237 self.join_secret_wiped_at.is_some()
238 }
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244 use chrono::TimeZone;
245 use zlayer_types::secrets::SecretMetadata;
246
247 fn make_identity(node_id: &str) -> NodeIdentity {
248 NodeIdentity {
249 node_id: node_id.to_string(),
250 secrets_pubkey: [0u8; 32],
251 wg_pubkey: format!("wg-{node_id}"),
252 joined_at: Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap(),
253 revoked_at: None,
254 }
255 }
256
257 fn make_wrapped_dek(generation: u64, node_ids: &[&str]) -> WrappedDek {
258 let mut wraps = HashMap::new();
259 for nid in node_ids {
260 wraps.insert((*nid).to_string(), vec![0xAB, 0xCD]);
261 }
262 WrappedDek {
263 dek_generation: generation,
264 wraps,
265 }
266 }
267
268 fn make_secret(name: &str, generation: u64) -> ReplicatedSecret {
269 ReplicatedSecret {
270 storage_key: format!("dep:{name}"),
271 ciphertext: vec![1, 2, 3, 4],
272 dek_generation: generation,
273 metadata: SecretMetadata::new(name),
274 node_affinity: None,
275 }
276 }
277
278 #[test]
279 fn apply_register_node_inserts() {
280 let mut state = SecretsState::default();
281 state
282 .apply(SecretsRaftOp::RegisterNode {
283 identity: make_identity("node-a"),
284 })
285 .expect("register should succeed");
286 assert_eq!(state.nodes.len(), 1);
287 assert!(state.nodes.contains_key("node-a"));
288 }
289
290 #[test]
291 fn apply_register_node_overwrites_existing() {
292 let mut state = SecretsState::default();
293 let mut first = make_identity("node-a");
294 first.wg_pubkey = "wg-original".to_string();
295 state
296 .apply(SecretsRaftOp::RegisterNode { identity: first })
297 .expect("first register");
298
299 let mut second = make_identity("node-a");
300 second.wg_pubkey = "wg-replaced".to_string();
301 state
302 .apply(SecretsRaftOp::RegisterNode { identity: second })
303 .expect("second register should not error");
304
305 assert_eq!(state.nodes.len(), 1);
306 assert_eq!(state.nodes["node-a"].wg_pubkey, "wg-replaced");
307 }
308
309 #[test]
310 fn apply_revoke_node_marks_revoked_at() {
311 let mut state = SecretsState::default();
312 state
313 .apply(SecretsRaftOp::RegisterNode {
314 identity: make_identity("node-a"),
315 })
316 .expect("register");
317 state
318 .apply(SecretsRaftOp::RevokeNode {
319 node_id: "node-a".to_string(),
320 })
321 .expect("revoke");
322 assert!(state.nodes["node-a"].revoked_at.is_some());
323
324 let original_ts = state.nodes["node-a"].revoked_at;
327 state
328 .apply(SecretsRaftOp::RevokeNode {
329 node_id: "node-a".to_string(),
330 })
331 .expect("revoke again");
332 assert_eq!(state.nodes["node-a"].revoked_at, original_ts);
333 }
334
335 #[test]
336 fn apply_revoke_unknown_node_errors() {
337 let mut state = SecretsState::default();
338 let err = state
339 .apply(SecretsRaftOp::RevokeNode {
340 node_id: "missing".to_string(),
341 })
342 .expect_err("revoke unknown should fail");
343 assert!(matches!(err, SecretsError::Provider(_)), "got: {err:?}");
344 }
345
346 #[test]
347 fn apply_rotate_dek_replaces_wraps() {
348 let mut state = SecretsState::default();
349 state
350 .apply(SecretsRaftOp::RotateDek {
351 new_wraps: make_wrapped_dek(1, &["node-a"]),
352 })
353 .expect("rotate 1");
354 state
355 .apply(SecretsRaftOp::RotateDek {
356 new_wraps: make_wrapped_dek(2, &["node-a", "node-b"]),
357 })
358 .expect("rotate 2");
359 let dek = state.wrapped_dek.as_ref().expect("dek present");
360 assert_eq!(dek.dek_generation, 2);
361 assert_eq!(dek.wraps.len(), 2);
362 assert!(dek.wraps.contains_key("node-a"));
363 assert!(dek.wraps.contains_key("node-b"));
364 }
365
366 #[test]
367 fn apply_put_secret_inserts_then_overwrites() {
368 let mut state = SecretsState::default();
369 let mut first = make_secret("api-key", 1);
370 first.ciphertext = vec![0xDE, 0xAD];
371 state
372 .apply(SecretsRaftOp::PutSecret {
373 secret: first.clone(),
374 })
375 .expect("put 1");
376 assert_eq!(state.secrets.len(), 1);
377 assert_eq!(
378 state.secrets[&first.storage_key].ciphertext,
379 vec![0xDE, 0xAD]
380 );
381
382 let mut second = make_secret("api-key", 2);
383 second.ciphertext = vec![0xBE, 0xEF];
384 state
385 .apply(SecretsRaftOp::PutSecret {
386 secret: second.clone(),
387 })
388 .expect("put 2");
389 assert_eq!(state.secrets.len(), 1);
390 assert_eq!(
391 state.secrets[&second.storage_key].ciphertext,
392 vec![0xBE, 0xEF]
393 );
394 assert_eq!(state.secrets[&second.storage_key].dek_generation, 2);
395 }
396
397 #[test]
398 fn apply_delete_secret_removes() {
399 let mut state = SecretsState::default();
400 let secret = make_secret("api-key", 1);
401 let key = secret.storage_key.clone();
402 state
403 .apply(SecretsRaftOp::PutSecret { secret })
404 .expect("put");
405 state
406 .apply(SecretsRaftOp::DeleteSecret {
407 storage_key: key.clone(),
408 })
409 .expect("delete");
410 assert!(state.secrets.is_empty());
411 }
412
413 #[test]
414 fn apply_delete_unknown_secret_errors() {
415 let mut state = SecretsState::default();
416 let err = state
417 .apply(SecretsRaftOp::DeleteSecret {
418 storage_key: "dep:nope".to_string(),
419 })
420 .expect_err("delete unknown should fail");
421 assert!(matches!(err, SecretsError::Provider(_)), "got: {err:?}");
422 }
423
424 #[test]
425 fn snapshot_round_trip() {
426 let mut state = SecretsState::default();
427 state
428 .apply(SecretsRaftOp::RegisterNode {
429 identity: make_identity("node-a"),
430 })
431 .expect("register a");
432 state
433 .apply(SecretsRaftOp::RegisterNode {
434 identity: make_identity("node-b"),
435 })
436 .expect("register b");
437 state
438 .apply(SecretsRaftOp::RotateDek {
439 new_wraps: make_wrapped_dek(7, &["node-a", "node-b"]),
440 })
441 .expect("rotate");
442 state
443 .apply(SecretsRaftOp::PutSecret {
444 secret: make_secret("api-key", 7),
445 })
446 .expect("put");
447 state
448 .apply(SecretsRaftOp::RevokeNode {
449 node_id: "node-b".to_string(),
450 })
451 .expect("revoke b");
452
453 let bytes = state.snapshot().expect("snapshot ok");
454 let restored = SecretsState::restore(&bytes).expect("restore ok");
455
456 let bytes2 = restored.snapshot().expect("snapshot restored ok");
461 let v1: serde_json::Value = serde_json::from_slice(&bytes).expect("parse v1");
462 let v2: serde_json::Value = serde_json::from_slice(&bytes2).expect("parse v2");
463 assert_eq!(v1, v2);
464
465 assert_eq!(restored.nodes.len(), state.nodes.len());
467 assert_eq!(restored.secrets.len(), state.secrets.len());
468 assert_eq!(
469 restored.wrapped_dek.as_ref().map(|w| w.dek_generation),
470 state.wrapped_dek.as_ref().map(|w| w.dek_generation),
471 );
472 }
473
474 #[test]
475 fn node_can_decrypt_reflects_wraps() {
476 let mut state = SecretsState::default();
477 assert!(!state.node_can_decrypt("node-a"));
478
479 state
480 .apply(SecretsRaftOp::RotateDek {
481 new_wraps: make_wrapped_dek(1, &["node-a"]),
482 })
483 .expect("rotate include");
484 assert!(state.node_can_decrypt("node-a"));
485 assert!(!state.node_can_decrypt("node-b"));
486
487 state
488 .apply(SecretsRaftOp::RotateDek {
489 new_wraps: make_wrapped_dek(2, &["node-b"]),
490 })
491 .expect("rotate exclude a");
492 assert!(!state.node_can_decrypt("node-a"));
493 assert!(state.node_can_decrypt("node-b"));
494 }
495
496 #[test]
497 fn revoke_token_inserts_entry() {
498 let mut state = SecretsState::default();
499 let expires_at = Utc::now() + chrono::Duration::hours(24);
500 state
501 .apply(SecretsRaftOp::RevokeToken {
502 token_hash: "abc123".to_string(),
503 expires_at,
504 })
505 .unwrap();
506 assert!(state.token_revoked("abc123"));
507 assert!(!state.token_revoked("def456"));
508 }
509
510 #[test]
511 fn revoke_token_is_idempotent() {
512 let mut state = SecretsState::default();
513 let expires_at = Utc::now() + chrono::Duration::hours(24);
514 let op = SecretsRaftOp::RevokeToken {
515 token_hash: "abc123".to_string(),
516 expires_at,
517 };
518 state.apply(op.clone()).unwrap();
519 state.apply(op).unwrap();
520 assert_eq!(state.revoked_tokens.len(), 1);
521 }
522
523 #[test]
524 fn revoke_token_skips_already_expired_input() {
525 let mut state = SecretsState::default();
526 let expired_at = Utc::now() - chrono::Duration::hours(1);
527 state
528 .apply(SecretsRaftOp::RevokeToken {
529 token_hash: "abc123".to_string(),
530 expires_at: expired_at,
531 })
532 .unwrap();
533 assert!(!state.token_revoked("abc123"));
535 }
536
537 #[test]
538 fn revoke_token_apply_prunes_expired_neighbors() {
539 let mut state = SecretsState::default();
540 let expired_at = Utc::now() - chrono::Duration::hours(1);
543 state.revoked_tokens.insert("stale".to_string(), expired_at);
544 let fresh_expires = Utc::now() + chrono::Duration::hours(24);
545 state
546 .apply(SecretsRaftOp::RevokeToken {
547 token_hash: "fresh".to_string(),
548 expires_at: fresh_expires,
549 })
550 .unwrap();
551 assert!(state.token_revoked("fresh"));
552 assert!(!state.token_revoked("stale"));
553 }
554
555 fn make_trust_bundle(
556 cluster_domain: &str,
557 ca_kid: &str,
558 ) -> zlayer_types::api::cluster::TrustBundle {
559 zlayer_types::api::cluster::TrustBundle {
560 v: zlayer_types::api::cluster::TRUST_BUNDLE_FORMAT_VERSION,
561 cluster_domain: cluster_domain.to_string(),
562 ca_public_key_b64: format!("pubkey-of-{cluster_domain}"),
563 ca_kid: ca_kid.to_string(),
564 generated_at: chrono::Utc::now().to_rfc3339(),
565 }
566 }
567
568 #[test]
569 fn import_trust_bundle_inserts_entry() {
570 let mut state = SecretsState::default();
571 let bundle = make_trust_bundle("prod-east", "deadbeef");
572 state
573 .apply(SecretsRaftOp::ImportTrustBundle {
574 bundle: bundle.clone(),
575 })
576 .unwrap();
577 let got = state
578 .trust_bundle_for("prod-east")
579 .expect("must be present");
580 assert_eq!(got.cluster_domain, "prod-east");
581 assert_eq!(got.ca_kid, "deadbeef");
582 assert!(state.trust_bundle_for("prod-west").is_none());
583 }
584
585 #[test]
586 fn import_trust_bundle_is_idempotent_overwriting_in_place() {
587 let mut state = SecretsState::default();
588 state
589 .apply(SecretsRaftOp::ImportTrustBundle {
590 bundle: make_trust_bundle("prod-east", "deadbeef"),
591 })
592 .unwrap();
593 state
594 .apply(SecretsRaftOp::ImportTrustBundle {
595 bundle: make_trust_bundle("prod-east", "newkid12"),
596 })
597 .unwrap();
598 let got = state.trust_bundle_for("prod-east").unwrap();
599 assert_eq!(got.ca_kid, "newkid12", "re-import must overwrite in place");
600 assert_eq!(state.trusted_bundles.len(), 1);
601 }
602
603 #[test]
604 fn remove_trust_bundle_drops_entry() {
605 let mut state = SecretsState::default();
606 state
607 .apply(SecretsRaftOp::ImportTrustBundle {
608 bundle: make_trust_bundle("prod-east", "deadbeef"),
609 })
610 .unwrap();
611 state
612 .apply(SecretsRaftOp::RemoveTrustBundle {
613 cluster_domain: "prod-east".into(),
614 })
615 .unwrap();
616 assert!(state.trust_bundle_for("prod-east").is_none());
617 }
618
619 #[test]
620 fn remove_trust_bundle_is_idempotent_for_unknown_domain() {
621 let mut state = SecretsState::default();
622 state
623 .apply(SecretsRaftOp::RemoveTrustBundle {
624 cluster_domain: "never-imported".into(),
625 })
626 .unwrap();
627 }
629
630 #[test]
631 fn set_jwt_algorithm_default_is_both() {
632 let state = SecretsState::default();
633 assert_eq!(
634 state.jwt_algorithm(),
635 zlayer_types::api::cluster::JwtAlgorithm::Both,
636 "default policy is Both for safety during migration"
637 );
638 }
639
640 #[test]
641 fn set_jwt_algorithm_flips_policy() {
642 let mut state = SecretsState::default();
643 state
644 .apply(SecretsRaftOp::SetJwtAlgorithm {
645 algorithm: zlayer_types::api::cluster::JwtAlgorithm::Eddsa,
646 })
647 .unwrap();
648 assert_eq!(
649 state.jwt_algorithm(),
650 zlayer_types::api::cluster::JwtAlgorithm::Eddsa
651 );
652 }
653
654 #[test]
655 fn set_jwt_algorithm_is_idempotent() {
656 let mut state = SecretsState::default();
657 state
658 .apply(SecretsRaftOp::SetJwtAlgorithm {
659 algorithm: zlayer_types::api::cluster::JwtAlgorithm::Hs256,
660 })
661 .unwrap();
662 state
663 .apply(SecretsRaftOp::SetJwtAlgorithm {
664 algorithm: zlayer_types::api::cluster::JwtAlgorithm::Hs256,
665 })
666 .unwrap();
667 assert_eq!(
668 state.jwt_algorithm(),
669 zlayer_types::api::cluster::JwtAlgorithm::Hs256
670 );
671 }
672
673 #[test]
674 fn wipe_join_secret_records_timestamp() {
675 let mut state = SecretsState::default();
676 assert!(!state.join_secret_wiped());
677 state.apply(SecretsRaftOp::WipeJoinSecret).unwrap();
678 assert!(state.join_secret_wiped());
679 assert!(state.join_secret_wiped_at.is_some());
680 }
681
682 #[test]
683 fn wipe_join_secret_is_idempotent_preserves_first_timestamp() {
684 let mut state = SecretsState::default();
685 state.apply(SecretsRaftOp::WipeJoinSecret).unwrap();
686 let first = state.join_secret_wiped_at.unwrap();
687 state.apply(SecretsRaftOp::WipeJoinSecret).unwrap();
689 assert_eq!(state.join_secret_wiped_at.unwrap(), first);
690 }
691}