use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use chrono::Utc;
use tokio::sync::RwLock;
use zlayer_secrets::{
ClusterDek, RaftSecretsHandle, RaftSecretsStore, RecipientPrivateKey, RecipientPublicKey,
Result as SecretsResult, Secret, SecretsError, SecretsProvider, SecretsState, SecretsStore,
};
use zlayer_types::api::internal::SecretsRaftOp;
use zlayer_types::storage::{NodeAffinity, NodeIdentity, ReplicatedSecret};
struct SharedRaftHandle {
state: Arc<RwLock<SecretsState>>,
}
impl SharedRaftHandle {
fn new(state: Arc<RwLock<SecretsState>>) -> Self {
Self { state }
}
}
#[async_trait]
impl RaftSecretsHandle for SharedRaftHandle {
async fn secrets_state(&self) -> SecretsState {
self.state.read().await.clone()
}
async fn propose_put_secret(&self, secret: ReplicatedSecret) -> SecretsResult<()> {
let mut guard = self.state.write().await;
guard
.apply(SecretsRaftOp::PutSecret { secret })
.map_err(|e| SecretsError::Provider(format!("apply PutSecret failed: {e}")))
}
async fn propose_delete_secret(&self, storage_key: &str) -> SecretsResult<()> {
let mut guard = self.state.write().await;
guard
.apply(SecretsRaftOp::DeleteSecret {
storage_key: storage_key.to_string(),
})
.map_err(|e| SecretsError::Provider(format!("apply DeleteSecret failed: {e}")))
}
}
struct ThreeNodeCluster {
a_id: String,
b_id: String,
#[allow(dead_code)]
c_id: String,
store_a: RaftSecretsStore,
store_b: RaftSecretsStore,
store_c: RaftSecretsStore,
#[allow(dead_code)]
state: Arc<RwLock<SecretsState>>,
}
fn setup_three_node_cluster() -> ThreeNodeCluster {
let a_id = "node-a".to_string();
let b_id = "node-b".to_string();
let c_id = "node-c".to_string();
let (sk_a, pk_a) = RecipientPrivateKey::generate();
let (sk_b, pk_b) = RecipientPrivateKey::generate();
let (sk_c, pk_c) = RecipientPrivateKey::generate();
let now = Utc::now();
let identity_a = NodeIdentity {
node_id: a_id.clone(),
secrets_pubkey: *pk_a.as_bytes(),
wg_pubkey: "wg-a".to_string(),
joined_at: now,
revoked_at: None,
};
let identity_b = NodeIdentity {
node_id: b_id.clone(),
secrets_pubkey: *pk_b.as_bytes(),
wg_pubkey: "wg-b".to_string(),
joined_at: now,
revoked_at: None,
};
let identity_c = NodeIdentity {
node_id: c_id.clone(),
secrets_pubkey: *pk_c.as_bytes(),
wg_pubkey: "wg-c".to_string(),
joined_at: now,
revoked_at: None,
};
let dek = ClusterDek::generate();
let mut recipients: HashMap<String, RecipientPublicKey> = HashMap::new();
recipients.insert(
a_id.clone(),
RecipientPublicKey::from_bytes(*pk_a.as_bytes()),
);
recipients.insert(
b_id.clone(),
RecipientPublicKey::from_bytes(*pk_b.as_bytes()),
);
recipients.insert(
c_id.clone(),
RecipientPublicKey::from_bytes(*pk_c.as_bytes()),
);
let envelope = dek
.rewrap_for_set(&recipients, 1)
.expect("rewrap DEK for the three-node set");
let mut initial = SecretsState::default();
initial
.apply(SecretsRaftOp::RegisterNode {
identity: identity_a,
})
.expect("seed register node-a");
initial
.apply(SecretsRaftOp::RegisterNode {
identity: identity_b,
})
.expect("seed register node-b");
initial
.apply(SecretsRaftOp::RegisterNode {
identity: identity_c,
})
.expect("seed register node-c");
initial
.apply(SecretsRaftOp::RotateDek {
new_wraps: envelope,
})
.expect("seed rotate DEK");
let state = Arc::new(RwLock::new(initial));
let handle_a: Arc<dyn RaftSecretsHandle> = Arc::new(SharedRaftHandle::new(state.clone()));
let handle_b: Arc<dyn RaftSecretsHandle> = Arc::new(SharedRaftHandle::new(state.clone()));
let handle_c: Arc<dyn RaftSecretsHandle> = Arc::new(SharedRaftHandle::new(state.clone()));
let store_a = RaftSecretsStore::new(sk_a, a_id.clone(), handle_a);
let store_b = RaftSecretsStore::new(sk_b, b_id.clone(), handle_b);
let store_c = RaftSecretsStore::new(sk_c, c_id.clone(), handle_c);
ThreeNodeCluster {
a_id,
b_id,
c_id,
store_a,
store_b,
store_c,
state,
}
}
fn assert_not_found(result: Result<Secret, SecretsError>, label: &str) {
match result {
Err(SecretsError::NotFound { .. }) => {}
Err(other) => panic!(
"{label}: expected SecretsError::NotFound (existence must not leak), got {other:?}",
),
Ok(_) => panic!("{label}: expected NotFound; got a Secret — affinity is not enforced!"),
}
}
#[tokio::test]
async fn affinity_nodes_a_only_blocks_b_and_c_from_reading() {
let cluster = setup_three_node_cluster();
cluster
.store_a
.set_secret_with_affinity(
"default",
"alpha",
&Secret::new("only-a-may-read"),
Some(&NodeAffinity::Nodes {
node_ids: vec![cluster.a_id.clone()],
}),
)
.await
.expect("write secret with affinity=[A]");
let got_a = cluster
.store_a
.get_secret("default", "alpha")
.await
.expect("A is in affinity set; must read plaintext");
assert_eq!(got_a.expose(), "only-a-may-read");
assert_not_found(
cluster.store_b.get_secret("default", "alpha").await,
"node-b read of A-affinity secret",
);
assert_not_found(
cluster.store_c.get_secret("default", "alpha").await,
"node-c read of A-affinity secret",
);
let a_exists = cluster
.store_a
.exists("default", "alpha")
.await
.expect("exists ok");
assert!(
a_exists,
"node-a is in the affinity set; exists must return true",
);
let b_exists = cluster
.store_b
.exists("default", "alpha")
.await
.expect("exists ok");
assert!(
!b_exists,
"node-b must not be told the affinity-A secret exists (exists returned true)",
);
let c_exists = cluster
.store_c
.exists("default", "alpha")
.await
.expect("exists ok");
assert!(
!c_exists,
"node-c must not be told the affinity-A secret exists (exists returned true)",
);
}
#[tokio::test]
async fn affinity_admin_token_does_not_bypass() {
let cluster = setup_three_node_cluster();
cluster
.store_a
.set_secret_with_affinity(
"default",
"alpha",
&Secret::new("only-a-even-for-admin"),
Some(&NodeAffinity::Nodes {
node_ids: vec![cluster.a_id.clone()],
}),
)
.await
.expect("write A-only secret");
assert_not_found(
cluster.store_b.get_secret("default", "alpha").await,
"node-b read with implicit-admin context",
);
let listed = cluster
.store_b
.list_secrets("default")
.await
.expect("list ok");
assert!(
listed.iter().all(|m| m.name != "alpha"),
"node-b list must not include the A-affinity secret (got: {:?})",
listed.iter().map(|m| &m.name).collect::<Vec<_>>(),
);
}
#[tokio::test]
async fn affinity_label_selector_allow_all_until_phase_1_5() {
let cluster = setup_three_node_cluster();
let mut labels = HashMap::new();
labels.insert("tier".to_string(), "prod".to_string());
cluster
.store_a
.set_secret_with_affinity(
"default",
"labelled",
&Secret::new("phase-1-allow-all"),
Some(&NodeAffinity::Labels { labels }),
)
.await
.expect("write secret with label selector");
for (label, store) in [
("A", &cluster.store_a),
("B", &cluster.store_b),
("C", &cluster.store_c),
] {
let got = store
.get_secret("default", "labelled")
.await
.unwrap_or_else(|e| {
panic!(
"node-{label}: Phase-1 label selectors are conservatively \
allow-all at the store layer; got error {e:?}",
);
});
assert_eq!(
got.expose(),
"phase-1-allow-all",
"node-{label} should read the labelled secret in Phase 1",
);
}
}
#[tokio::test]
async fn affinity_none_allows_all_nodes() {
let cluster = setup_three_node_cluster();
cluster
.store_a
.set_secret("default", "open", &Secret::new("anyone-may-read"))
.await
.expect("write unrestricted secret");
for (label, store) in [
("A", &cluster.store_a),
("B", &cluster.store_b),
("C", &cluster.store_c),
] {
let got = store
.get_secret("default", "open")
.await
.unwrap_or_else(|e| panic!("node-{label}: unrestricted read failed: {e:?}"));
assert_eq!(got.expose(), "anyone-may-read", "node-{label} read");
}
}
#[tokio::test]
async fn affinity_listing_filters_secrets_per_node() {
let cluster = setup_three_node_cluster();
cluster
.store_a
.set_secret_with_affinity(
"default",
"alpha",
&Secret::new("alpha-value"),
Some(&NodeAffinity::Nodes {
node_ids: vec![cluster.a_id.clone()],
}),
)
.await
.expect("write alpha");
cluster
.store_a
.set_secret_with_affinity("default", "beta", &Secret::new("beta-value"), None)
.await
.expect("write beta");
cluster
.store_a
.set_secret_with_affinity(
"default",
"gamma",
&Secret::new("gamma-value"),
Some(&NodeAffinity::Nodes {
node_ids: vec![cluster.a_id.clone(), cluster.b_id.clone()],
}),
)
.await
.expect("write gamma");
let names_for = |list: Vec<zlayer_secrets::SecretMetadata>| -> Vec<String> {
let mut names: Vec<String> = list.into_iter().map(|m| m.name).collect();
names.sort();
names
};
let list_a = cluster
.store_a
.list_secrets("default")
.await
.expect("list A");
assert_eq!(
names_for(list_a),
vec!["alpha".to_string(), "beta".to_string(), "gamma".to_string()],
"node-a should see all three secrets (alpha, beta, gamma)",
);
let list_b = cluster
.store_b
.list_secrets("default")
.await
.expect("list B");
assert_eq!(
names_for(list_b),
vec!["beta".to_string(), "gamma".to_string()],
"node-b should see beta and gamma (alpha is A-only)",
);
let list_c = cluster
.store_c
.list_secrets("default")
.await
.expect("list C");
assert_eq!(
names_for(list_c),
vec!["beta".to_string()],
"node-c should see only beta (alpha is A-only, gamma is A+B)",
);
let beta_on_c = cluster
.store_c
.get_secret("default", "beta")
.await
.expect("C reads unrestricted beta");
assert_eq!(beta_on_c.expose(), "beta-value");
assert_not_found(
cluster.store_c.get_secret("default", "alpha").await,
"node-c read of alpha (A-only)",
);
assert_not_found(
cluster.store_c.get_secret("default", "gamma").await,
"node-c read of gamma (A+B-only)",
);
}