use std::collections::HashMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::auth::Identity;
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum OwnershipError {
#[error("backend error: {message}")]
Backend { message: String },
#[error("not found: {entity}")]
NotFound { entity: String },
}
pub trait OwnershipProvider: Send + Sync + 'static {
fn owns(
&self,
identity: &Identity,
resource_type: &str,
resource_id: &str,
action: &str,
) -> bool;
fn owned_resources(&self, identity: &Identity, resource_type: &str) -> Vec<String>;
fn owns_any(&self, identity: &Identity, resource_type: &str) -> bool;
}
#[async_trait]
pub trait OwnershipStore: Send + Sync + 'static {
async fn record(
&self,
identity: &Identity,
resource_type: &str,
resource_id: &str,
) -> Result<(), OwnershipError>;
async fn revoke(&self, resource_type: &str, resource_id: &str) -> Result<(), OwnershipError>;
}
pub struct InMemoryOwnershipStore {
inner: RwLock<HashMap<(String, String), Identity>>,
}
impl InMemoryOwnershipStore {
pub fn new() -> Self {
Self {
inner: RwLock::new(HashMap::new()),
}
}
}
impl Default for InMemoryOwnershipStore {
fn default() -> Self {
Self::new()
}
}
impl OwnershipProvider for InMemoryOwnershipStore {
fn owns(
&self,
identity: &Identity,
resource_type: &str,
resource_id: &str,
_action: &str,
) -> bool {
let inner = self.inner.read().unwrap_or_else(|e| e.into_inner());
inner
.get(&(resource_type.to_string(), resource_id.to_string()))
.map(|owner| owner.id == identity.id)
.unwrap_or(false)
}
fn owned_resources(&self, identity: &Identity, resource_type: &str) -> Vec<String> {
let inner = self.inner.read().unwrap_or_else(|e| e.into_inner());
inner
.iter()
.filter(|((rt, _), owner)| rt == resource_type && owner.id == identity.id)
.map(|((_, rid), _)| rid.clone())
.collect()
}
fn owns_any(&self, identity: &Identity, resource_type: &str) -> bool {
let inner = self.inner.read().unwrap_or_else(|e| e.into_inner());
inner
.iter()
.any(|((rt, _), owner)| rt == resource_type && owner.id == identity.id)
}
}
#[async_trait]
impl OwnershipStore for InMemoryOwnershipStore {
async fn record(
&self,
identity: &Identity,
resource_type: &str,
resource_id: &str,
) -> Result<(), OwnershipError> {
let mut inner = self.inner.write().unwrap_or_else(|e| e.into_inner());
inner.insert(
(resource_type.to_string(), resource_id.to_string()),
identity.clone(),
);
Ok(())
}
async fn revoke(&self, resource_type: &str, resource_id: &str) -> Result<(), OwnershipError> {
let mut inner = self.inner.write().unwrap_or_else(|e| e.into_inner());
inner.remove(&(resource_type.to_string(), resource_id.to_string()));
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_identity(id: &str) -> Identity {
Identity {
id: id.to_string(),
scopes: vec![],
resources: HashMap::new(),
}
}
#[tokio::test]
async fn record_owns_revoke_round_trip() {
let store = InMemoryOwnershipStore::new();
let owner = make_identity("worker-a");
assert!(!store.owns(&owner, "container", "c1", "exec"));
assert!(!store.owns_any(&owner, "container"));
assert!(store.owned_resources(&owner, "container").is_empty());
store.record(&owner, "container", "c1").await.unwrap();
assert!(store.owns(&owner, "container", "c1", "exec"));
assert!(store.owns(&owner, "container", "c1", "logs"));
assert!(store.owns_any(&owner, "container"));
assert_eq!(store.owned_resources(&owner, "container"), vec!["c1"]);
store.revoke("container", "c1").await.unwrap();
assert!(!store.owns(&owner, "container", "c1", "exec"));
assert!(!store.owns_any(&owner, "container"));
assert!(store.owned_resources(&owner, "container").is_empty());
}
#[tokio::test]
async fn owned_resources_returns_all_for_owner_with_multiple() {
let store = InMemoryOwnershipStore::new();
let owner = make_identity("worker-a");
store.record(&owner, "container", "c1").await.unwrap();
store.record(&owner, "container", "c2").await.unwrap();
store.record(&owner, "container", "c3").await.unwrap();
let mut owned = store.owned_resources(&owner, "container");
owned.sort();
assert_eq!(owned, vec!["c1", "c2", "c3"]);
}
#[tokio::test]
async fn owned_resources_filters_by_resource_type() {
let store = InMemoryOwnershipStore::new();
let owner = make_identity("worker-a");
store.record(&owner, "container", "c1").await.unwrap();
store.record(&owner, "tty", "t1").await.unwrap();
let owned_containers = store.owned_resources(&owner, "container");
assert_eq!(owned_containers, vec!["c1"]);
let owned_ttys = store.owned_resources(&owner, "tty");
assert_eq!(owned_ttys, vec!["t1"]);
}
#[tokio::test]
async fn owns_any_returns_false_for_owner_with_no_resources_of_type() {
let store = InMemoryOwnershipStore::new();
let owner = make_identity("worker-a");
store.record(&owner, "container", "c1").await.unwrap();
assert!(store.owns_any(&owner, "container"));
assert!(!store.owns_any(&owner, "tty"));
}
#[tokio::test]
async fn revoke_on_non_existent_resource_is_no_op() {
let store = InMemoryOwnershipStore::new();
store.revoke("container", "never-existed").await.unwrap();
}
#[tokio::test]
async fn owns_returns_false_for_different_identity() {
let store = InMemoryOwnershipStore::new();
let owner = make_identity("worker-a");
let other = make_identity("worker-b");
store.record(&owner, "container", "c1").await.unwrap();
assert!(store.owns(&owner, "container", "c1", "exec"));
assert!(!store.owns(&other, "container", "c1", "exec"));
assert!(!store.owns_any(&other, "container"));
assert!(store.owned_resources(&other, "container").is_empty());
}
#[tokio::test]
async fn record_replaces_existing_owner() {
let store = InMemoryOwnershipStore::new();
let owner_a = make_identity("worker-a");
let owner_b = make_identity("worker-b");
store.record(&owner_a, "container", "c1").await.unwrap();
store.record(&owner_b, "container", "c1").await.unwrap();
assert!(!store.owns(&owner_a, "container", "c1", "exec"));
assert!(store.owns(&owner_b, "container", "c1", "exec"));
}
#[tokio::test]
async fn default_is_empty_store() {
let store = InMemoryOwnershipStore::default();
let owner = make_identity("worker-a");
assert!(store.owned_resources(&owner, "container").is_empty());
assert!(!store.owns_any(&owner, "container"));
}
#[test]
fn ownership_error_display_formatting() {
let backend = OwnershipError::Backend {
message: "disk full".to_string(),
};
assert_eq!(backend.to_string(), "backend error: disk full");
let not_found = OwnershipError::NotFound {
entity: "container:c1".to_string(),
};
assert_eq!(not_found.to_string(), "not found: container:c1");
}
#[test]
fn ownership_error_is_non_exhaustive() {
let err = OwnershipError::Backend {
message: "x".to_string(),
};
let _ = err.to_string();
}
}