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

652 lines
22 KiB
Rust

//! 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>,
admin_token: Option<String>,
}
impl IamTokenService {
/// Create a new token service
pub fn new(
token_service: Arc<InternalTokenService>,
principal_store: Arc<PrincipalStore>,
token_store: Arc<TokenStore>,
admin_token: Option<String>,
) -> Self {
Self {
token_service,
principal_store,
token_store,
admin_token,
}
}
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)
}
fn admin_token_valid(metadata: &tonic::metadata::MetadataMap, token: &str) -> bool {
if let Some(value) = metadata.get("x-iam-admin-token") {
if let Ok(raw) = value.to_str() {
if raw.trim() == token {
return true;
}
}
}
if let Some(value) = metadata.get("authorization") {
if let Ok(raw) = value.to_str() {
let raw = raw.trim();
if let Some(rest) = raw.strip_prefix("Bearer ") {
return rest.trim() == token;
}
if let Some(rest) = raw.strip_prefix("bearer ") {
return rest.trim() == token;
}
}
}
false
}
fn require_admin_token<T>(&self, request: &Request<T>) -> Result<(), Status> {
if let Some(token) = self.admin_token.as_deref() {
if !Self::admin_token_valid(request.metadata(), token) {
return Err(Status::unauthenticated(
"missing or invalid IAM admin token",
));
}
}
Ok(())
}
fn validate_scope_for_principal(
principal: &iam_types::Principal,
scope: &TypesScope,
) -> Result<(), Status> {
let principal_org = principal.org_id.as_deref();
let principal_project = principal.project_id.as_deref();
match principal.kind {
TypesPrincipalKind::ServiceAccount => match scope {
TypesScope::System => Err(Status::permission_denied(
"service accounts cannot mint system-scoped tokens",
)),
TypesScope::Org { id } => {
if principal_org == Some(id.as_str()) {
Ok(())
} else {
Err(Status::permission_denied(
"service account token scope must match principal tenant",
))
}
}
TypesScope::Project { id, org_id } => {
if principal_org == Some(org_id.as_str())
&& principal_project == Some(id.as_str())
{
Ok(())
} else {
Err(Status::permission_denied(
"service account token scope must match principal tenant",
))
}
}
TypesScope::Resource {
project_id, org_id, ..
} => {
if principal_org == Some(org_id.as_str())
&& principal_project == Some(project_id.as_str())
{
Ok(())
} else {
Err(Status::permission_denied(
"service account token scope must match principal tenant",
))
}
}
},
_ => {
if let Some(org_id) = principal_org {
match scope {
TypesScope::System => {}
TypesScope::Org { id } if id == org_id => {}
TypesScope::Project {
org_id: scope_org, ..
}
| TypesScope::Resource {
org_id: scope_org, ..
} if scope_org == org_id => {}
_ => {
return Err(Status::permission_denied(
"token scope must match principal tenant",
))
}
}
}
Ok(())
}
}
}
}
#[tonic::async_trait]
impl IamToken for IamTokenService {
async fn issue_token(
&self,
request: Request<IssueTokenRequest>,
) -> Result<Response<IssueTokenResponse>, Status> {
self.require_admin_token(&request)?;
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);
Self::validate_scope_for_principal(&principal, &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, None);
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(),
None,
);
// 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(),
None,
);
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"));
}
#[tokio::test]
async fn test_service_account_system_scope_is_rejected() {
let (token_service, principal_store, token_store) = test_setup();
let service = IamTokenService::new(
token_service,
principal_store.clone(),
token_store.clone(),
None,
);
let principal = Principal::new_service_account("svc-1", "Service 1", "org-1", "proj-1");
principal_store.create(&principal).await.unwrap();
let result = service
.issue_token(Request::new(IssueTokenRequest {
principal_id: "svc-1".into(),
principal_kind: PrincipalKind::ServiceAccount as i32,
roles: vec![],
scope: Some(IamTokenService::convert_scope_to_proto(&TypesScope::System)),
ttl_seconds: 3600,
}))
.await;
assert!(result.is_err());
assert_eq!(result.unwrap_err().code(), tonic::Code::PermissionDenied);
}
}