photoncloud-monorepo/iam/crates/iam-api/src/token_service.rs

//! IAM Token gRPC service implementation

use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use sha2::{Digest, Sha256};
use tonic::{Request, Response, Status};

use iam_authn::InternalTokenService;
use iam_store::{PrincipalStore, TokenStore};
use iam_types::{
    InternalTokenClaims, PrincipalKind as TypesPrincipalKind, PrincipalRef, Scope as TypesScope,
    TokenMetadata, TokenType,
};

use crate::proto::{
    self, iam_token_server::IamToken, scope, IssueTokenRequest, IssueTokenResponse, PrincipalKind,
    RefreshTokenRequest, RefreshTokenResponse, RevokeTokenRequest, RevokeTokenResponse, Scope,
    ValidateTokenRequest, ValidateTokenResponse,
};

/// IAM Token service implementation
pub struct IamTokenService {
    token_service: Arc<InternalTokenService>,
    principal_store: Arc<PrincipalStore>,
    token_store: Arc<TokenStore>,
}

impl IamTokenService {
    /// Create a new token service
    pub fn new(
        token_service: Arc<InternalTokenService>,
        principal_store: Arc<PrincipalStore>,
        token_store: Arc<TokenStore>,
    ) -> Self {
        Self {
            token_service,
            principal_store,
            token_store,
        }
    }

    fn convert_scope(proto_scope: &Option<Scope>) -> TypesScope {
        match proto_scope {
            Some(s) => match &s.scope {
                Some(scope::Scope::System(true)) => TypesScope::System,
                Some(scope::Scope::Org(org)) => TypesScope::org(&org.id),
                Some(scope::Scope::Project(proj)) => TypesScope::project(&proj.id, &proj.org_id),
                Some(scope::Scope::Resource(res)) => {
                    TypesScope::resource(&res.id, &res.project_id, &res.org_id)
                }
                _ => TypesScope::System,
            },
            None => TypesScope::System,
        }
    }

    fn convert_scope_to_proto(scope: &TypesScope) -> Scope {
        Scope {
            scope: Some(match scope {
                TypesScope::System => scope::Scope::System(true),
                TypesScope::Org { id } => scope::Scope::Org(proto::OrgScope { id: id.clone() }),
                TypesScope::Project { id, org_id } => scope::Scope::Project(proto::ProjectScope {
                    id: id.clone(),
                    org_id: org_id.clone(),
                }),
                TypesScope::Resource {
                    id,
                    project_id,
                    org_id,
                } => scope::Scope::Resource(proto::ResourceScope {
                    id: id.clone(),
                    project_id: project_id.clone(),
                    org_id: org_id.clone(),
                }),
            }),
        }
    }

    fn compute_token_id(token: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(token.as_bytes());
        let digest = hasher.finalize();
        URL_SAFE_NO_PAD.encode(digest)
    }

    fn now_ts() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs()
    }

    async fn persist_metadata(
        &self,
        claims: &InternalTokenClaims,
        token: &str,
    ) -> Result<String, Status> {
        let token_id = Self::compute_token_id(token);
        let meta = TokenMetadata::new(
            &token_id,
            &claims.principal_id,
            TokenType::Access,
            claims.iat,
            claims.exp,
        );
        if let Err(e) = self.token_store.put(&meta).await {
            return Err(Status::internal(format!(
                "failed to persist token metadata: {}",
                e
            )));
        }
        // ensure we mirror revocation flags if claims already mark them
        if meta.revoked {
            return Ok(token_id);
        }
        Ok(token_id)
    }
}

#[tonic::async_trait]
impl IamToken for IamTokenService {
    async fn issue_token(
        &self,
        request: Request<IssueTokenRequest>,
    ) -> Result<Response<IssueTokenResponse>, Status> {
        let req = request.into_inner();

        // Get principal kind
        let principal_kind = match PrincipalKind::try_from(req.principal_kind) {
            Ok(PrincipalKind::User) => TypesPrincipalKind::User,
            Ok(PrincipalKind::ServiceAccount) => TypesPrincipalKind::ServiceAccount,
            Ok(PrincipalKind::Group) => TypesPrincipalKind::Group,
            _ => return Err(Status::invalid_argument("invalid principal kind")),
        };

        // Get principal from store
        let principal_ref = PrincipalRef::new(principal_kind, &req.principal_id);
        let principal = self
            .principal_store
            .get(&principal_ref)
            .await
            .map_err(|e| Status::internal(format!("Failed to get principal: {}", e)))?
            .ok_or_else(|| Status::not_found("Principal not found"))?;
        if !principal.enabled {
            return Err(Status::permission_denied("Principal is disabled"));
        }

        // Convert scope
        let scope = Self::convert_scope(&req.scope);

        // Determine TTL
        let ttl = if req.ttl_seconds > 0 {
            Some(Duration::from_secs(req.ttl_seconds))
        } else {
            None
        };

        // Issue token
        let issued = self
            .token_service
            .issue(&principal, req.roles.clone(), scope, ttl)
            .await
            .map_err(|e| Status::internal(format!("Failed to issue token: {}", e)))?;

        // Persist metadata for revocation tracking
        self.persist_metadata(&issued.claims, &issued.token).await?;

        Ok(Response::new(IssueTokenResponse {
            token: issued.token,
            expires_at: issued.expires_at,
            session_id: issued.claims.session_id,
        }))
    }

    async fn validate_token(
        &self,
        request: Request<ValidateTokenRequest>,
    ) -> Result<Response<ValidateTokenResponse>, Status> {
        let req = request.into_inner();

        match self.token_service.verify(&req.token).await {
            Ok(claims) => {
                // Check revocation list
                let token_id = Self::compute_token_id(&req.token);
                if let Some((meta, _)) = self
                    .token_store
                    .get(&claims.principal_id, &token_id)
                    .await
                    .map_err(|e| {
                        Status::internal(format!("Failed to read token metadata: {}", e))
                    })?
                {
                    if meta.revoked {
                        return Ok(Response::new(ValidateTokenResponse {
                            valid: false,
                            claims: None,
                            reason: "token revoked".into(),
                        }));
                    }
                }

                // Ensure principal still exists and is enabled
                let principal_ref =
                    PrincipalRef::new(claims.principal_kind.clone(), &claims.principal_id);
                let principal = self
                    .principal_store
                    .get(&principal_ref)
                    .await
                    .map_err(|e| Status::internal(format!("Failed to read principal: {}", e)))?;
                let Some(principal) = principal else {
                    return Ok(Response::new(ValidateTokenResponse {
                        valid: false,
                        claims: None,
                        reason: "principal not found".into(),
                    }));
                };
                if !principal.enabled {
                    return Ok(Response::new(ValidateTokenResponse {
                        valid: false,
                        claims: None,
                        reason: "principal disabled".into(),
                    }));
                }

                let proto_claims = crate::proto::InternalTokenClaims {
                    principal_id: claims.principal_id.clone(),
                    principal_kind: match claims.principal_kind {
                        TypesPrincipalKind::User => PrincipalKind::User as i32,
                        TypesPrincipalKind::ServiceAccount => PrincipalKind::ServiceAccount as i32,
                        TypesPrincipalKind::Group => PrincipalKind::Group as i32,
                    },
                    principal_name: claims.principal_name.clone(),
                    roles: claims.roles.clone(),
                    scope: Some(Self::convert_scope_to_proto(&claims.scope)),
                    org_id: claims.org_id.clone(),
                    project_id: claims.project_id.clone(),
                    node_id: claims.node_id.clone(),
                    iat: claims.iat,
                    exp: claims.exp,
                    session_id: claims.session_id.clone(),
                    auth_method: claims.auth_method.to_string(),
                };

                Ok(Response::new(ValidateTokenResponse {
                    valid: true,
                    claims: Some(proto_claims),
                    reason: String::new(),
                }))
            }
            Err(e) => Ok(Response::new(ValidateTokenResponse {
                valid: false,
                claims: None,
                reason: e.to_string(),
            })),
        }
    }

    async fn revoke_token(
        &self,
        request: Request<RevokeTokenRequest>,
    ) -> Result<Response<RevokeTokenResponse>, Status> {
        let req = request.into_inner();

        let claims = self
            .token_service
            .verify(&req.token)
            .await
            .map_err(|e| Status::invalid_argument(format!("Invalid token: {}", e)))?;

        let token_id = Self::compute_token_id(&req.token);
        let now = Self::now_ts();

        let existing = self
            .token_store
            .get(&claims.principal_id, &token_id)
            .await
            .map_err(|e| Status::internal(format!("Failed to read token metadata: {}", e)))?;

        // Try revoking existing metadata; if it does not exist, create a revoked record
        let revoked = if let Some((meta, _)) = existing {
            if meta.revoked {
                true
            } else {
                self.token_store
                    .revoke(&claims.principal_id, &token_id, &req.reason, now)
                    .await
                    .map_err(|e| Status::internal(format!("Failed to revoke token: {}", e)))?
            }
        } else {
            let mut meta = TokenMetadata::new(
                &token_id,
                &claims.principal_id,
                TokenType::Access,
                claims.iat,
                claims.exp,
            );
            meta.revoke(now, &req.reason);
            self.token_store
                .put(&meta)
                .await
                .map_err(|e| Status::internal(format!("Failed to write revocation: {}", e)))?;
            true
        };

        Ok(Response::new(RevokeTokenResponse { success: revoked }))
    }

    async fn refresh_token(
        &self,
        request: Request<RefreshTokenRequest>,
    ) -> Result<Response<RefreshTokenResponse>, Status> {
        let req = request.into_inner();

        // Validate current token
        let claims = self
            .token_service
            .verify(&req.token)
            .await
            .map_err(|e| Status::unauthenticated(format!("Invalid token: {}", e)))?;

        // Check revocation list
        let token_id = Self::compute_token_id(&req.token);
        if let Some((meta, _)) = self
            .token_store
            .get(&claims.principal_id, &token_id)
            .await
            .map_err(|e| Status::internal(format!("Failed to read token metadata: {}", e)))?
        {
            if meta.revoked {
                return Err(Status::permission_denied("token revoked"));
            }
        }

        // Get principal
        let principal_kind = claims.principal_kind.clone();
        let principal_ref = PrincipalRef::new(principal_kind, &claims.principal_id);
        let principal = self
            .principal_store
            .get(&principal_ref)
            .await
            .map_err(|e| Status::internal(format!("Failed to get principal: {}", e)))?
            .ok_or_else(|| Status::not_found("Principal not found"))?;
        if !principal.enabled {
            return Err(Status::permission_denied("Principal is disabled"));
        }

        // Determine new TTL
        let ttl = if req.ttl_seconds > 0 {
            Some(Duration::from_secs(req.ttl_seconds))
        } else {
            None
        };

        // Issue new token
        let issued = self
            .token_service
            .issue(&principal, claims.roles.clone(), claims.scope.clone(), ttl)
            .await
            .map_err(|e| Status::internal(format!("Failed to issue token: {}", e)))?;

        self.persist_metadata(&issued.claims, &issued.token).await?;

        Ok(Response::new(RefreshTokenResponse {
            token: issued.token,
            expires_at: issued.expires_at,
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use iam_authn::{InternalTokenConfig, SigningKey};
    use iam_store::{Backend, TokenStore};
    use iam_types::Principal;

    fn test_setup() -> (
        Arc<InternalTokenService>,
        Arc<PrincipalStore>,
        Arc<TokenStore>,
    ) {
        let backend = Arc::new(Backend::memory());
        let principal_store = Arc::new(PrincipalStore::new(backend.clone()));
        let token_store = Arc::new(TokenStore::new(backend));

        let signing_key = SigningKey::generate("test-key");
        let config = InternalTokenConfig::new(signing_key, "iam-test");
        let token_service = Arc::new(InternalTokenService::new(config));

        (token_service, principal_store, token_store)
    }

    #[tokio::test]
    async fn test_issue_token_principal_not_found() {
        let (token_service, principal_store, token_store) = test_setup();
        let service = IamTokenService::new(token_service, principal_store, token_store);

        let req = IssueTokenRequest {
            principal_id: "nonexistent".into(),
            principal_kind: PrincipalKind::User as i32,
            roles: vec![],
            scope: None,
            ttl_seconds: 3600,
        };

        let result = service.issue_token(Request::new(req)).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_revoke_and_validate_blocklist() {
        let (token_service, principal_store, token_store) = test_setup();
        let service =
            IamTokenService::new(token_service, principal_store.clone(), token_store.clone());

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

        // issue token
        let issue_resp = service
            .issue_token(Request::new(IssueTokenRequest {
                principal_id: "alice".into(),
                principal_kind: PrincipalKind::User as i32,
                roles: vec!["roles/ProjectAdmin".into()],
                scope: None,
                ttl_seconds: 3600,
            }))
            .await
            .unwrap()
            .into_inner();

        // validate OK
        let valid_resp = service
            .validate_token(Request::new(ValidateTokenRequest {
                token: issue_resp.token.clone(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(valid_resp.valid);

        // revoke
        let revoke_resp = service
            .revoke_token(Request::new(RevokeTokenRequest {
                token: issue_resp.token.clone(),
                reason: "test".into(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(revoke_resp.success);

        // validate should be rejected
        let invalid_resp = service
            .validate_token(Request::new(ValidateTokenRequest {
                token: issue_resp.token.clone(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(!invalid_resp.valid);
        assert_eq!(invalid_resp.reason, "token revoked");
    }

    #[tokio::test]
    async fn test_validate_token_principal_disabled() {
        let (token_service, principal_store, token_store) = test_setup();
        let service =
            IamTokenService::new(token_service, principal_store.clone(), token_store.clone());

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

        let issue_resp = service
            .issue_token(Request::new(IssueTokenRequest {
                principal_id: "alice".into(),
                principal_kind: PrincipalKind::User as i32,
                roles: vec!["roles/ProjectAdmin".into()],
                scope: None,
                ttl_seconds: 3600,
            }))
            .await
            .unwrap()
            .into_inner();

        let (mut stored, version) = principal_store
            .get_with_version(&principal.to_ref())
            .await
            .unwrap()
            .unwrap();
        stored.enabled = false;
        principal_store.update(&stored, version).await.unwrap();

        let valid_resp = service
            .validate_token(Request::new(ValidateTokenRequest {
                token: issue_resp.token,
            }))
            .await
            .unwrap()
            .into_inner();

        assert!(!valid_resp.valid);
        assert!(valid_resp.reason.contains("disabled"));
    }
}