use auth::models::users::user_events::UserEventType; use auth::models::users::{user_ssh_keys, users}; use auth::protos::auth::common::v1::{AuthMethod, PageResponse}; use auth::protos::auth::v1::*; use export::KenAuthError; use sea_orm::sea_query::{Expr, ExprTrait}; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder, QuerySelect, TransactionTrait}; use uuid::Uuid; use crate::services::ssh_key::SshKeyService; use auth::errors::AuthError; use auth::share::{auth_convert as convert, auth_error as error, auth_store as store}; #[derive(Clone)] pub struct SshKeys { svc: SshKeyService, } impl SshKeys { pub fn new(svc: SshKeyService) -> Self { Self { svc } } } fn not_found() -> LookupSshKeyResponse { LookupSshKeyResponse { found: false, ..Default::default() } } impl SshKeys { #[tracing::instrument(skip_all, fields(request_id = %export::middleware::resolve_request_id(request_id)))] pub async fn lookup_ssh_key_inner( &self, request_id: Option<&str>, request: LookupSshKeyRequest, ) -> Result { let fingerprint = if request.fingerprint.trim().is_empty() { let public_key = request.public_key.trim(); if public_key.is_empty() { return Err(error::invalid_argument( "fingerprint or public_key is required", request_id, )); } crate::services::ssh_key::service::parse_public_key(public_key) .map_err(|err| error::status(&err, request_id))? .1 } else { request.fingerprint.trim().to_owned() }; let row = user_ssh_keys::Entity::find() .filter(user_ssh_keys::Column::Fingerprint.eq(fingerprint)) .one(self.svc.accounts.db().read()) .await .map_err(|err| error::status(&AuthError::from(err), request_id))?; let Some(row) = row else { return Ok(not_found()); }; if request.require_active && (row.revoked_at.is_some() || row .expires_at .is_some_and(|expires| expires <= store::now())) { return Ok(not_found()); } let user = if request.require_active { match store::load_active_user(self.svc.accounts.db(), row.user_id).await { Ok(user) => user, Err( AuthError::AccountUnavailable { .. } | AuthError::LoginForbidden { .. } | AuthError::UserNotFound, ) => return Ok(not_found()), Err(err) => return Err(error::status(&err, request_id)), } } else { store::load_user(self.svc.accounts.db(), row.user_id) .await .map_err(|err| error::status(&err, request_id))? }; Ok(LookupSshKeyResponse { found: true, key: Some(crate::share::ssh_keys::ssh_key_info(&row)), principal: Some(store::principal(&user, AuthMethod::SshKey)), current_security_version: user.security_version, }) } #[tracing::instrument(skip_all, fields(request_id = %export::middleware::resolve_request_id(request_id)))] pub async fn add_ssh_key_inner( &self, request_id: Option<&str>, request: AddSshKeyRequest, ) -> Result { let created = self .svc .add_ssh_key_for_user( request.user_id, crate::services::ssh_key::types::SshKeyAddRequest { name: request.name, description: Some(request.description).filter(|text| !text.is_empty()), public_key: request.public_key, expires_at: None, }, None, Some(request.client_ip).filter(|ip| !ip.is_empty()), request_id.unwrap_or_default().to_owned(), ) .await .map_err(|err| error::status(&err, request_id))?; let row = user_ssh_keys::Entity::find_by_id(created.id) .one(self.svc.accounts.db().read()) .await .map_err(|err| error::status(&AuthError::from(err), request_id))? .ok_or_else(|| error::unavailable("ssh key vanished after insert", request_id))?; Ok(AddSshKeyResponse { key: Some(crate::share::ssh_keys::ssh_key_info(&row)), }) } #[tracing::instrument(skip_all, fields(request_id = %export::middleware::resolve_request_id(request_id)))] pub async fn list_user_ssh_keys_inner( &self, request_id: Option<&str>, request: ListUserSshKeysRequest, ) -> Result { let page = request.page.unwrap_or_default(); let offset = match store::parse_page_token(&page.page_token) { Some(offset) => offset, None => { return Err(error::invalid_argument( "page_token is malformed", request_id, )); } }; let page_size = convert::page_size(page.page_size); let mut query = user_ssh_keys::Entity::find().filter(user_ssh_keys::Column::UserId.eq(request.user_id)); if !request.include_revoked { query = query.filter(user_ssh_keys::Column::RevokedAt.is_null()); } let rows = query .order_by_desc(user_ssh_keys::Column::CreatedAt) .limit(page_size as u64) .offset(offset) .all(self.svc.accounts.db().read()) .await .map_err(|err| error::status(&AuthError::from(err), request_id))?; let next_page_token = if convert::has_more(offset, page_size, rows.len()) { store::page_token(offset + rows.len() as u64) } else { String::new() }; let keys = rows .iter() .map(crate::share::ssh_keys::ssh_key_info) .collect(); Ok(ListUserSshKeysResponse { keys, page: Some(PageResponse { next_page_token }), }) } #[tracing::instrument(skip_all, fields(request_id = %export::middleware::resolve_request_id(request_id)))] pub async fn revoke_ssh_key_inner( &self, request_id: Option<&str>, request: RevokeSshKeyRequest, ) -> Result { let key_id = Uuid::parse_str(request.key_id.trim()) .map_err(|_| error::invalid_argument("key_id is not a valid uuid", request_id))?; let row = user_ssh_keys::Entity::find_by_id(key_id) .filter(user_ssh_keys::Column::UserId.eq(request.user_id)) .one(self.svc.accounts.db().write()) .await .map_err(|err| error::status(&AuthError::from(err), request_id))?; let Some(row) = row else { let user = store::load_user(self.svc.accounts.db(), request.user_id) .await .map_err(|err| error::status(&err, request_id))?; return Ok(RevokeSshKeyResponse { revoked: false, security_version: user.security_version, }); }; if row.revoked_at.is_some() { let user = store::load_user(self.svc.accounts.db(), request.user_id) .await .map_err(|err| error::status(&err, request_id))?; return Ok(RevokeSshKeyResponse { revoked: false, security_version: user.security_version, }); } let revoked_at = store::now(); // 撤销与版本提升必须在同一事务:分身只提交一半会让已撤销的 key 继续可用。 let txn = self .svc .accounts .db() .write() .begin() .await .map_err(|err| error::status(&AuthError::from(err), request_id))?; let updated = user_ssh_keys::Entity::update_many() .col_expr(user_ssh_keys::Column::RevokedAt, Expr::value(revoked_at)) .col_expr( user_ssh_keys::Column::Version, Expr::col(user_ssh_keys::Column::Version).add(1), ) .filter(user_ssh_keys::Column::Id.eq(key_id)) .filter(user_ssh_keys::Column::RevokedAt.is_null()) .exec(&txn) .await .map_err(|err| error::status(&AuthError::from(err), request_id))?; if updated.rows_affected != 1 { let _ = txn.rollback().await; let user = store::load_user(self.svc.accounts.db(), request.user_id) .await .map_err(|err| error::status(&err, request_id))?; return Ok(RevokeSshKeyResponse { revoked: false, security_version: user.security_version, }); } if request.bump_security_version { users::Entity::update_many() .col_expr( users::Column::SecurityVersion, Expr::col(users::Column::SecurityVersion).add(1), ) .col_expr(users::Column::UpdatedAt, Expr::value(revoked_at)) .filter(users::Column::Id.eq(request.user_id)) .exec(&txn) .await .map_err(|err| error::status(&AuthError::from(err), request_id))?; } txn.commit() .await .map_err(|err| error::status(&AuthError::from(err), request_id))?; if request.bump_security_version { auth::services::me::invalidate_me_cache(self.svc.accounts.cache(), request.user_id) .await; } if let Err(err) = store::audit( self.svc.accounts.db(), request.user_id, None, UserEventType::PublicKeyRemoved, Some(serde_json::json!({ "ssh_key_id": key_id.to_string(), "name": row.name, "fingerprint": row.fingerprint, "reason": request.reason, "security_version_bumped": request.bump_security_version, })), request_id, None, None, ) .await { // 撤销与版本提升已提交,审计失败不再回滚。 tracing::warn!(%err, ssh_key_id = %key_id, "failed to record ssh key revocation"); } let user = store::load_user(self.svc.accounts.db(), request.user_id) .await .map_err(|err| error::status(&err, request_id))?; Ok(RevokeSshKeyResponse { revoked: true, security_version: user.security_version, }) } }