538d47b
use std::time::Duration;
use futures_util::StreamExt;
use crate::protos::auth::internal::v1::cache_maintenance_service_server::{
CacheMaintenanceService, CacheMaintenanceServiceServer,
};
use crate::protos::auth::internal::v1::*;
use tonic::{Request, Response, Status};
use crate::AuthSvc;
use crate::errors::AuthError;
use crate::internal::{context, error, store};
use crate::me::types::MeResponse;
const USER_SCOPED_PREFIXES: [&str; 3] = ["auth:me:", "auth:admin:", "auth:rate:mfa:"];
const REBUILD_USER_LIMIT: usize = 1000;
const REBUILD_ME_TTL_SECS: u64 = 300;
const REBUILD_CONCURRENCY: usize = 8;
const REBUILD_TOTAL_TIMEOUT_SECS: u64 = 30;
fn scope_prefixes(scope: CacheScope, user_id: i64) -> Vec<String> {
let prefixes: Vec<&str> = match scope {
CacheScope::All => vec![
"auth:me:",
"auth:admin:",
"auth:svc-token:",
"auth:rate:mfa:",
"auth:code:",
"auth:rate:",
"auth:pow:used:",
],
CacheScope::Sessions => vec!["auth:me:", "auth:admin:"],
CacheScope::Introspection => vec!["auth:svc-token:"],
CacheScope::Authorization => vec!["auth:admin:", "auth:me:"],
CacheScope::Mfa => vec!["auth:rate:mfa:"],
CacheScope::Codes => vec!["auth:code:", "auth:rate:"],
CacheScope::Pow => vec!["auth:pow:used:"],
CacheScope::Unspecified => Vec::new(),
};
let mut keys = Vec::with_capacity(prefixes.len());
for prefix in prefixes {
if user_id > 0 && USER_SCOPED_PREFIXES.contains(&prefix) {
keys.push(format!("{prefix}{user_id}"));
} else {
keys.push(prefix.to_owned());
}
}
keys
}
#[derive(Clone)]
pub struct CacheMaintenanceApi {
svc: AuthSvc,
}
impl CacheMaintenanceApi {
pub fn new(svc: AuthSvc) -> Self {
Self { svc }
}
pub fn server(
self,
) -> tonic::service::interceptor::InterceptedService<
CacheMaintenanceServiceServer<Self>,
context::Interceptor,
> {
let gate = context::Interceptor::new(self.svc.auth().rpc_key.clone());
CacheMaintenanceServiceServer::with_interceptor(self, gate)
}
}
#[tonic::async_trait]
impl CacheMaintenanceService for CacheMaintenanceApi {
async fn flush_auth_cache(
&self,
request: Request<FlushAuthCacheRequest>,
) -> Result<Response<FlushAuthCacheResponse>, Status> {
let meta = context::meta(&request).await?;
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.CacheMaintenanceService");
let result = self
.flush_auth_cache_inner(meta.request_id(), request.into_inner())
.await;
timer.finish("FlushAuthCache", &result);
result
}
async fn rebuild_identity_cache(
&self,
request: Request<RebuildIdentityCacheRequest>,
) -> Result<Response<RebuildIdentityCacheResponse>, Status> {
let meta = context::meta(&request).await?;
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.CacheMaintenanceService");
let result = self
.rebuild_identity_cache_inner(meta.request_id(), request.into_inner())
.await;
timer.finish("RebuildIdentityCache", &result);
result
}
}
impl CacheMaintenanceApi {
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn flush_auth_cache_inner(
&self,
request_id: Option<&str>,
request: FlushAuthCacheRequest,
) -> Result<Response<FlushAuthCacheResponse>, Status> {
let request_id = request_id;
if request.scope == CacheScope::Unspecified as i32 {
return Err(error::invalid_argument("scope is required", request_id));
}
let scope = match CacheScope::try_from(request.scope) {
Ok(scope) => scope,
Err(_) => return Err(error::invalid_argument("scope is invalid", request_id)),
};
let prefixes = if request.key_prefix.is_empty() {
scope_prefixes(scope, request.user_id)
} else {
vec![request.key_prefix.clone()]
};
let cache = self.svc.cache();
let mut matched_keys = 0u64;
let mut deleted_keys = 0u64;
if request.user_id > 0 {
for prefix in &prefixes {
if request.dry_run {
let exists = cache
.exists(prefix)
.await
.map_err(|err| error::status(&AuthError::from(err), request_id))?;
if exists {
matched_keys += 1;
}
} else {
let deleted = cache
.delete(prefix)
.await
.map_err(|err| error::status(&AuthError::from(err), request_id))?;
if deleted {
matched_keys += 1;
deleted_keys += 1;
}
}
}
} else {
if cache.is_cluster() {
return Err(error::unavailable(
"global cache flush requires a non-cluster cache",
request_id,
));
}
for prefix in &prefixes {
let count = cache
.count_prefix(prefix)
.await
.map_err(|err| error::status(&AuthError::from(err), request_id))?
as u64;
matched_keys += count;
if !request.dry_run {
deleted_keys += cache
.delete_prefix(prefix)
.await
.map_err(|err| error::status(&AuthError::from(err), request_id))?
as u64;
}
}
}
Ok(Response::new(FlushAuthCacheResponse {
matched_keys,
deleted_keys,
dry_run: request.dry_run,
}))
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn rebuild_identity_cache_inner(
&self,
request_id: Option<&str>,
request: RebuildIdentityCacheRequest,
) -> Result<Response<RebuildIdentityCacheResponse>, Status> {
let request_id = request_id;
if request.user_ids.len() > REBUILD_USER_LIMIT {
return Err(error::invalid_argument(
"user_ids exceeds the limit of 1000",
request_id,
));
}
let queued_users = request.user_ids.len() as u32;
if request.dry_run {
return Ok(Response::new(RebuildIdentityCacheResponse {
queued_users,
rebuilt_users: 0,
skipped_users: 0,
}));
}
let svc = self.svc.clone();
let rebuild = async {
let results = futures_util::stream::iter(request.user_ids)
.map(|user_id| {
let svc = svc.clone();
async move {
let user = match store::load_user(svc.db(), user_id).await {
Ok(user) if user.deleted_at.is_none() => user,
_ => return (user_id, false),
};
let me = MeResponse::from(user);
match svc
.cache()
.set_json_ex(
&format!("auth:me:{user_id}"),
&me,
Duration::from_secs(REBUILD_ME_TTL_SECS),
)
.await
{
Ok(()) => (user_id, true),
Err(err) => {
tracing::warn!(user_id, %err, "identity cache rebuild failed");
(user_id, false)
}
}
}
})
.buffer_unordered(REBUILD_CONCURRENCY)
.collect::<Vec<(i64, bool)>>()
.await;
let rebuilt_users = results.iter().filter(|(_, ok)| *ok).count() as u32;
let skipped_users = results.len() as u32 - rebuilt_users;
Ok(Response::new(RebuildIdentityCacheResponse {
queued_users,
rebuilt_users,
skipped_users,
}))
};
match tokio::time::timeout(Duration::from_secs(REBUILD_TOTAL_TIMEOUT_SECS), rebuild).await {
Ok(response) => response,
Err(_) => Err(error::status_of(
crate::protos::pb::auth::common::v1::Code::Unavailable,
"identity cache rebuild did not finish in time",
request_id,
)),
}
}
}