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

//! Principal storage
//!
//! Stores and retrieves principals (users, service accounts, groups).

use std::sync::Arc;

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

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

/// Key prefixes for principal storage
mod keys {
    /// Primary key: by kind and ID
    /// Format: iam/principals/{kind}/{id}
    pub const PRINCIPALS: &str = "iam/principals/";

    /// Secondary index: by org
    /// Format: iam/principals/by-org/{org_id}/{kind}/{id}
    pub const BY_ORG: &str = "iam/principals/by-org/";

    /// Secondary index: by project (for service accounts)
    /// Format: iam/principals/by-project/{project_id}/{id}
    pub const BY_PROJECT: &str = "iam/principals/by-project/";

    /// Secondary index: by email
    /// Format: iam/principals/by-email/{email}
    pub const BY_EMAIL: &str = "iam/principals/by-email/";

    /// Secondary index: by OIDC subject
    /// Format: iam/principals/by-oidc/{iss_hash}/{sub}
    pub const BY_OIDC: &str = "iam/principals/by-oidc/";
}

/// Store for principals
pub struct PrincipalStore {
    backend: Arc<Backend>,
}

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

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

    /// Create a new principal
    pub async fn create(&self, principal: &Principal) -> Result<u64> {
        let key = self.make_primary_key(&principal.kind, &principal.id);

        let bytes =
            serde_json::to_vec(principal).map_err(|e| Error::Serialization(e.to_string()))?;

        match self.backend.cas(key.as_bytes(), 0, &bytes).await? {
            CasResult::Success(version) => {
                // Create secondary indexes
                self.create_indexes(principal).await?;
                Ok(version)
            }
            CasResult::Conflict { .. } => Err(Error::Iam(IamError::PrincipalAlreadyExists(
                principal.to_ref().to_string(),
            ))),
            CasResult::NotFound => Err(Error::Internal("Unexpected CAS result".into())),
        }
    }

    /// Get a principal by reference
    pub async fn get(&self, principal_ref: &PrincipalRef) -> Result<Option<Principal>> {
        let key = self.make_primary_key(&principal_ref.kind, &principal_ref.id);
        match self.get_json::<Principal>(key.as_bytes()).await? {
            Some((principal, _)) => Ok(Some(principal)),
            None => Ok(None),
        }
    }

    /// Get a principal with version
    pub async fn get_with_version(
        &self,
        principal_ref: &PrincipalRef,
    ) -> Result<Option<(Principal, u64)>> {
        let key = self.make_primary_key(&principal_ref.kind, &principal_ref.id);
        self.get_json::<Principal>(key.as_bytes()).await
    }

    /// Get a principal by email
    pub async fn get_by_email(&self, email: &str) -> Result<Option<Principal>> {
        let index_key = format!("{}{}", keys::BY_EMAIL, email.to_lowercase());

        if let Some((ref_bytes, _)) = self.backend.get(index_key.as_bytes()).await? {
            let principal_ref: PrincipalRef = serde_json::from_slice(&ref_bytes)
                .map_err(|e| Error::Serialization(e.to_string()))?;
            return self.get(&principal_ref).await;
        }

        Ok(None)
    }

    /// Get a principal by OIDC subject
    pub async fn get_by_oidc(&self, issuer: &str, subject: &str) -> Result<Option<Principal>> {
        let iss_hash = self.hash_issuer(issuer);
        let index_key = format!("{}{}/{}", keys::BY_OIDC, iss_hash, subject);

        if let Some((ref_bytes, _)) = self.backend.get(index_key.as_bytes()).await? {
            let principal_ref: PrincipalRef = serde_json::from_slice(&ref_bytes)
                .map_err(|e| Error::Serialization(e.to_string()))?;
            return self.get(&principal_ref).await;
        }

        Ok(None)
    }

    /// Update a principal
    pub async fn update(&self, principal: &Principal, expected_version: u64) -> Result<u64> {
        let key = self.make_primary_key(&principal.kind, &principal.id);
        let bytes =
            serde_json::to_vec(principal).map_err(|e| Error::Serialization(e.to_string()))?;

        match self
            .backend
            .cas(key.as_bytes(), expected_version, &bytes)
            .await?
        {
            CasResult::Success(version) => {
                // Update indexes if needed (email, oidc changes)
                // For simplicity, recreate all indexes
                self.create_indexes(principal).await?;
                Ok(version)
            }
            CasResult::Conflict { expected, actual } => {
                Err(Error::Storage(iam_types::StorageError::CasConflict {
                    expected,
                    actual,
                }))
            }
            CasResult::NotFound => Err(Error::Iam(IamError::PrincipalNotFound(
                principal.to_ref().to_string(),
            ))),
        }
    }

    /// Delete a principal
    pub async fn delete(&self, principal_ref: &PrincipalRef) -> Result<bool> {
        // First get the principal to know what indexes to delete
        if let Some(principal) = self.get(principal_ref).await? {
            let key = self.make_primary_key(&principal.kind, &principal.id);

            let deleted = self.backend.delete(key.as_bytes()).await?;

            if deleted {
                self.delete_indexes(&principal).await?;
            }

            Ok(deleted)
        } else {
            Ok(false)
        }
    }

    /// List principals by kind
    pub async fn list_by_kind(&self, kind: &PrincipalKind) -> Result<Vec<Principal>> {
        let prefix = format!("{}{}/", keys::PRINCIPALS, kind);
        let pairs = self.backend.scan_prefix(prefix.as_bytes(), 10000).await?;

        let mut principals = Vec::new();
        for pair in pairs {
            let principal: Principal = serde_json::from_slice(&pair.value)
                .map_err(|e| Error::Serialization(e.to_string()))?;
            principals.push(principal);
        }

        Ok(principals)
    }

    /// List principals by organization
    pub async fn list_by_org(&self, org_id: &str) -> Result<Vec<Principal>> {
        let prefix = format!("{}{}/", keys::BY_ORG, org_id);
        let pairs = self.backend.scan_prefix(prefix.as_bytes(), 10000).await?;

        let mut principals = Vec::new();
        for pair in pairs {
            let principal_ref: PrincipalRef = serde_json::from_slice(&pair.value)
                .map_err(|e| Error::Serialization(e.to_string()))?;
            if let Some(principal) = self.get(&principal_ref).await? {
                principals.push(principal);
            }
        }

        Ok(principals)
    }

    /// List service accounts by project
    pub async fn list_by_project(&self, project_id: &str) -> Result<Vec<Principal>> {
        let prefix = format!("{}{}/", keys::BY_PROJECT, project_id);
        let pairs = self.backend.scan_prefix(prefix.as_bytes(), 10000).await?;

        let mut principals = Vec::new();
        for pair in pairs {
            let principal_ref: PrincipalRef = serde_json::from_slice(&pair.value)
                .map_err(|e| Error::Serialization(e.to_string()))?;
            if let Some(principal) = self.get(&principal_ref).await? {
                principals.push(principal);
            }
        }

        Ok(principals)
    }

    /// Check if a principal exists
    pub async fn exists(&self, principal_ref: &PrincipalRef) -> Result<bool> {
        let key = self.make_primary_key(&principal_ref.kind, &principal_ref.id);
        Ok(self.backend.get(key.as_bytes()).await?.is_some())
    }

    // Helper methods

    fn make_primary_key(&self, kind: &PrincipalKind, id: &str) -> String {
        format!("{}{}/{}", keys::PRINCIPALS, kind, id)
    }

    fn hash_issuer(&self, issuer: &str) -> String {
        // Simple hash for issuer to avoid special characters in keys
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        let mut hasher = DefaultHasher::new();
        issuer.hash(&mut hasher);
        format!("{:016x}", hasher.finish())
    }

    async fn create_indexes(&self, principal: &Principal) -> Result<()> {
        let ref_bytes = serde_json::to_vec(&principal.to_ref())
            .map_err(|e| Error::Serialization(e.to_string()))?;

        // Org index
        if let Some(org_id) = &principal.org_id {
            let key = format!(
                "{}{}/{}/{}",
                keys::BY_ORG,
                org_id,
                principal.kind,
                principal.id
            );
            self.backend.put(key.as_bytes(), &ref_bytes).await?;
        }

        // Project index (for service accounts)
        if let Some(project_id) = &principal.project_id {
            let key = format!("{}{}/{}", keys::BY_PROJECT, project_id, principal.id);
            self.backend.put(key.as_bytes(), &ref_bytes).await?;
        }

        // Email index
        if let Some(email) = &principal.email {
            let key = format!("{}{}", keys::BY_EMAIL, email.to_lowercase());
            self.backend.put(key.as_bytes(), &ref_bytes).await?;
        }

        // OIDC index
        if let Some(oidc_sub) = &principal.oidc_sub {
            // Assume we store issuer in metadata
            if let Some(issuer) = principal.metadata.get("oidc_issuer") {
                let iss_hash = self.hash_issuer(issuer);
                let key = format!("{}{}/{}", keys::BY_OIDC, iss_hash, oidc_sub);
                self.backend.put(key.as_bytes(), &ref_bytes).await?;
            }
        }

        Ok(())
    }

    async fn delete_indexes(&self, principal: &Principal) -> Result<()> {
        // Org index
        if let Some(org_id) = &principal.org_id {
            let key = format!(
                "{}{}/{}/{}",
                keys::BY_ORG,
                org_id,
                principal.kind,
                principal.id
            );
            self.backend.delete(key.as_bytes()).await?;
        }

        // Project index
        if let Some(project_id) = &principal.project_id {
            let key = format!("{}{}/{}", keys::BY_PROJECT, project_id, principal.id);
            self.backend.delete(key.as_bytes()).await?;
        }

        // Email index
        if let Some(email) = &principal.email {
            let key = format!("{}{}", keys::BY_EMAIL, email.to_lowercase());
            self.backend.delete(key.as_bytes()).await?;
        }

        // OIDC index
        if let Some(oidc_sub) = &principal.oidc_sub {
            if let Some(issuer) = principal.metadata.get("oidc_issuer") {
                let iss_hash = self.hash_issuer(issuer);
                let key = format!("{}{}/{}", keys::BY_OIDC, iss_hash, oidc_sub);
                self.backend.delete(key.as_bytes()).await?;
            }
        }

        Ok(())
    }
}

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

    fn test_backend() -> Arc<Backend> {
        Arc::new(Backend::memory())
    }

    #[tokio::test]
    async fn test_principal_crud() {
        let store = PrincipalStore::new(test_backend());

        let mut principal = Principal::new_user("alice", "Alice Smith");
        principal.email = Some("alice@example.com".into());
        principal.org_id = Some("org-1".into());

        // Create
        let version = store.create(&principal).await.unwrap();
        assert!(version > 0);

        // Get
        let fetched = store.get(&PrincipalRef::user("alice")).await.unwrap();
        assert!(fetched.is_some());
        assert_eq!(fetched.unwrap().name, "Alice Smith");

        // Get by email
        let fetched = store.get_by_email("alice@example.com").await.unwrap();
        assert!(fetched.is_some());
        assert_eq!(fetched.unwrap().id, "alice");

        // Delete
        let deleted = store.delete(&PrincipalRef::user("alice")).await.unwrap();
        assert!(deleted);

        // Verify deleted
        let fetched = store.get(&PrincipalRef::user("alice")).await.unwrap();
        assert!(fetched.is_none());
    }

    #[tokio::test]
    async fn test_service_account() {
        let store = PrincipalStore::new(test_backend());

        let sa = Principal::new_service_account("compute-agent", "Compute Agent", "proj-1");

        store.create(&sa).await.unwrap();

        // List by project
        let sas = store.list_by_project("proj-1").await.unwrap();
        assert_eq!(sas.len(), 1);
        assert_eq!(sas[0].id, "compute-agent");
    }

    #[tokio::test]
    async fn test_list_by_kind() {
        let store = PrincipalStore::new(test_backend());

        store
            .create(&Principal::new_user("user1", "User 1"))
            .await
            .unwrap();
        store
            .create(&Principal::new_user("user2", "User 2"))
            .await
            .unwrap();
        store
            .create(&Principal::new_service_account("sa1", "SA 1", "proj-1"))
            .await
            .unwrap();

        let users = store.list_by_kind(&PrincipalKind::User).await.unwrap();
        assert_eq!(users.len(), 2);

        let sas = store
            .list_by_kind(&PrincipalKind::ServiceAccount)
            .await
            .unwrap();
        assert_eq!(sas.len(), 1);
    }

    #[tokio::test]
    async fn test_duplicate_create() {
        let store = PrincipalStore::new(test_backend());

        let principal = Principal::new_user("alice", "Alice");
        store.create(&principal).await.unwrap();

        // Try to create again
        let result = store.create(&principal).await;
        assert!(result.is_err());
    }
}