aegis-delegate 0.2.0

// AEGIS Delegate — Delegation Store Trait Abstraction
//
// Reference: AEGIS Specification v1.0.0 §9.1
//
// Defines a pluggable storage interface for delegation persistence.
// Includes an in-memory implementation for development and testing.

use std::collections::HashMap;
use std::sync::RwLock;

use async_trait::async_trait;

use openagent_aegis_core::{Delegation, DelegationError, Pagination};

// ---------------------------------------------------------------------------
// DelegationStore Trait
// ---------------------------------------------------------------------------

/// Pluggable storage backend for delegation records.
///
/// Implementations may use in-memory storage, databases, or distributed
/// caches. All operations are async to accommodate network-backed stores.
///
/// This trait provides CRUD operations on individual delegations, separate
/// from the tree-based structural operations in [`crate::tree::DelegationTree`].
#[async_trait]
pub trait DelegationStore: Send + Sync {
    /// Persist a delegation. Overwrites any existing delegation with the same ID.
    async fn store_delegation(&self, delegation: &Delegation) -> Result<(), DelegationError>;

    /// Retrieve a delegation by its unique identifier.
    ///
    /// Returns `None` if no delegation with the given ID exists.
    async fn get_delegation(&self, id: &str) -> Result<Option<Delegation>, DelegationError>;

    /// List delegations granted by a specific delegator DID with pagination.
    async fn list_by_delegator(
        &self,
        delegator_did: &str,
        pagination: Pagination,
    ) -> Result<Vec<Delegation>, DelegationError>;

    /// List delegations received by a specific delegate DID with pagination.
    async fn list_by_delegate(
        &self,
        delegate_did: &str,
        pagination: Pagination,
    ) -> Result<Vec<Delegation>, DelegationError>;

    /// Delete a delegation by ID.
    ///
    /// Returns `true` if a delegation was found and removed, `false` if no
    /// delegation with the given ID existed.
    async fn delete_delegation(&self, id: &str) -> Result<bool, DelegationError>;
}

// ---------------------------------------------------------------------------
// InMemoryDelegationStore
// ---------------------------------------------------------------------------

/// In-memory delegation store backed by a `RwLock<HashMap>`.
///
/// Suitable for development, testing, and single-instance deployments.
/// For production multi-node deployments, use a database-backed
/// implementation of [`DelegationStore`].
pub struct InMemoryDelegationStore {
    delegations: RwLock<HashMap<String, Delegation>>,
}

impl InMemoryDelegationStore {
    /// Creates a new empty in-memory delegation store.
    pub fn new() -> Self {
        Self {
            delegations: RwLock::new(HashMap::new()),
        }
    }
}

impl Default for InMemoryDelegationStore {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl DelegationStore for InMemoryDelegationStore {
    async fn store_delegation(&self, delegation: &Delegation) -> Result<(), DelegationError> {
        let mut store = self
            .delegations
            .write()
            .map_err(|e| DelegationError::InvalidProof {
                reason: format!("delegation store lock poisoned: {e}"),
            })?;
        store.insert(delegation.id.clone(), delegation.clone());
        Ok(())
    }

    async fn get_delegation(&self, id: &str) -> Result<Option<Delegation>, DelegationError> {
        let store = self
            .delegations
            .read()
            .map_err(|e| DelegationError::InvalidProof {
                reason: format!("delegation store lock poisoned: {e}"),
            })?;
        Ok(store.get(id).cloned())
    }

    async fn list_by_delegator(
        &self,
        delegator_did: &str,
        pagination: Pagination,
    ) -> Result<Vec<Delegation>, DelegationError> {
        let store = self
            .delegations
            .read()
            .map_err(|e| DelegationError::InvalidProof {
                reason: format!("delegation store lock poisoned: {e}"),
            })?;
        let matching: Vec<Delegation> = store
            .values()
            .filter(|d| d.delegator == delegator_did)
            .skip(pagination.offset as usize)
            .take(pagination.limit as usize)
            .cloned()
            .collect();
        Ok(matching)
    }

    async fn list_by_delegate(
        &self,
        delegate_did: &str,
        pagination: Pagination,
    ) -> Result<Vec<Delegation>, DelegationError> {
        let store = self
            .delegations
            .read()
            .map_err(|e| DelegationError::InvalidProof {
                reason: format!("delegation store lock poisoned: {e}"),
            })?;
        let matching: Vec<Delegation> = store
            .values()
            .filter(|d| d.delegate == delegate_did)
            .skip(pagination.offset as usize)
            .take(pagination.limit as usize)
            .cloned()
            .collect();
        Ok(matching)
    }

    async fn delete_delegation(&self, id: &str) -> Result<bool, DelegationError> {
        let mut store = self
            .delegations
            .write()
            .map_err(|e| DelegationError::InvalidProof {
                reason: format!("delegation store lock poisoned: {e}"),
            })?;
        Ok(store.remove(id).is_some())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use openagent_aegis_core::{DelegationProof, DelegationScope};
    use chrono::Utc;

    fn make_delegation(id: &str, delegator: &str, delegate: &str) -> Delegation {
        Delegation {
            id: id.to_string(),
            delegator: delegator.to_string(),
            delegate: delegate.to_string(),
            scope: DelegationScope {
                actions: vec!["transfer".to_string()],
                resources: vec![],
                chains: vec![],
                limits: None,
                temporal: None,
            },
            created: Utc::now(),
            expires: None,
            revocable: true,
            proof: DelegationProof {
                proof_type: "AegisDelegationProof2025".to_string(),
                verification_method: format!("{delegator}#key"),
                created: Utc::now(),
                jws: "test-jws".to_string(),
            },
        }
    }

    #[tokio::test]
    async fn store_and_retrieve() {
        let store = InMemoryDelegationStore::new();
        let d = make_delegation("d1", "did:root", "did:agent");

        let result = store.store_delegation(&d).await;
        assert!(result.is_ok());

        let retrieved = store.get_delegation("d1").await;
        assert!(retrieved.is_ok());
        let retrieved = retrieved.unwrap_or(None);
        assert!(retrieved.is_some());
        if let Some(del) = retrieved {
            assert_eq!(del.id, "d1");
            assert_eq!(del.delegator, "did:root");
            assert_eq!(del.delegate, "did:agent");
        }
    }

    #[tokio::test]
    async fn get_missing_returns_none() {
        let store = InMemoryDelegationStore::new();
        let result = store.get_delegation("nonexistent").await;
        assert!(result.is_ok());
        assert!(result
            .unwrap_or(Some(make_delegation("x", "x", "x")))
            .is_none());
    }

    #[tokio::test]
    async fn list_by_delegator() {
        let store = InMemoryDelegationStore::new();
        let d1 = make_delegation("d1", "did:root", "did:a");
        let d2 = make_delegation("d2", "did:root", "did:b");
        let d3 = make_delegation("d3", "did:other", "did:c");
        let _ = store.store_delegation(&d1).await;
        let _ = store.store_delegation(&d2).await;
        let _ = store.store_delegation(&d3).await;

        let root_delegations = store
            .list_by_delegator("did:root", Pagination::default())
            .await;
        assert!(root_delegations.is_ok());
        assert_eq!(root_delegations.unwrap_or_default().len(), 2);

        let other_delegations = store
            .list_by_delegator("did:other", Pagination::default())
            .await;
        assert!(other_delegations.is_ok());
        assert_eq!(other_delegations.unwrap_or_default().len(), 1);

        let empty = store
            .list_by_delegator("did:unknown", Pagination::default())
            .await;
        assert!(empty.is_ok());
        assert!(empty.unwrap_or_default().is_empty());
    }

    #[tokio::test]
    async fn list_by_delegate() {
        let store = InMemoryDelegationStore::new();
        let d1 = make_delegation("d1", "did:root", "did:agent");
        let d2 = make_delegation("d2", "did:other", "did:agent");
        let d3 = make_delegation("d3", "did:root", "did:different");
        let _ = store.store_delegation(&d1).await;
        let _ = store.store_delegation(&d2).await;
        let _ = store.store_delegation(&d3).await;

        let agent_delegations = store
            .list_by_delegate("did:agent", Pagination::default())
            .await;
        assert!(agent_delegations.is_ok());
        assert_eq!(agent_delegations.unwrap_or_default().len(), 2);

        let different = store
            .list_by_delegate("did:different", Pagination::default())
            .await;
        assert!(different.is_ok());
        assert_eq!(different.unwrap_or_default().len(), 1);
    }

    #[tokio::test]
    async fn delete_existing() {
        let store = InMemoryDelegationStore::new();
        let d = make_delegation("d1", "did:root", "did:agent");
        let _ = store.store_delegation(&d).await;

        let deleted = store.delete_delegation("d1").await;
        assert!(deleted.is_ok());
        assert_eq!(deleted.unwrap_or(false), true);

        let after = store.get_delegation("d1").await;
        assert!(after.is_ok());
        assert!(after
            .unwrap_or(Some(make_delegation("x", "x", "x")))
            .is_none());
    }

    #[tokio::test]
    async fn delete_nonexistent_returns_false() {
        let store = InMemoryDelegationStore::new();
        let deleted = store.delete_delegation("nonexistent").await;
        assert!(deleted.is_ok());
        assert_eq!(deleted.unwrap_or(true), false);
    }

    #[tokio::test]
    async fn store_overwrites_existing() {
        let store = InMemoryDelegationStore::new();
        let d1 = make_delegation("d1", "did:root", "did:agent-a");
        let _ = store.store_delegation(&d1).await;

        let d1_updated = make_delegation("d1", "did:root", "did:agent-b");
        let _ = store.store_delegation(&d1_updated).await;

        let retrieved = store.get_delegation("d1").await;
        assert!(retrieved.is_ok());
        if let Some(del) = retrieved.unwrap_or(None) {
            assert_eq!(del.delegate, "did:agent-b");
        }
    }

    #[tokio::test]
    async fn default_constructor() {
        let store = InMemoryDelegationStore::default();
        let result = store.get_delegation("any").await;
        assert!(result.is_ok());
        assert!(result
            .unwrap_or(Some(make_delegation("x", "x", "x")))
            .is_none());
    }
}