i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 17:37:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: harden bulk wallet balance adjustment

Browse Source
This commit is contained in:
RWDai
2026-09-23 16:29:47 +08:00
parent e25e240d16
commit 30bb0c3130
12 changed files with 370 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/users/batch.rs
+140 -28
View File
@@ -1,6 +1,7 @@
use super::{
build_admin_users_bad_request_response, build_admin_users_permission_denied_response,
build_admin_users_read_only_response, disabled_user_policy_detail, disabled_user_policy_field,
management_token_may_adjust_admin_wallet_balance,
management_token_may_administer_user_accounts, normalize_admin_user_role,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
@@ -173,6 +174,13 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
"当前为只读模式，无法批量更新用户钱包",
));
}
if mutation.wallet_balance_adjustment.is_some()
&& !management_token_may_adjust_admin_wallet_balance(request_context)
{
return Ok(build_admin_users_permission_denied_response(
request_context,
));
}
if mutation.wallet_balance_adjustment.is_some() && !state.has_auth_wallet_write_capability() {
return Ok(build_admin_users_read_only_response(
"当前为只读模式，无法批量调整用户钱包余额",
@@ -195,9 +203,29 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
.iter()
.map(|user_id| json!({ "user_id": user_id, "reason": "用户不存在或已删除" }))
.collect::<Vec<_>>();
let mut completed_user_ids = Vec::new();
let mut uncertain_user_ids = Vec::new();
let mut unprocessed_user_ids = Vec::new();
let mut interrupted = false;
for item in &resolved.items {
if state.find_user_auth_by_id(&item.user_id).await?.is_none() {
for (item_index, item) in resolved.items.iter().enumerate() {
let user = match state.find_user_auth_by_id(&item.user_id).await {
Ok(user) => user,
Err(_) => {
record_batch_action_interruption(
&resolved.items,
item_index,
false,
"读取用户状态失败，批次已中止，该用户未执行",
&mut failures,
&mut uncertain_user_ids,
&mut unprocessed_user_ids,
);
interrupted = true;
break;
}
};
if user.is_none() {
failures.push(json!({
"user_id": item.user_id,
"reason": "用户不存在或已删除",
@@ -220,34 +248,66 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
}
if let Some(unlimited) = mutation.unlimited {
if !apply_batch_user_wallet_limit_mode(state, &item.user_id, unlimited).await? {
failures.push(json!({
"user_id": item.user_id,
"reason": "用户钱包不可用",
}));
continue;
match apply_batch_user_wallet_limit_mode(state, &item.user_id, unlimited).await {
Ok(true) => {}
Ok(false) => {
failures.push(json!({
"user_id": item.user_id,
"reason": "用户钱包不可用",
}));
continue;
}
Err(_) => {
record_batch_action_interruption(
&resolved.items,
item_index,
true,
"用户钱包更新结果未确认，批次已中止，请核对钱包后再重试",
&mut failures,
&mut uncertain_user_ids,
&mut unprocessed_user_ids,
);
interrupted = true;
break;
}
}
}
if let Some(adjustment) = mutation.wallet_balance_adjustment {
if !apply_batch_user_wallet_balance_adjustment(
match apply_batch_user_wallet_balance_adjustment(
state,
&item.user_id,
adjustment,
current_admin_user_id,
)
.await?
.await
{
failures.push(json!({
"user_id": item.user_id,
"reason": "用户钱包不可用",
}));
continue;
Ok(true) => {}
Ok(false) => {
failures.push(json!({
"user_id": item.user_id,
"reason": "用户钱包不可用",
}));
continue;
}
Err(_) => {
record_batch_action_interruption(
&resolved.items,
item_index,
true,
"余额调整结果未确认，批次已中止，请核对钱包后再重试",
&mut failures,
&mut uncertain_user_ids,
&mut unprocessed_user_ids,
);
interrupted = true;
break;
}
}
}
if mutation.has_auth_user_fields()
&& state
if mutation.has_auth_user_fields() {
let updated_user = match state
.update_local_auth_user_admin_fields(
&item.user_id,
mutation.role.clone(),
@@ -261,22 +321,39 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
None,
mutation.is_active,
)
.await?
.is_none()
{
failures.push(json!({
"user_id": item.user_id,
"reason": "用户不存在或已删除",
}));
continue;
.await
{
Ok(user) => user,
Err(_) => {
record_batch_action_interruption(
&resolved.items,
item_index,
true,
"用户更新结果未确认，批次已中止，请核对后再重试",
&mut failures,
&mut uncertain_user_ids,
&mut unprocessed_user_ids,
);
interrupted = true;
break;
}
};
if updated_user.is_none() {
failures.push(json!({
"user_id": item.user_id,
"reason": "用户不存在或已删除",
}));
continue;
}
}
success += 1;
completed_user_ids.push(item.user_id.clone());
}
let failed = failures.len();
let total = success + failed;
let response = Json(json!({
let mut response_payload = json!({
"total": total,
"success": success,
"failed": failed,
@@ -284,8 +361,14 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
"warnings": resolved.warnings,
"action": request.action.trim().to_ascii_lowercase(),
"modified_fields": mutation.modified_fields,
}))
.into_response();
"interrupted": interrupted,
});
if interrupted {
response_payload["completed_user_ids"] = json!(completed_user_ids);
response_payload["uncertain_user_ids"] = json!(uncertain_user_ids);
response_payload["unprocessed_user_ids"] = json!(unprocessed_user_ids);
}
let response = Json(response_payload).into_response();
Ok(attach_admin_audit_response(
response,
@@ -296,6 +379,35 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
))
}
fn record_batch_action_interruption(
items: &[AdminUserSelectionItem],
item_index: usize,
current_result_uncertain: bool,
reason: &str,
failures: &mut Vec<Value>,
uncertain_user_ids: &mut Vec<String>,
unprocessed_user_ids: &mut Vec<String>,
) {
let current_item = &items[item_index];
failures.push(json!({
"user_id": current_item.user_id,
"reason": reason,
}));
if current_result_uncertain {
uncertain_user_ids.push(current_item.user_id.clone());
} else {
unprocessed_user_ids.push(current_item.user_id.clone());
}
for item in items.iter().skip(item_index + 1) {
failures.push(json!({
"user_id": item.user_id,
"reason": "因前序错误未执行",
}));
unprocessed_user_ids.push(item.user_id.clone());
}
}
fn parse_resolve_selection_request(
request_body: Option<&Bytes>,
) -> Result<AdminUserSelectionRequest, String> {
apps/aether-gateway/src/handlers/admin/users/mod.rs
+5 -4
View File
@@ -51,10 +51,11 @@ use self::shared::{
build_admin_users_data_unavailable_response, build_admin_users_permission_denied_response,
build_admin_users_read_only_response, disabled_user_policy_detail, disabled_user_policy_field,
format_optional_datetime_iso8601, legacy_admin_list_policy_mode,
legacy_admin_rate_limit_policy_mode, management_token_may_administer_user_accounts,
normalize_admin_optional_user_email, normalize_admin_user_group_ids, normalize_admin_user_role,
normalize_admin_username, validate_admin_user_password, AdminCreateUserApiKeyRequest,
AdminCreateUserRequest, AdminToggleUserApiKeyLockRequest, AdminUpdateUserApiKeyRequest,
legacy_admin_rate_limit_policy_mode, management_token_may_adjust_admin_wallet_balance,
management_token_may_administer_user_accounts, normalize_admin_optional_user_email,
normalize_admin_user_group_ids, normalize_admin_user_role, normalize_admin_username,
validate_admin_user_password, AdminCreateUserApiKeyRequest, AdminCreateUserRequest,
AdminToggleUserApiKeyLockRequest, AdminUpdateUserApiKeyRequest,
};
pub(crate) use self::shared::{
normalize_admin_list_policy_mode, normalize_admin_rate_limit_policy_mode,
apps/aether-gateway/src/handlers/admin/users/shared.rs
+12
View File
@@ -165,6 +165,18 @@ pub(super) fn management_token_may_administer_user_accounts(
})
}
pub(super) fn management_token_may_adjust_admin_wallet_balance(
request_context: &crate::handlers::admin::request::AdminRequestContext<'_>,
) -> bool {
request_context.decision().is_some_and(|decision| {
crate::control::management_token_principal_has_permission(decision, "admin:wallets:write")
|| crate::control::management_token_principal_has_permission(
decision,
"admin:wallets:admin",
)
})
}
pub(super) fn build_admin_users_permission_denied_response(
request_context: &crate::handlers::admin::request::AdminRequestContext<'_>,
) -> Response<Body> {
apps/aether-gateway/src/state/app.rs
+2
View File
@@ -492,6 +492,8 @@ pub struct AppState {
Arc<StdMutex<HashMap<String, aether_data::repository::wallet::StoredWalletSnapshot>>>,
>,
#[cfg(test)]
pub(crate) auth_wallet_adjustment_error_for_tests: Option<String>,
#[cfg(test)]
pub(crate) admin_wallet_payment_order_store:
Option<Arc<StdMutex<HashMap<String, AdminWalletPaymentOrderRecord>>>>,
#[cfg(test)]
apps/aether-gateway/src/state/core.rs
+2
View File
@@ -466,6 +466,8 @@ impl AppState {
#[cfg(test)]
auth_wallet_store: Some(Arc::new(StdMutex::new(HashMap::new()))),
#[cfg(test)]
auth_wallet_adjustment_error_for_tests: None,
#[cfg(test)]
admin_wallet_payment_order_store: Some(Arc::new(StdMutex::new(HashMap::new()))),
#[cfg(test)]
admin_payment_callback_store: Some(Arc::new(StdMutex::new(HashMap::new()))),
apps/aether-gateway/src/state/runtime/wallet/balance_mutations.rs
+7
View File
@@ -18,6 +18,13 @@ impl AppState {
)>,
GatewayError,
> {
#[cfg(test)]
if self.auth_wallet_adjustment_error_for_tests.as_deref() == Some(wallet_id) {
return Err(GatewayError::Internal(
"injected test wallet adjustment failure".to_string(),
));
}
#[cfg(test)]
if let Some(store) = self.auth_wallet_store.as_ref() {
let mut guard = store.lock().expect("auth wallet store should lock");
apps/aether-gateway/src/state/testing.rs
+8
View File
@@ -472,6 +472,14 @@ impl AppState {
self
}
pub(crate) fn fail_auth_wallet_adjustment_for_tests(
mut self,
wallet_id: impl Into<String>,
) -> Self {
self.auth_wallet_adjustment_error_for_tests = Some(wallet_id.into());
self
}
pub(crate) fn with_admin_wallet_payment_orders_for_tests<I>(mut self, orders: I) -> Self
where
I: IntoIterator<Item = crate::AdminWalletPaymentOrderRecord>,
apps/aether-gateway/src/tests/control/admin/users_batch.rs
+154 -3
View File
@@ -1,11 +1,17 @@
use aether_data::repository::users::StoredUserAuthRecord;
use std::sync::Arc;
use aether_data::repository::management_tokens::InMemoryManagementTokenRepository;
use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserAuthRecord};
use aether_data::repository::wallet::StoredWalletSnapshot;
use axum::http::StatusCode;
use chrono::Utc;
use reqwest::{Client, RequestBuilder, Response};
use serde_json::{json, Value};
use super::super::{build_router_with_state, start_server, AppState};
use super::super::{
build_router_with_state, hash_management_token, sample_management_token, start_server, AppState,
};
use crate::data::GatewayDataState;
fn admin_headers(request: RequestBuilder) -> RequestBuilder {
request
@@ -19,13 +25,17 @@ fn admin_headers(request: RequestBuilder) -> RequestBuilder {
}
fn sample_user(user_id: &str) -> StoredUserAuthRecord {
sample_user_with_role(user_id, "user")
}
fn sample_user_with_role(user_id: &str, role: &str) -> StoredUserAuthRecord {
StoredUserAuthRecord::new(
user_id.to_string(),
Some(format!("{user_id}@example.com")),
true,
user_id.to_string(),
Some("hash".to_string()),
"user".to_string(),
role.to_string(),
"local".to_string(),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
@@ -162,6 +172,147 @@ async fn gateway_batches_wallet_addition_deduction_and_clamped_deduction_per_use
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_requires_wallet_write_permission_for_batch_balance_adjustments() {
let users_write_token = "ae-batch-users-write-only";
let wallet_write_token = "ae-batch-users-wallet-write";
let token_owner = sample_user_with_role("token-owner", "admin");
let target_user = sample_user("user-1");
let mut users_only = sample_management_token(
"token-users-write-only",
&token_owner.id,
&token_owner.username,
true,
);
users_only.token.allowed_ips = None;
users_only.token.permissions = Some(json!(["admin:users:write"]));
let mut users_and_wallets = sample_management_token(
"token-users-and-wallet-write",
&token_owner.id,
&token_owner.username,
true,
);
users_and_wallets.token.allowed_ips = None;
users_and_wallets.token.permissions = Some(json!(["admin:users:write", "admin:wallets:write"]));
let token_repository = Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
vec![users_only, users_and_wallets],
vec![
(
hash_management_token(users_write_token),
"token-users-write-only".to_string(),
),
(
hash_management_token(wallet_write_token),
"token-users-and-wallet-write".to_string(),
),
],
));
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
token_owner.clone(),
target_user.clone(),
]));
let data = GatewayDataState::with_management_token_repository_for_tests(token_repository)
.with_user_reader(user_repository);
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data)
.with_auth_users_for_tests([token_owner, target_user])
.with_auth_wallets_for_tests([sample_wallet("user-1", 10.0, 0.0)]);
let (gateway_url, gateway_handle) = start_server(build_router_with_state(state)).await;
let client = Client::new();
let payload = json!({
"selection": { "user_ids": ["user-1"] },
"action": "adjust_wallet_balance",
"payload": { "operation": "add", "amount": 5.0 }
});
let denied = client
.post(format!("{gateway_url}/api/admin/users/batch-action"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.bearer_auth(users_write_token)
.json(&payload)
.send()
.await
.expect("users-only management token request should complete");
assert_eq!(denied.status(), StatusCode::FORBIDDEN);
assert_eq!(
wallet_detail(&client, &gateway_url, "user-1").await["balance"],
10.0
);
let allowed = client
.post(format!("{gateway_url}/api/admin/users/batch-action"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.bearer_auth(wallet_write_token)
.json(&payload)
.send()
.await
.expect("wallet-write management token request should complete");
assert_eq!(allowed.status(), StatusCode::OK);
let result: Value = allowed.json().await.expect("response should parse");
assert_eq!(result["success"], 1);
assert_eq!(
wallet_detail(&client, &gateway_url, "user-1").await["balance"],
15.0
);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_reports_completed_uncertain_and_unprocessed_users_after_adjustment_error() {
let state = AppState::new()
.expect("gateway should build")
.with_auth_users_for_tests([
sample_user("user-1"),
sample_user("user-2"),
sample_user("user-3"),
])
.with_auth_wallets_for_tests([
sample_wallet("user-1", 10.0, 0.0),
sample_wallet("user-2", 20.0, 0.0),
sample_wallet("user-3", 30.0, 0.0),
])
.fail_auth_wallet_adjustment_for_tests("wallet-user-2");
let (gateway_url, gateway_handle) = start_server(build_router_with_state(state)).await;
let client = Client::new();
let response = post_batch_action(
&client,
&gateway_url,
json!({
"selection": { "user_ids": ["user-1", "user-2", "user-3"] },
"action": "adjust_wallet_balance",
"payload": { "operation": "add", "amount": 5.0 }
}),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let result: Value = response.json().await.expect("response should parse");
assert_eq!(result["interrupted"], true);
assert_eq!(result["success"], 1);
assert_eq!(result["failed"], 2);
assert_eq!(result["completed_user_ids"], json!(["user-1"]));
assert_eq!(result["uncertain_user_ids"], json!(["user-2"]));
assert_eq!(result["unprocessed_user_ids"], json!(["user-3"]));
assert_eq!(result["failures"][0]["user_id"], "user-2");
assert_eq!(result["failures"][1]["user_id"], "user-3");
assert_eq!(
wallet_detail(&client, &gateway_url, "user-1").await["balance"],
15.0
);
assert_eq!(
wallet_detail(&client, &gateway_url, "user-2").await["balance"],
20.0
);
assert_eq!(
wallet_detail(&client, &gateway_url, "user-3").await["balance"],
30.0
);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_reports_missing_wallet_and_skips_zero_delta_for_non_positive_balance() {
let state = AppState::new()