aegis_delegate/

tree.rs

// AEGIS Delegate — Delegation Trees
//
// Reference: AEGIS Specification v1.0.0 §9.2
//
// Maintains a delegation tree rooted at human roots with configurable
// maximum depth. Enforces no-amplification via scope subsetting.

use openagent_aegis_core::{Delegation, DelegationError, DelegationScope};
use std::collections::HashMap;

use crate::scope::{intersect_scopes, is_scope_subset};

/// A delegation tree rooted at a human root.
///
/// Stores delegations indexed by delegate DID and by delegation ID.
/// Max depth is configurable (default 8).
pub struct DelegationTree {
    /// All delegations indexed by delegate DID.
    delegations: HashMap<String, Vec<Delegation>>,
    /// Index from delegation ID to (delegator, delegate) for quick lookup.
    id_index: HashMap<String, (String, String)>,
    /// Maximum delegation depth.
    max_depth: u32,
}

impl DelegationTree {
    /// Creates a new empty delegation tree with the given maximum depth.
    pub fn new(max_depth: u32) -> Self {
        Self {
            delegations: HashMap::new(),
            id_index: HashMap::new(),
            max_depth,
        }
    }

    /// Adds a delegation to the tree.
    ///
    /// Validates:
    /// - The delegation depth does not exceed max_depth
    /// - The delegation scope is a subset of the delegator's effective scope
    ///   (no-amplification rule, §9.6)
    pub fn add_delegation(&mut self, delegation: Delegation) -> Result<(), DelegationError> {
        // Calculate the depth for the new delegation
        let current_depth = self.depth(&delegation.delegator);
        let new_depth = current_depth + 1;

        if new_depth > self.max_depth {
            return Err(DelegationError::MaxDepthExceeded {
                depth: new_depth,
                max_depth: self.max_depth,
            });
        }

        // Check scope subset if the delegator is itself a delegate (not a root)
        if let Some(parent_scope) = self.effective_scope(&delegation.delegator) {
            if !is_scope_subset(&delegation.scope, &parent_scope) {
                return Err(DelegationError::ScopeAmplification {
                    reason: format!(
                        "delegate {} scope exceeds delegator {} effective scope",
                        delegation.delegate, delegation.delegator,
                    ),
                });
            }
        }

        // Index by ID
        self.id_index.insert(
            delegation.id.clone(),
            (delegation.delegator.clone(), delegation.delegate.clone()),
        );

        // Index by delegate DID
        self.delegations
            .entry(delegation.delegate.clone())
            .or_default()
            .push(delegation);

        Ok(())
    }

    /// Gets the delegation chain from a delegate up to the root.
    ///
    /// Returns delegations in order from the immediate delegation (closest to
    /// the delegate) up to the root delegation.
    pub fn get_delegation_chain(&self, delegate_did: &str) -> Vec<&Delegation> {
        let mut chain = Vec::new();
        let mut current = delegate_did;

        loop {
            match self.delegations.get(current) {
                Some(delegations) if !delegations.is_empty() => {
                    // Take the first (most recent) delegation for this delegate
                    let delegation = &delegations[0];
                    chain.push(delegation);

                    // Walk up to the delegator
                    let delegator = delegation.delegator.as_str();
                    if delegator == current {
                        // Self-delegation, stop to prevent infinite loop
                        break;
                    }
                    current = delegator;
                }
                _ => break,
            }
        }

        chain
    }

    /// Computes the effective scope for a delegate.
    ///
    /// This is the intersection of all scopes in the delegation chain,
    /// implementing the intersection narrowing rule (§9.6 rule 2).
    pub fn effective_scope(&self, delegate_did: &str) -> Option<DelegationScope> {
        let chain = self.get_delegation_chain(delegate_did);
        if chain.is_empty() {
            return None;
        }

        let mut effective = chain[0].scope.clone();
        for delegation in chain.iter().skip(1) {
            effective = intersect_scopes(&effective, &delegation.scope);
        }

        Some(effective)
    }

    /// Returns the depth of a delegate in the tree.
    ///
    /// A root entity (not a delegate of anyone) has depth 0.
    /// A direct delegate of a root has depth 1, and so on.
    pub fn depth(&self, delegate_did: &str) -> u32 {
        self.get_delegation_chain(delegate_did).len() as u32
    }

    /// Revokes a delegation and all sub-delegations (cascade revocation).
    ///
    /// Returns the list of all revoked delegation IDs, including the original
    /// and all transitively dependent delegations.
    pub fn revoke(&mut self, delegation_id: &str) -> Result<Vec<String>, DelegationError> {
        // Find the delegation to revoke
        let (_, delegate_did) =
            self.id_index
                .get(delegation_id)
                .cloned()
                .ok_or_else(|| DelegationError::Revoked {
                    delegation_id: delegation_id.to_string(),
                })?;

        let mut revoked_ids = Vec::new();

        // Collect all DIDs that need revocation (BFS from delegate)
        let mut dids_to_process = vec![delegate_did.clone()];
        let mut processed_dids = std::collections::HashSet::new();

        while let Some(did) = dids_to_process.pop() {
            if !processed_dids.insert(did.clone()) {
                continue;
            }

            // Find all delegations FROM this DID (where this DID is a delegator)
            // by scanning all entries
            for (sub_delegate_did, delegations) in &self.delegations {
                for d in delegations {
                    if d.delegator == did && !processed_dids.contains(sub_delegate_did.as_str()) {
                        dids_to_process.push(sub_delegate_did.clone());
                    }
                }
            }
        }

        // Remove the original delegation by ID
        if let Some(delegations) = self.delegations.get_mut(&delegate_did) {
            delegations.retain(|d| {
                if d.id == delegation_id {
                    revoked_ids.push(d.id.clone());
                    false
                } else {
                    true
                }
            });
            if delegations.is_empty() {
                self.delegations.remove(&delegate_did);
            }
        }
        self.id_index.remove(delegation_id);

        // Remove all sub-delegations for processed DIDs (except the original DID
        // which may have other delegations)
        for did in &processed_dids {
            if *did == delegate_did {
                continue; // Already handled above
            }
            if let Some(delegations) = self.delegations.get_mut(did) {
                for d in delegations.iter() {
                    revoked_ids.push(d.id.clone());
                    self.id_index.remove(&d.id);
                }
            }
            self.delegations.remove(did);
        }

        Ok(revoked_ids)
    }
}

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

    fn make_delegation(
        id: &str,
        delegator: &str,
        delegate: &str,
        actions: Vec<&str>,
    ) -> Delegation {
        Delegation {
            id: id.to_string(),
            delegator: delegator.to_string(),
            delegate: delegate.to_string(),
            scope: DelegationScope {
                actions: actions.into_iter().map(String::from).collect(),
                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(),
            },
        }
    }

    #[test]
    fn single_delegation_depth() {
        let mut tree = DelegationTree::new(8);
        let d = make_delegation("d1", "did:root", "did:agent1", vec!["transfer"]);
        let result = tree.add_delegation(d);
        assert!(result.is_ok());
        assert_eq!(tree.depth("did:agent1"), 1);
        assert_eq!(tree.depth("did:root"), 0);
    }

    #[test]
    fn chain_depth_and_walk() {
        let mut tree = DelegationTree::new(8);

        let d1 = make_delegation("d1", "did:root", "did:a", vec!["transfer", "approve"]);
        let d2 = make_delegation("d2", "did:a", "did:b", vec!["transfer"]);
        let d3 = make_delegation("d3", "did:b", "did:c", vec!["transfer"]);

        assert!(tree.add_delegation(d1).is_ok());
        assert!(tree.add_delegation(d2).is_ok());
        assert!(tree.add_delegation(d3).is_ok());

        assert_eq!(tree.depth("did:c"), 3);

        let chain = tree.get_delegation_chain("did:c");
        assert_eq!(chain.len(), 3);
        assert_eq!(chain[0].delegate, "did:c");
        assert_eq!(chain[1].delegate, "did:b");
        assert_eq!(chain[2].delegate, "did:a");
    }

    #[test]
    fn max_depth_enforced() {
        let mut tree = DelegationTree::new(2);

        let d1 = make_delegation("d1", "did:root", "did:a", vec!["transfer"]);
        let d2 = make_delegation("d2", "did:a", "did:b", vec!["transfer"]);
        let d3 = make_delegation("d3", "did:b", "did:c", vec!["transfer"]);

        assert!(tree.add_delegation(d1).is_ok());
        assert!(tree.add_delegation(d2).is_ok());

        // This should fail: depth would be 3, max is 2
        let result = tree.add_delegation(d3);
        assert!(result.is_err());
        match result {
            Err(DelegationError::MaxDepthExceeded { depth, max_depth }) => {
                assert_eq!(depth, 3);
                assert_eq!(max_depth, 2);
            }
            _ => panic!("expected MaxDepthExceeded"),
        }
    }

    #[test]
    fn scope_amplification_rejected() {
        let mut tree = DelegationTree::new(8);

        let d1 = make_delegation("d1", "did:root", "did:a", vec!["transfer"]);

        // did:a tries to delegate "approve" which it doesn't have
        let d2 = make_delegation("d2", "did:a", "did:b", vec!["transfer", "approve"]);

        assert!(tree.add_delegation(d1).is_ok());

        let result = tree.add_delegation(d2);
        assert!(result.is_err());
        match result {
            Err(DelegationError::ScopeAmplification { .. }) => {}
            other => panic!("expected ScopeAmplification, got {other:?}"),
        }
    }

    #[test]
    fn effective_scope_is_intersection() {
        let mut tree = DelegationTree::new(8);

        let d1 = make_delegation(
            "d1",
            "did:root",
            "did:a",
            vec!["transfer", "approve", "stake"],
        );
        let d2 = make_delegation("d2", "did:a", "did:b", vec!["transfer", "approve"]);

        assert!(tree.add_delegation(d1).is_ok());
        assert!(tree.add_delegation(d2).is_ok());

        let scope = tree.effective_scope("did:b");
        assert!(scope.is_some());
        if let Some(s) = scope {
            // Intersection of [transfer,approve] and [transfer,approve,stake]
            assert!(s.actions.contains(&"transfer".to_string()));
            assert!(s.actions.contains(&"approve".to_string()));
            assert!(!s.actions.contains(&"stake".to_string()));
        }
    }

    #[test]
    fn revocation_cascades() {
        let mut tree = DelegationTree::new(8);

        let d1 = make_delegation("d1", "did:root", "did:a", vec!["transfer"]);
        let d2 = make_delegation("d2", "did:a", "did:b", vec!["transfer"]);
        let d3 = make_delegation("d3", "did:b", "did:c", vec!["transfer"]);

        assert!(tree.add_delegation(d1).is_ok());
        assert!(tree.add_delegation(d2).is_ok());
        assert!(tree.add_delegation(d3).is_ok());

        // Revoking d1 should cascade to d2 and d3
        let revoked = tree.revoke("d1");
        assert!(revoked.is_ok());
        if let Ok(ids) = revoked {
            assert!(ids.contains(&"d1".to_string()));
            assert!(ids.contains(&"d2".to_string()));
            assert!(ids.contains(&"d3".to_string()));
        }

        // Tree should be empty now
        assert_eq!(tree.depth("did:a"), 0);
        assert_eq!(tree.depth("did:b"), 0);
        assert_eq!(tree.depth("did:c"), 0);
    }

    #[test]
    fn revoke_nonexistent_returns_error() {
        let mut tree = DelegationTree::new(8);
        let result = tree.revoke("nonexistent");
        assert!(result.is_err());
    }

    #[test]
    fn no_chain_for_unknown_did() {
        let tree = DelegationTree::new(8);
        let chain = tree.get_delegation_chain("did:unknown");
        assert!(chain.is_empty());
        assert!(tree.effective_scope("did:unknown").is_none());
    }
}