photoncloud-monorepo/iam/crates/iam-store/src/group_store.rs

//! Group membership storage
//!
//! Tracks which principals are members of which groups.

use std::sync::Arc;

use iam_types::{Error, PrincipalKind, PrincipalRef, Result};

use crate::backend::{Backend, CasResult, JsonStore, KvPair, StorageBackend};

/// Key prefixes for group membership storage
mod keys {
    /// Members of a group
    /// Format: iam/groups/{group_id}/members/{kind}/{member_id}
    pub const GROUP_MEMBERS: &str = "iam/groups/";

    /// Groups a principal belongs to (reverse index)
    /// Format: iam/memberships/{kind}/{principal_id}/{group_id}
    pub const PRINCIPAL_GROUPS: &str = "iam/memberships/";
}

/// Store for group memberships
pub struct GroupStore {
    backend: Arc<Backend>,
}

impl JsonStore for GroupStore {
    fn backend(&self) -> &Backend {
        &self.backend
    }
}

impl GroupStore {
    /// Create a new group store
    pub fn new(backend: Arc<Backend>) -> Self {
        Self { backend }
    }

    /// Add a member to a group
    pub async fn add_member(&self, group_id: &str, member: &PrincipalRef) -> Result<()> {
        // Create forward index (group -> member)
        let forward_key = format!(
            "{}{}/members/{}/{}",
            keys::GROUP_MEMBERS,
            group_id,
            member.kind,
            member.id
        );

        // Create reverse index (member -> group)
        let reverse_key = format!(
            "{}{}/{}/{}",
            keys::PRINCIPAL_GROUPS,
            member.kind,
            member.id,
            group_id
        );

        // Store membership marker (empty value, just presence matters)
        let marker = b"1";

        // Create forward index
        match self.backend.cas(forward_key.as_bytes(), 0, marker).await? {
            CasResult::Success(_) => {}
            CasResult::Conflict { .. } => {
                // Already a member, that's fine
                return Ok(());
            }
            CasResult::NotFound => return Err(Error::Internal("Unexpected CAS result".into())),
        }

        // Create reverse index
        self.backend
            .cas(reverse_key.as_bytes(), 0, marker)
            .await
            .ok();

        Ok(())
    }

    /// Remove a member from a group
    pub async fn remove_member(&self, group_id: &str, member: &PrincipalRef) -> Result<bool> {
        let forward_key = format!(
            "{}{}/members/{}/{}",
            keys::GROUP_MEMBERS,
            group_id,
            member.kind,
            member.id
        );
        let reverse_key = format!(
            "{}{}/{}/{}",
            keys::PRINCIPAL_GROUPS,
            member.kind,
            member.id,
            group_id
        );

        // Delete forward index
        let deleted = self.backend.delete(forward_key.as_bytes()).await?;

        // Delete reverse index
        self.backend.delete(reverse_key.as_bytes()).await.ok();

        Ok(deleted)
    }

    /// List all members of a group
    pub async fn list_members(&self, group_id: &str) -> Result<Vec<PrincipalRef>> {
        let prefix = format!("{}{}/members/", keys::GROUP_MEMBERS, group_id);
        let pairs = self.backend.scan_prefix(prefix.as_bytes(), 1000).await?;

        let mut members = Vec::new();
        for KvPair { key, .. } in pairs {
            // Key format: iam/groups/{group_id}/members/{kind}/{member_id}
            let key_str = String::from_utf8_lossy(&key);
            let parts: Vec<&str> = key_str.split('/').collect();
            if parts.len() >= 6 {
                let kind_str = parts[4];
                let member_id = parts[5];

                let kind = match kind_str {
                    "user" => PrincipalKind::User,
                    "service_account" => PrincipalKind::ServiceAccount,
                    "group" => PrincipalKind::Group,
                    _ => continue,
                };

                members.push(PrincipalRef::new(kind, member_id));
            }
        }

        Ok(members)
    }

    /// List all groups a principal belongs to
    pub async fn list_groups(&self, principal: &PrincipalRef) -> Result<Vec<String>> {
        let prefix = format!(
            "{}{}/{}/",
            keys::PRINCIPAL_GROUPS,
            principal.kind,
            principal.id
        );
        let pairs = self.backend.scan_prefix(prefix.as_bytes(), 1000).await?;

        let mut groups = Vec::new();
        for KvPair { key, .. } in pairs {
            // Key format: iam/memberships/{kind}/{principal_id}/{group_id}
            let key_str = String::from_utf8_lossy(&key);
            let parts: Vec<&str> = key_str.split('/').collect();
            if parts.len() >= 5 {
                groups.push(parts[4].to_string());
            }
        }

        Ok(groups)
    }

    /// Check if a principal is a member of a group
    pub async fn is_member(&self, group_id: &str, member: &PrincipalRef) -> Result<bool> {
        let key = format!(
            "{}{}/members/{}/{}",
            keys::GROUP_MEMBERS,
            group_id,
            member.kind,
            member.id
        );
        let result = self.backend.get(key.as_bytes()).await?;
        Ok(result.is_some())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Backend;

    #[tokio::test]
    async fn test_add_and_list_members() {
        let backend = Arc::new(Backend::memory());
        let store = GroupStore::new(backend);

        let alice = PrincipalRef::user("alice");
        let bob = PrincipalRef::user("bob");

        store.add_member("devs", &alice).await.unwrap();
        store.add_member("devs", &bob).await.unwrap();

        let members = store.list_members("devs").await.unwrap();
        assert_eq!(members.len(), 2);
    }

    #[tokio::test]
    async fn test_remove_member() {
        let backend = Arc::new(Backend::memory());
        let store = GroupStore::new(backend);

        let alice = PrincipalRef::user("alice");

        store.add_member("devs", &alice).await.unwrap();
        assert!(store.is_member("devs", &alice).await.unwrap());

        store.remove_member("devs", &alice).await.unwrap();
        assert!(!store.is_member("devs", &alice).await.unwrap());
    }

    #[tokio::test]
    async fn test_list_groups_for_principal() {
        let backend = Arc::new(Backend::memory());
        let store = GroupStore::new(backend);

        let alice = PrincipalRef::user("alice");

        store.add_member("devs", &alice).await.unwrap();
        store.add_member("admins", &alice).await.unwrap();

        let groups = store.list_groups(&alice).await.unwrap();
        assert_eq!(groups.len(), 2);
        assert!(groups.contains(&"devs".to_string()));
        assert!(groups.contains(&"admins".to_string()));
    }
}