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

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

use apigateway_api::proto::{AuthorizeRequest, AuthorizeResponse, Subject};
use apigateway_api::GatewayAuthService;
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use iam_authn::InternalTokenService;
use iam_authz::{AuthzContext, AuthzDecision, AuthzRequest, PolicyEvaluator};
use iam_store::{PrincipalStore, TokenStore};
use iam_types::{InternalTokenClaims, Principal, PrincipalRef, Resource};
use sha2::{Digest, Sha256};
use tonic::{Request, Response, Status};

pub struct GatewayAuthServiceImpl {
    token_service: Arc<InternalTokenService>,
    principal_store: Arc<PrincipalStore>,
    token_store: Arc<TokenStore>,
    evaluator: Arc<PolicyEvaluator>,
}

impl GatewayAuthServiceImpl {
    pub fn new(
        token_service: Arc<InternalTokenService>,
        principal_store: Arc<PrincipalStore>,
        token_store: Arc<TokenStore>,
        evaluator: Arc<PolicyEvaluator>,
    ) -> Self {
        Self {
            token_service,
            principal_store,
            token_store,
            evaluator,
        }
    }

    async fn check_token_revoked(
        &self,
        principal_id: &str,
        token: &str,
    ) -> Result<Option<String>, Status> {
        let token_id = compute_token_id(token);
        let meta = self
            .token_store
            .get(principal_id, &token_id)
            .await
            .map_err(|e| Status::internal(format!("token store error: {}", e)))?;

        if let Some((meta, _)) = meta {
            if meta.revoked {
                let reason = meta
                    .revocation_reason
                    .unwrap_or_else(|| "token revoked".to_string());
                return Ok(Some(reason));
            }
        }

        Ok(None)
    }
}

#[tonic::async_trait]
impl GatewayAuthService for GatewayAuthServiceImpl {
    async fn authorize(
        &self,
        request: Request<AuthorizeRequest>,
    ) -> Result<Response<AuthorizeResponse>, Status> {
        let req = request.into_inner();
        let mut token = req.token.trim();
        let mut parts = token.split_whitespace();
        if let Some(scheme) = parts.next() {
            if scheme.eq_ignore_ascii_case("bearer") {
                if let Some(value) = parts.next() {
                    if parts.next().is_none() {
                        token = value.trim();
                    }
                }
            }
        }

        if token.is_empty() {
            return Ok(Response::new(deny_response("missing token")));
        }

        let claims = match self.token_service.verify(token).await {
            Ok(claims) => claims,
            Err(err) => return Ok(Response::new(deny_response(err.to_string()))),
        };

        if let Some(reason) = self
            .check_token_revoked(&claims.principal_id, token)
            .await?
        {
            return Ok(Response::new(deny_response(reason)));
        }

        let principal_ref = PrincipalRef::new(claims.principal_kind.clone(), &claims.principal_id);
        let principal = match self.principal_store.get(&principal_ref).await {
            Ok(Some(principal)) => principal,
            Ok(None) => return Ok(Response::new(deny_response("principal not found"))),
            Err(err) => {
                return Err(Status::internal(format!(
                    "failed to read principal: {}",
                    err
                )))
            }
        };

        if !principal.enabled {
            return Ok(Response::new(deny_response("principal disabled")));
        }

        let (action, resource, context, org_id, project_id) =
            match build_authz_request(&req, &claims, &principal) {
                Ok(values) => values,
                Err(reason) => return Ok(Response::new(deny_response(reason))),
            };
        let authz_request =
            AuthzRequest::new(principal.clone(), action, resource).with_context(context);
        let decision = self
            .evaluator
            .evaluate(&authz_request)
            .await
            .map_err(|e| Status::internal(format!("authz evaluation failed: {}", e)))?;

        match decision {
            AuthzDecision::Allow => {}
            AuthzDecision::Deny { reason } => {
                return Ok(Response::new(deny_response(reason)));
            }
        }

        let subject = Subject {
            subject_id: claims.principal_id.clone(),
            org_id,
            project_id,
            roles: claims.roles.clone(),
            scopes: vec![claims.scope.to_string()],
        };

        let ttl_seconds = ttl_from_claims(claims.exp);
        let mut headers = HashMap::new();
        headers.insert("x-iam-session-id".to_string(), claims.session_id.clone());
        headers.insert(
            "x-iam-principal-kind".to_string(),
            claims.principal_kind.to_string(),
        );
        headers.insert(
            "x-iam-auth-method".to_string(),
            claims.auth_method.to_string(),
        );

        Ok(Response::new(AuthorizeResponse {
            allow: true,
            reason: String::new(),
            subject: Some(subject),
            headers,
            ttl_seconds,
        }))
    }
}

fn deny_response(reason: impl Into<String>) -> AuthorizeResponse {
    AuthorizeResponse {
        allow: false,
        reason: reason.into(),
        subject: None,
        headers: HashMap::new(),
        ttl_seconds: 0,
    }
}

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 ttl_from_claims(exp: u64) -> u32 {
    let now = now_ts();
    let remaining = exp.saturating_sub(now);
    u32::try_from(remaining).unwrap_or(u32::MAX)
}

fn build_authz_request(
    req: &AuthorizeRequest,
    claims: &InternalTokenClaims,
    principal: &Principal,
) -> Result<(String, Resource, AuthzContext, String, String), String> {
    let action = action_for_request(req);
    let (org_id, project_id) = resolve_org_project(req, claims, principal)?;
    let mut resource = Resource::new(
        "gateway_route",
        resource_id_for_request(req),
        org_id.clone(),
        project_id.clone(),
    );
    resource = resource
        .with_tag("route", req.route_name.clone())
        .with_tag("method", req.method.clone())
        .with_tag("path", req.path.clone());
    if !req.raw_query.is_empty() {
        resource = resource.with_tag("raw_query", req.raw_query.clone());
    }

    let mut context = AuthzContext::new()
        .with_http_method(req.method.clone())
        .with_request_path(req.path.clone())
        .with_metadata("route", req.route_name.clone())
        .with_metadata("request_id", req.request_id.clone())
        .with_metadata("org_id", org_id.clone())
        .with_metadata("project_id", project_id.clone());
    if !req.raw_query.is_empty() {
        context = context.with_metadata("raw_query", req.raw_query.clone());
    }
    if let Ok(ip) = req.client_ip.parse() {
        context = context.with_source_ip(ip);
    }

    Ok((action, resource, context, org_id, project_id))
}

fn action_for_request(req: &AuthorizeRequest) -> String {
    let route = if req.route_name.trim().is_empty() {
        "gateway"
    } else {
        req.route_name.trim()
    };
    let verb = method_to_verb(&req.method);
    format!("gateway:{}:{}", normalize_action_component(route), verb)
}

fn method_to_verb(method: &str) -> &'static str {
    match method.trim().to_uppercase().as_str() {
        "GET" | "HEAD" => "read",
        "POST" => "create",
        "PUT" | "PATCH" => "update",
        "DELETE" => "delete",
        "OPTIONS" => "list",
        _ => "execute",
    }
}

fn normalize_action_component(value: &str) -> String {
    value
        .chars()
        .map(|ch| {
            if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
                ch.to_ascii_lowercase()
            } else {
                '_'
            }
        })
        .collect()
}

fn resource_id_for_request(req: &AuthorizeRequest) -> String {
    if !req.route_name.trim().is_empty() {
        return req.route_name.trim().to_string();
    }
    let path = req.path.trim_matches('/');
    if path.is_empty() {
        "root".to_string()
    } else {
        path.replace('/', ":")
    }
}

fn resolve_org_project(
    req: &AuthorizeRequest,
    claims: &InternalTokenClaims,
    principal: &Principal,
) -> Result<(String, String), String> {
    let allow_header_override = allow_header_tenant_override();
    let org_id = claims
        .org_id
        .clone()
        .or_else(|| claims.scope.org_id().map(|value| value.to_string()))
        .or_else(|| principal.org_id.clone())
        .or_else(|| {
            if allow_header_override {
                header_value(&req.headers, "x-org-id")
            } else {
                None
            }
        })
        .ok_or_else(|| "tenant resolution failed: missing org_id".to_string())?;

    let project_id = claims
        .project_id
        .clone()
        .or_else(|| claims.scope.project_id().map(|value| value.to_string()))
        .or_else(|| principal.project_id.clone())
        .or_else(|| {
            if allow_header_override {
                header_value(&req.headers, "x-project-id")
            } else {
                None
            }
        })
        .ok_or_else(|| "tenant resolution failed: missing project_id".to_string())?;

    Ok((org_id, project_id))
}

fn allow_header_tenant_override() -> bool {
    std::env::var("IAM_GATEWAY_ALLOW_HEADER_TENANT")
        .or_else(|_| std::env::var("PHOTON_IAM_GATEWAY_ALLOW_HEADER_TENANT"))
        .ok()
        .map(|value| {
            matches!(
                value.trim().to_lowercase().as_str(),
                "1" | "true" | "yes" | "y" | "on"
            )
        })
        .unwrap_or(false)
}

fn header_value(headers: &HashMap<String, String>, key: &str) -> Option<String> {
    headers
        .get(&key.to_ascii_lowercase())
        .map(|value| value.trim().to_string())
        .filter(|value| !value.is_empty())
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use iam_authn::{InternalTokenConfig, SigningKey};
    use iam_authz::{PolicyCache, PolicyEvaluator};
    use iam_store::{Backend, BackendConfig, BindingStore, PrincipalStore, RoleStore, TokenStore};
    use iam_types::{
        Permission, PolicyBinding, Principal, PrincipalRef, Role, Scope, TokenMetadata, TokenType,
    };
    use std::time::Duration;

    fn make_request(token: &str) -> AuthorizeRequest {
        AuthorizeRequest {
            request_id: "req-1".into(),
            token: token.to_string(),
            method: "GET".into(),
            path: "/v1/example".into(),
            raw_query: "".into(),
            headers: HashMap::new(),
            client_ip: "127.0.0.1".into(),
            route_name: "example".into(),
        }
    }

    async fn build_service() -> (
        GatewayAuthServiceImpl,
        Arc<InternalTokenService>,
        Arc<RoleStore>,
        Arc<BindingStore>,
        Arc<TokenStore>,
        Principal,
    ) {
        let backend = Arc::new(Backend::new(BackendConfig::Memory).await.unwrap());
        let principal_store = Arc::new(PrincipalStore::new(backend.clone()));
        let role_store = Arc::new(RoleStore::new(backend.clone()));
        let binding_store = Arc::new(BindingStore::new(backend.clone()));
        let token_store = Arc::new(TokenStore::new(backend));
        let signing_key = SigningKey::generate("test-key-1");
        let token_config = InternalTokenConfig::new(signing_key, "iam-test")
            .with_default_ttl(Duration::from_secs(3600))
            .with_max_ttl(Duration::from_secs(7200));
        let token_service = Arc::new(InternalTokenService::new(token_config));
        let mut principal = Principal::new_user("user-1", "User One");
        principal.org_id = Some("org-1".into());
        principal.project_id = Some("proj-1".into());
        principal_store.create(&principal).await.unwrap();
        let cache = Arc::new(PolicyCache::default_config());
        let evaluator = Arc::new(PolicyEvaluator::new(
            binding_store.clone(),
            role_store.clone(),
            cache,
        ));
        let service = GatewayAuthServiceImpl::new(
            token_service.clone(),
            principal_store.clone(),
            token_store.clone(),
            evaluator,
        );
        (
            service,
            token_service,
            role_store,
            binding_store,
            token_store,
            principal,
        )
    }

    #[tokio::test]
    async fn test_authorize_missing_token_denies() {
        let (service, _, _, _, _, _) = build_service().await;
        let response = service
            .authorize(Request::new(make_request("")))
            .await
            .unwrap()
            .into_inner();
        assert!(!response.allow);
        assert!(response.reason.contains("missing token"));
    }

    #[tokio::test]
    async fn test_authorize_valid_token_allows() {
        let (service, token_service, role_store, binding_store, _, principal) =
            build_service().await;
        let role = Role::new(
            "GatewayReader",
            Scope::project("proj-1", "org-1"),
            vec![Permission::new("gateway:example:read", "*")],
        );
        role_store.create(&role).await.unwrap();
        let binding = PolicyBinding::new(
            "binding-1",
            PrincipalRef::new(principal.kind.clone(), principal.id.clone()),
            role.to_ref(),
            Scope::project("proj-1", "org-1"),
        );
        binding_store.create(&binding).await.unwrap();
        let issued = token_service
            .issue(&principal, vec!["role-1".into()], Scope::system(), None)
            .await
            .unwrap();
        let response = service
            .authorize(Request::new(make_request(&issued.token)))
            .await
            .unwrap()
            .into_inner();
        assert!(response.allow);
        let subject = response.subject.expect("subject");
        assert_eq!(subject.subject_id, principal.id);
        assert_eq!(subject.roles, vec!["role-1".to_string()]);
        assert_eq!(subject.scopes, vec!["system".to_string()]);
        assert!(response.ttl_seconds > 0);
    }

    #[tokio::test]
    async fn test_authorize_revoked_token_denies() {
        let (service, token_service, _, _, token_store, principal) = build_service().await;
        let issued = token_service
            .issue(&principal, vec![], Scope::system(), None)
            .await
            .unwrap();
        let token_id = compute_token_id(&issued.token);
        let meta = TokenMetadata::new(
            &token_id,
            &issued.claims.principal_id,
            TokenType::Access,
            issued.claims.iat,
            issued.claims.exp,
        );
        token_store.put(&meta).await.unwrap();
        token_store
            .revoke(
                &issued.claims.principal_id,
                &token_id,
                "test revoke",
                now_ts(),
            )
            .await
            .unwrap();

        let response = service
            .authorize(Request::new(make_request(&issued.token)))
            .await
            .unwrap()
            .into_inner();
        assert!(!response.allow);
        assert!(response.reason.contains("revoke"));
    }
}