13aafa3
use crate::models::users::user_events::UserEventType;
use crate::models::users::{user_ssh_keys, users};
use crate::protos::auth::common::v1::{AuthMethod, PageResponse};
use crate::protos::auth::v1::ssh_key_service_server::{SshKeyService, SshKeyServiceServer};
use crate::protos::auth::v1::*;
use sea_orm::sea_query::{Expr, ExprTrait};
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder, QuerySelect, TransactionTrait};
use tonic::{Request, Response, Status};
use uuid::Uuid;
use crate::AuthService;
use crate::api::{context, convert, error, store};
use crate::errors::AuthError;
#[derive(Clone)]
pub struct SshKeyApi {
svc: AuthService,
}
impl SshKeyApi {
pub fn new(svc: AuthService) -> Self {
Self { svc }
}
pub fn server(
self,
) -> tonic::service::interceptor::InterceptedService<
SshKeyServiceServer<Self>,
context::Interceptor,
> {
let gate = context::Interceptor::new(self.svc.auth().rpc_key.clone());
SshKeyServiceServer::with_interceptor(self, gate)
}
}
fn not_found() -> LookupSshKeyResponse {
LookupSshKeyResponse {
found: false,
..Default::default()
}
}
#[tonic::async_trait]
impl SshKeyService for SshKeyApi {
async fn lookup_ssh_key(
&self,
request: Request<LookupSshKeyRequest>,
) -> Result<Response<LookupSshKeyResponse>, Status> {
let meta = context::meta(&request).await?;
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.SshKeyService");
let result = self
.lookup_ssh_key_inner(meta.request_id(), request.into_inner())
.await;
timer.finish("LookupSshKey", &result);
result
}
async fn add_ssh_key(
&self,
request: Request<AddSshKeyRequest>,
) -> Result<Response<AddSshKeyResponse>, Status> {
let meta = context::meta(&request).await?;
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.SshKeyService");
let result = self
.add_ssh_key_inner(meta.request_id(), request.into_inner())
.await;
timer.finish("AddSshKey", &result);
result
}
async fn list_user_ssh_keys(
&self,
request: Request<ListUserSshKeysRequest>,
) -> Result<Response<ListUserSshKeysResponse>, Status> {
let meta = context::meta(&request).await?;
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.SshKeyService");
let result = self
.list_user_ssh_keys_inner(meta.request_id(), request.into_inner())
.await;
timer.finish("ListUserSshKeys", &result);
result
}
async fn revoke_ssh_key(
&self,
request: Request<RevokeSshKeyRequest>,
) -> Result<Response<RevokeSshKeyResponse>, Status> {
let meta = context::meta(&request).await?;
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.SshKeyService");
let result = self
.revoke_ssh_key_inner(meta.request_id(), request.into_inner())
.await;
timer.finish("RevokeSshKey", &result);
result
}
}
impl SshKeyApi {
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn lookup_ssh_key_inner(
&self,
request_id: Option<&str>,
request: LookupSshKeyRequest,
) -> Result<Response<LookupSshKeyResponse>, Status> {
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.db().read())
.await
.map_err(|err| error::status(&AuthError::from(err), request_id))?;
let Some(row) = row else {
return Ok(Response::new(not_found()));
};
if request.require_active
&& (row.revoked_at.is_some()
|| row
.expires_at
.is_some_and(|expires| expires <= store::now()))
{
return Ok(Response::new(not_found()));
}
let user = if request.require_active {
match store::load_active_user(self.svc.db(), row.user_id).await {
Ok(user) => user,
Err(
AuthError::AccountUnavailable { .. }
| AuthError::LoginForbidden { .. }
| AuthError::UserNotFound,
) => return Ok(Response::new(not_found())),
Err(err) => return Err(error::status(&err, request_id)),
}
} else {
store::load_user(self.svc.db(), row.user_id)
.await
.map_err(|err| error::status(&err, request_id))?
};
Ok(Response::new(LookupSshKeyResponse {
found: true,
key: Some(convert::ssh_key_info(&row)),
principal: Some(store::principal(&user, AuthMethod::SshKey)),
current_security_version: user.security_version,
}))
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn add_ssh_key_inner(
&self,
request_id: Option<&str>,
request: AddSshKeyRequest,
) -> Result<Response<AddSshKeyResponse>, Status> {
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.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(Response::new(AddSshKeyResponse {
key: Some(convert::ssh_key_info(&row)),
}))
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn list_user_ssh_keys_inner(
&self,
request_id: Option<&str>,
request: ListUserSshKeysRequest,
) -> Result<Response<ListUserSshKeysResponse>, Status> {
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.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(convert::ssh_key_info).collect();
Ok(Response::new(ListUserSshKeysResponse {
keys,
page: Some(PageResponse { next_page_token }),
}))
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn revoke_ssh_key_inner(
&self,
request_id: Option<&str>,
request: RevokeSshKeyRequest,
) -> Result<Response<RevokeSshKeyResponse>, Status> {
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.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.db(), request.user_id)
.await
.map_err(|err| error::status(&err, request_id))?;
return Ok(Response::new(RevokeSshKeyResponse {
revoked: false,
security_version: user.security_version,
}));
};
if row.revoked_at.is_some() {
let user = store::load_user(self.svc.db(), request.user_id)
.await
.map_err(|err| error::status(&err, request_id))?;
return Ok(Response::new(RevokeSshKeyResponse {
revoked: false,
security_version: user.security_version,
}));
}
let revoked_at = store::now();
// 撤销与版本提升必须在同一事务：分身只提交一半会让已撤销的 key 继续可用。
let txn = self
.svc
.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.db(), request.user_id)
.await
.map_err(|err| error::status(&err, request_id))?;
return Ok(Response::new(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 {
crate::services::me::invalidate_me_cache(self.svc.cache(), request.user_id).await;
}
if let Err(err) = store::audit(
self.svc.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.db(), request.user_id)
.await
.map_err(|err| error::status(&err, request_id))?;
Ok(Response::new(RevokeSshKeyResponse {
revoked: true,
security_version: user.security_version,
}))
}
}