mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
Merge upstream main
This commit is contained in:
@@ -799,6 +799,7 @@ mod tests {
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
api_key_ip_rules: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -633,6 +633,7 @@ mod tests {
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
api_key_ip_rules: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -336,6 +336,7 @@ mod tests {
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
api_key_ip_rules: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -847,6 +847,18 @@ fn build_openai_image_provider_body_from_openai_responses_body(
|
||||
} else if let Some(value) = object.get("stream") {
|
||||
body.insert("stream".to_string(), value.clone());
|
||||
}
|
||||
let image_tool = tool.clone().unwrap_or_else(|| {
|
||||
let mut tool = serde_json::Map::new();
|
||||
tool.insert(
|
||||
"type".to_string(),
|
||||
Value::String("image_generation".to_string()),
|
||||
);
|
||||
tool
|
||||
});
|
||||
body.insert(
|
||||
"tools".to_string(),
|
||||
Value::Array(vec![Value::Object(image_tool)]),
|
||||
);
|
||||
|
||||
let mut summary = serde_json::Map::new();
|
||||
summary.insert(
|
||||
@@ -1230,7 +1242,7 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn openai_responses_image_bridge_body_does_not_inject_tools() {
|
||||
fn openai_responses_image_bridge_body_preserves_image_generation_tool() {
|
||||
let body_json = json!({
|
||||
"model": "gpt-image-2",
|
||||
"input": "Draw a glass city",
|
||||
@@ -1253,7 +1265,9 @@ mod tests {
|
||||
)
|
||||
.expect("responses image body should convert");
|
||||
|
||||
assert!(provider_body.get("tools").is_none());
|
||||
assert_eq!(provider_body["tools"][0]["type"], "image_generation");
|
||||
assert_eq!(provider_body["tools"][0]["size"], "1024x1024");
|
||||
assert_eq!(provider_body["tools"][0]["output_format"], "png");
|
||||
assert_eq!(provider_body["model"], "gpt-image-2");
|
||||
assert_eq!(provider_body["input"], "Draw a glass city");
|
||||
assert_eq!(provider_body["stream"], true);
|
||||
|
||||
@@ -317,6 +317,12 @@ pub(crate) fn build_local_auth_rejection_response(
|
||||
StatusCode::FORBIDDEN,
|
||||
&format!("{ACCESS_POLICY_SUBJECT}不允许访问模型 {model}"),
|
||||
),
|
||||
GatewayLocalAuthRejection::IpNotAllowed { remote_ip } => build_local_http_error_response(
|
||||
trace_id,
|
||||
control_decision,
|
||||
StatusCode::UNAUTHORIZED,
|
||||
&format!("API Key 不允许从当前 IP 访问: {remote_ip}"),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ pub(crate) enum GatewayLocalAuthRejection {
|
||||
ProviderNotAllowed { provider: String },
|
||||
ApiFormatNotAllowed { api_format: String },
|
||||
ModelNotAllowed { model: String },
|
||||
IpNotAllowed { remote_ip: String },
|
||||
}
|
||||
|
||||
pub(crate) fn trusted_auth_local_rejection(
|
||||
@@ -574,6 +575,7 @@ mod tests {
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: None,
|
||||
allowed_models: Some(allowed_models),
|
||||
ip_rules: None,
|
||||
});
|
||||
decision
|
||||
}
|
||||
|
||||
@@ -48,6 +48,8 @@ pub(crate) struct GatewayControlAuthContext {
|
||||
pub(crate) local_rejection: Option<GatewayLocalAuthRejection>,
|
||||
#[serde(skip)]
|
||||
pub(crate) allowed_models: Option<Vec<String>>,
|
||||
#[serde(skip)]
|
||||
pub(crate) ip_rules: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -207,6 +209,9 @@ fn log_local_auth_rejection(trace_id: &str, decision: &GatewayControlDecision) {
|
||||
GatewayLocalAuthRejection::ModelNotAllowed { model } => {
|
||||
("model_not_allowed", model.clone())
|
||||
}
|
||||
GatewayLocalAuthRejection::IpNotAllowed { remote_ip } => {
|
||||
("ip_not_allowed", remote_ip.clone())
|
||||
}
|
||||
};
|
||||
info!(
|
||||
event_name = "local_auth_rejected",
|
||||
@@ -581,6 +586,7 @@ pub(super) async fn resolve_data_backed_auth_context(
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey),
|
||||
allowed_models: None,
|
||||
ip_rules: None,
|
||||
}));
|
||||
};
|
||||
|
||||
@@ -638,6 +644,7 @@ async fn resolve_trusted_auth_context(
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey),
|
||||
allowed_models: None,
|
||||
ip_rules: None,
|
||||
}));
|
||||
};
|
||||
|
||||
@@ -728,6 +735,7 @@ async fn build_data_backed_auth_context(
|
||||
&& !snapshot.api_key_is_standalone,
|
||||
local_rejection,
|
||||
allowed_models,
|
||||
ip_rules: snapshot.api_key_ip_rules,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -177,6 +177,17 @@ pub(super) fn classify_admin_endpoints_family_route(
|
||||
"admin:endpoints_manage",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/endpoints/providers/")
|
||||
&& normalized_path.ends_with("/key-balance")
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"endpoints_manage",
|
||||
"query_key_balance",
|
||||
"admin:endpoints_manage",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/endpoints/providers/")
|
||||
&& normalized_path.ends_with("/keys")
|
||||
|
||||
@@ -380,6 +380,20 @@ fn classifies_admin_refresh_provider_quota_as_admin_proxy_route() {
|
||||
assert!(!decision.is_execution_runtime_candidate());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_query_provider_key_balance_as_admin_proxy_route() {
|
||||
let headers = http::HeaderMap::new();
|
||||
let uri: Uri = "/api/admin/endpoints/providers/provider-newapi/key-balance"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
|
||||
.expect("decision should resolve");
|
||||
assert_eq!(decision.route_class.as_deref(), Some("admin_proxy"));
|
||||
assert_eq!(decision.route_family.as_deref(), Some("endpoints_manage"));
|
||||
assert_eq!(decision.route_kind.as_deref(), Some("query_key_balance"));
|
||||
assert!(!decision.is_execution_runtime_candidate());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_refresh_provider_quota_buffers_request_body_for_key_selection() {
|
||||
let headers = headers(&[]);
|
||||
@@ -399,6 +413,25 @@ fn admin_refresh_provider_quota_buffers_request_body_for_key_selection() {
|
||||
assert!(local_proxy_route_requires_buffered_body(&context));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_query_provider_key_balance_buffers_request_body_for_key_secret() {
|
||||
let headers = headers(&[]);
|
||||
let uri: Uri = "/api/admin/endpoints/providers/provider-newapi/key-balance"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
|
||||
.expect("decision should resolve");
|
||||
let context = GatewayPublicRequestContext::from_request_parts(
|
||||
"trace-key-balance",
|
||||
&http::Method::POST,
|
||||
&uri,
|
||||
&headers,
|
||||
Some(decision),
|
||||
);
|
||||
|
||||
assert!(local_proxy_route_requires_buffered_body(&context));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_default_body_rules_as_admin_proxy_route() {
|
||||
let headers = headers(&[]);
|
||||
|
||||
@@ -1512,6 +1512,16 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_export_users(
|
||||
&self,
|
||||
query: &aether_data::repository::users::UserExportListQuery,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.count_export_users(query).await,
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_export_users(
|
||||
&self,
|
||||
) -> Result<aether_data::repository::users::UserExportSummary, DataLayerError> {
|
||||
|
||||
@@ -10,7 +10,8 @@ use crate::handlers::admin::users::{
|
||||
default_admin_user_api_key_name, format_optional_unix_secs_iso8601,
|
||||
generate_admin_user_api_key_plaintext, hash_admin_user_api_key, masked_user_api_key_display,
|
||||
normalize_admin_feature_settings, normalize_admin_optional_api_key_name,
|
||||
normalize_admin_user_api_formats, normalize_admin_user_string_list,
|
||||
normalize_admin_user_api_formats, normalize_admin_user_ip_rules,
|
||||
normalize_admin_user_string_list,
|
||||
};
|
||||
use crate::handlers::shared::normalize_optional_api_key_concurrent_limit;
|
||||
use crate::GatewayError;
|
||||
@@ -109,6 +110,10 @@ pub(super) async fn build_admin_create_api_key_response(
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(build_admin_api_keys_bad_request_response(detail)),
|
||||
};
|
||||
let ip_rules = match normalize_admin_user_ip_rules(payload.ip_rules) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(build_admin_api_keys_bad_request_response(detail)),
|
||||
};
|
||||
if payload.rate_limit.is_some_and(|value| value < 0) {
|
||||
return Ok(build_admin_api_keys_bad_request_response(
|
||||
"rate_limit 必须大于等于 0",
|
||||
@@ -162,6 +167,7 @@ pub(super) async fn build_admin_create_api_key_response(
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
ip_rules,
|
||||
rate_limit: payload.rate_limit,
|
||||
concurrent_limit,
|
||||
force_capabilities: None,
|
||||
@@ -327,6 +333,19 @@ pub(super) async fn build_admin_update_api_key_response(
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let ip_rules_present =
|
||||
field_presence.contains("ip_rules") || field_presence.contains("allowed_ips");
|
||||
let ip_rules = if ip_rules_present {
|
||||
match payload.ip_rules {
|
||||
Some(value) => match normalize_admin_user_ip_rules(value) {
|
||||
Ok(value) => Some(value),
|
||||
Err(detail) => return Ok(build_admin_api_keys_bad_request_response(detail)),
|
||||
},
|
||||
None => Some(None),
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let effective_expires_at_unix_secs = if field_presence.contains("expires_at") {
|
||||
match parse_standalone_api_key_expires_at(payload.expires_at.as_deref()) {
|
||||
Ok(value) => value,
|
||||
@@ -394,6 +413,7 @@ pub(super) async fn build_admin_update_api_key_response(
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
ip_rules,
|
||||
expires_at_present: field_presence.contains("expires_at"),
|
||||
expires_at_unix_secs: if field_presence.contains("expires_at") {
|
||||
effective_expires_at_unix_secs
|
||||
|
||||
@@ -3,6 +3,7 @@ use crate::handlers::admin::shared::{query_param_value, AdminTypedObjectPatch};
|
||||
use crate::handlers::admin::users::{
|
||||
format_optional_unix_secs_iso8601, masked_user_api_key_display,
|
||||
};
|
||||
use crate::handlers::shared::deserialize_optional_string_list_patch;
|
||||
use aether_admin::system::serialize_admin_system_users_export_wallet;
|
||||
use axum::{
|
||||
body::Body,
|
||||
@@ -20,6 +21,8 @@ pub(super) struct AdminStandaloneApiKeyCreateRequest {
|
||||
pub(super) allowed_providers: Option<Vec<String>>,
|
||||
pub(super) allowed_api_formats: Option<Vec<String>>,
|
||||
pub(super) allowed_models: Option<Vec<String>>,
|
||||
#[serde(default, alias = "allowed_ips")]
|
||||
pub(super) ip_rules: Option<Vec<String>>,
|
||||
pub(super) rate_limit: Option<i32>,
|
||||
pub(super) concurrent_limit: Option<i32>,
|
||||
pub(super) initial_balance_usd: Option<f64>,
|
||||
@@ -36,6 +39,12 @@ pub(super) struct AdminStandaloneApiKeyUpdateRequest {
|
||||
pub(super) allowed_providers: Option<Vec<String>>,
|
||||
pub(super) allowed_api_formats: Option<Vec<String>>,
|
||||
pub(super) allowed_models: Option<Vec<String>>,
|
||||
#[serde(
|
||||
default,
|
||||
alias = "allowed_ips",
|
||||
deserialize_with = "deserialize_optional_string_list_patch"
|
||||
)]
|
||||
pub(super) ip_rules: Option<Option<Vec<String>>>,
|
||||
pub(super) rate_limit: Option<i32>,
|
||||
pub(super) concurrent_limit: Option<i32>,
|
||||
pub(super) initial_balance_usd: Option<f64>,
|
||||
@@ -161,6 +170,7 @@ pub(super) fn build_admin_api_key_list_item_payload(
|
||||
"allowed_providers": record.allowed_providers,
|
||||
"allowed_api_formats": record.allowed_api_formats,
|
||||
"allowed_models": record.allowed_models,
|
||||
"ip_rules": record.ip_rules,
|
||||
"last_used_at": format_optional_unix_secs_iso8601(record.last_used_at_unix_secs),
|
||||
"expires_at": format_optional_unix_secs_iso8601(record.expires_at_unix_secs),
|
||||
"created_at": format_optional_unix_secs_iso8601(record.created_at_unix_secs),
|
||||
@@ -191,6 +201,7 @@ pub(super) fn build_admin_api_key_detail_payload(
|
||||
"allowed_providers": record.allowed_providers,
|
||||
"allowed_api_formats": record.allowed_api_formats,
|
||||
"allowed_models": record.allowed_models,
|
||||
"ip_rules": record.ip_rules,
|
||||
"last_used_at": format_optional_unix_secs_iso8601(record.last_used_at_unix_secs),
|
||||
"expires_at": format_optional_unix_secs_iso8601(record.expires_at_unix_secs),
|
||||
"created_at": format_optional_unix_secs_iso8601(record.created_at_unix_secs),
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod balance;
|
||||
mod mutations;
|
||||
mod quota;
|
||||
mod reads;
|
||||
@@ -18,6 +19,10 @@ pub(crate) async fn maybe_build_local_admin_endpoints_keys_response(
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
if let Some(response) = balance::maybe_handle(state, request_context, request_body).await? {
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
if let Some(response) = mutations::maybe_handle(state, request_context, request_body).await? {
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -32,6 +32,68 @@ pub(super) fn admin_provider_ops_is_valid_action_type(action_type: &str) -> bool
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_ops_saved_connector_credentials(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
) -> serde_json::Map<String, serde_json::Value> {
|
||||
admin_provider_ops_decrypted_credentials(
|
||||
state,
|
||||
admin_provider_ops_config_object(provider)
|
||||
.and_then(admin_provider_ops_connector_object)
|
||||
.and_then(|connector| connector.get("credentials")),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn admin_provider_ops_query_balance_response_for_credentials(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
architecture_id: &str,
|
||||
base_url: &str,
|
||||
provider_ops_config: &serde_json::Map<String, serde_json::Value>,
|
||||
connector_config: &serde_json::Map<String, serde_json::Value>,
|
||||
credentials: &serde_json::Map<String, serde_json::Value>,
|
||||
request_config: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
) -> serde_json::Value {
|
||||
let architecture_id = normalize_architecture_id(architecture_id);
|
||||
let Some(architecture) = get_architecture(architecture_id) else {
|
||||
return responses::admin_provider_ops_action_not_supported(
|
||||
"query_balance",
|
||||
ADMIN_PROVIDER_OPS_ACTION_RUST_ONLY_MESSAGE,
|
||||
);
|
||||
};
|
||||
let headers = match build_headers(architecture.architecture_id, connector_config, credentials) {
|
||||
Ok(headers) => headers,
|
||||
Err(message) => {
|
||||
return responses::admin_provider_ops_action_not_configured("query_balance", message);
|
||||
}
|
||||
};
|
||||
let Some(action_config) = resolve_action_config(
|
||||
architecture_id,
|
||||
provider_ops_config,
|
||||
"query_balance",
|
||||
request_config,
|
||||
) else {
|
||||
return responses::admin_provider_ops_action_not_supported(
|
||||
"query_balance",
|
||||
ADMIN_PROVIDER_OPS_ACTION_RUST_ONLY_MESSAGE,
|
||||
);
|
||||
};
|
||||
|
||||
query_balance::admin_provider_ops_run_query_balance_action(
|
||||
state,
|
||||
provider_id,
|
||||
provider,
|
||||
&architecture,
|
||||
base_url,
|
||||
&action_config,
|
||||
&headers,
|
||||
credentials,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn admin_provider_ops_local_action_response(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
|
||||
@@ -111,11 +111,14 @@ pub(super) async fn admin_provider_ops_run_query_balance_action(
|
||||
|
||||
if status != http::StatusCode::OK {
|
||||
let cookie_auth = architecture.query_balance_cookie_auth_errors;
|
||||
let new_api_token_auth = architecture.architecture_id == "new_api";
|
||||
return match status {
|
||||
http::StatusCode::UNAUTHORIZED => admin_provider_ops_action_error(
|
||||
"auth_failed",
|
||||
"query_balance",
|
||||
if cookie_auth {
|
||||
if new_api_token_auth {
|
||||
"访问令牌无效,请使用 New API 个人安全设置里的访问令牌"
|
||||
} else if cookie_auth {
|
||||
"Cookie 已失效,请重新配置"
|
||||
} else {
|
||||
"认证失败"
|
||||
@@ -125,7 +128,9 @@ pub(super) async fn admin_provider_ops_run_query_balance_action(
|
||||
http::StatusCode::FORBIDDEN => admin_provider_ops_action_error(
|
||||
"auth_failed",
|
||||
"query_balance",
|
||||
if cookie_auth {
|
||||
if new_api_token_auth {
|
||||
"访问令牌无效或无权限,请使用 New API 个人安全设置里的访问令牌"
|
||||
} else if cookie_auth {
|
||||
"Cookie 已失效或无权限"
|
||||
} else {
|
||||
"无权限访问"
|
||||
|
||||
@@ -8,7 +8,9 @@ use super::super::responses::{
|
||||
};
|
||||
use super::super::support::admin_provider_ops_json_object_map;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use aether_admin::provider::ops::parse_sub2api_balance_payload;
|
||||
use aether_admin::provider::ops::{
|
||||
parse_sub2api_api_key_usage_payload, parse_sub2api_balance_payload,
|
||||
};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
use serde_json::{json, Value};
|
||||
@@ -24,6 +26,23 @@ pub(super) async fn admin_provider_ops_sub2api_balance_payload(
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
) -> serde_json::Value {
|
||||
let start = std::time::Instant::now();
|
||||
if let Some(api_key) = credentials
|
||||
.get("api_key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return admin_provider_ops_sub2api_api_key_balance_payload(
|
||||
state,
|
||||
provider_id,
|
||||
base_url,
|
||||
action_config,
|
||||
api_key,
|
||||
proxy_snapshot,
|
||||
start,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let (access_token, updated_credentials, _frontend_updated_credentials) =
|
||||
match admin_provider_ops_sub2api_exchange_token(
|
||||
state,
|
||||
@@ -191,6 +210,117 @@ pub(super) async fn admin_provider_ops_sub2api_balance_payload(
|
||||
)
|
||||
}
|
||||
|
||||
async fn admin_provider_ops_sub2api_api_key_balance_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
base_url: &str,
|
||||
action_config: &serde_json::Map<String, serde_json::Value>,
|
||||
api_key: &str,
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
start: std::time::Instant,
|
||||
) -> serde_json::Value {
|
||||
let usage_endpoint = action_config
|
||||
.get("api_key_usage_endpoint")
|
||||
.or_else(|| action_config.get("usage_endpoint"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("/v1/usage");
|
||||
let usage_url = admin_provider_ops_sub2api_request_url(base_url, usage_endpoint);
|
||||
let auth_value = match reqwest::header::HeaderValue::from_str(&format!("Bearer {api_key}")) {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
return admin_provider_ops_action_error(
|
||||
"parse_error",
|
||||
"query_balance",
|
||||
"API Key 格式无效",
|
||||
Some(start.elapsed().as_millis() as u64),
|
||||
);
|
||||
}
|
||||
};
|
||||
let auth_headers = reqwest::header::HeaderMap::from_iter([
|
||||
(reqwest::header::AUTHORIZATION, auth_value),
|
||||
(
|
||||
reqwest::header::ACCEPT,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
]);
|
||||
let request_id = format!("provider-ops-action:sub2api:usage:{provider_id}");
|
||||
let result = admin_provider_ops_execute_json_request(
|
||||
state,
|
||||
&request_id,
|
||||
reqwest::Method::GET,
|
||||
&usage_url,
|
||||
&auth_headers,
|
||||
None,
|
||||
proxy_snapshot,
|
||||
)
|
||||
.await;
|
||||
let response_time_ms = Some(start.elapsed().as_millis() as u64);
|
||||
let (status, response_json) = match result {
|
||||
Ok(result) => result,
|
||||
Err(AdminProviderOpsExecuteJsonError::InvalidJson(message))
|
||||
| Err(AdminProviderOpsExecuteJsonError::Transport(message)) => {
|
||||
return admin_provider_ops_action_error(
|
||||
"network_error",
|
||||
"query_balance",
|
||||
network_error_message(&message),
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
};
|
||||
if matches!(
|
||||
status,
|
||||
http::StatusCode::UNAUTHORIZED | http::StatusCode::FORBIDDEN
|
||||
) {
|
||||
return admin_provider_ops_action_error(
|
||||
"auth_failed",
|
||||
"query_balance",
|
||||
"认证失败,请检查 API Key",
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
if status != http::StatusCode::OK {
|
||||
return admin_provider_ops_action_error(
|
||||
"unknown_error",
|
||||
"query_balance",
|
||||
format!(
|
||||
"HTTP {}: {}",
|
||||
status.as_u16(),
|
||||
status.canonical_reason().unwrap_or("Unknown")
|
||||
),
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
|
||||
let data = match parse_sub2api_api_key_usage_payload(action_config, &response_json) {
|
||||
Ok(payload) => payload,
|
||||
Err(message) => {
|
||||
return admin_provider_ops_action_error(
|
||||
if message.contains("无效") {
|
||||
"auth_failed"
|
||||
} else if message == "响应格式无效" {
|
||||
"parse_error"
|
||||
} else {
|
||||
"unknown_error"
|
||||
},
|
||||
"query_balance",
|
||||
message,
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
admin_provider_ops_action_response(
|
||||
"success",
|
||||
"query_balance",
|
||||
data,
|
||||
None,
|
||||
response_time_ms,
|
||||
86400,
|
||||
)
|
||||
}
|
||||
|
||||
fn network_error_message(error: &str) -> String {
|
||||
let normalized = error.trim();
|
||||
let lower = normalized.to_ascii_lowercase();
|
||||
|
||||
@@ -250,6 +250,9 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
payload: AdminProviderOpsSaveConfigRequest,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
let architecture_id =
|
||||
admin_provider_ops_pure::normalize_architecture_id(payload.architecture_id.as_str())
|
||||
.to_string();
|
||||
let auth_type = payload.connector.auth_type.trim().to_string();
|
||||
if auth_type.is_empty() || !admin_provider_ops_is_supported_auth_type(auth_type.as_str()) {
|
||||
return Err("connector.auth_type 必须是合法的认证类型".to_string());
|
||||
@@ -257,7 +260,7 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
|
||||
|
||||
let merged_credentials = admin_provider_ops_merge_credentials(
|
||||
state,
|
||||
payload.architecture_id.as_str(),
|
||||
architecture_id.as_str(),
|
||||
provider,
|
||||
payload.connector.credentials,
|
||||
);
|
||||
@@ -278,7 +281,7 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
|
||||
.collect::<serde_json::Map<String, serde_json::Value>>();
|
||||
|
||||
Ok(json!({
|
||||
"architecture_id": payload.architecture_id,
|
||||
"architecture_id": architecture_id,
|
||||
"base_url": payload.base_url,
|
||||
"connector": {
|
||||
"auth_type": auth_type,
|
||||
@@ -328,14 +331,16 @@ pub(super) fn build_admin_provider_ops_config_payload(
|
||||
});
|
||||
};
|
||||
let connector = admin_provider_ops_connector_object(provider_ops_config);
|
||||
let architecture_id = provider_ops_config
|
||||
.get("architecture_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(admin_provider_ops_pure::normalize_architecture_id)
|
||||
.unwrap_or("generic_api");
|
||||
|
||||
json!({
|
||||
"provider_id": provider_id,
|
||||
"is_configured": true,
|
||||
"architecture_id": provider_ops_config
|
||||
.get("architecture_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("generic_api"),
|
||||
"architecture_id": architecture_id,
|
||||
"base_url": resolve_admin_provider_ops_base_url(
|
||||
provider,
|
||||
endpoints,
|
||||
|
||||
@@ -6,7 +6,8 @@ use crate::handlers::admin::provider::ops::providers::config::persist_admin_prov
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use aether_admin::provider::ops::{
|
||||
admin_provider_ops_frontend_updated_credentials, admin_provider_ops_verify_failure,
|
||||
parse_verify_payload, ADMIN_PROVIDER_OPS_USER_AGENT,
|
||||
admin_provider_ops_verify_success, admin_provider_ops_verify_user_payload,
|
||||
parse_sub2api_api_key_usage_payload, parse_verify_payload, ADMIN_PROVIDER_OPS_USER_AGENT,
|
||||
};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
@@ -21,6 +22,21 @@ pub(super) async fn admin_provider_ops_local_sub2api_verify_response(
|
||||
credentials: &Map<String, Value>,
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
) -> Value {
|
||||
if let Some(api_key) = credentials
|
||||
.get("api_key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return admin_provider_ops_local_sub2api_api_key_verify_response(
|
||||
state,
|
||||
base_url,
|
||||
api_key,
|
||||
proxy_snapshot,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let (access_token, updated_credentials, frontend_updated_credentials) =
|
||||
match admin_provider_ops_sub2api_exchange_token(
|
||||
state,
|
||||
@@ -93,6 +109,78 @@ pub(super) async fn admin_provider_ops_local_sub2api_verify_response(
|
||||
)
|
||||
}
|
||||
|
||||
async fn admin_provider_ops_local_sub2api_api_key_verify_response(
|
||||
state: &AdminAppState<'_>,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
) -> Value {
|
||||
let usage_url = admin_provider_ops_sub2api_request_url(base_url, "/v1/usage");
|
||||
let auth_value = match reqwest::header::HeaderValue::from_str(&format!("Bearer {api_key}")) {
|
||||
Ok(value) => value,
|
||||
Err(_) => return admin_provider_ops_verify_failure("API Key 格式无效"),
|
||||
};
|
||||
let auth_headers = reqwest::header::HeaderMap::from_iter([
|
||||
(reqwest::header::AUTHORIZATION, auth_value),
|
||||
(
|
||||
reqwest::header::ACCEPT,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
]);
|
||||
let auth_headers =
|
||||
admin_provider_ops_headers_with_transport_controls(&auth_headers, None, true);
|
||||
let (status, response_json) = match admin_provider_ops_execute_json_request(
|
||||
state,
|
||||
"provider-ops-verify:sub2api:api-key",
|
||||
reqwest::Method::GET,
|
||||
&usage_url,
|
||||
&auth_headers,
|
||||
None,
|
||||
proxy_snapshot,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(AdminProviderOpsExecuteJsonError::InvalidJson(message))
|
||||
| Err(AdminProviderOpsExecuteJsonError::Transport(message)) => {
|
||||
return admin_provider_ops_verify_failure(
|
||||
admin_provider_ops_verify_execution_error_message(&message),
|
||||
);
|
||||
}
|
||||
};
|
||||
if matches!(
|
||||
status,
|
||||
http::StatusCode::UNAUTHORIZED | http::StatusCode::FORBIDDEN
|
||||
) {
|
||||
return admin_provider_ops_verify_failure("认证失败:API Key 无效或已过期");
|
||||
}
|
||||
if status != http::StatusCode::OK {
|
||||
return admin_provider_ops_verify_failure(format!("验证失败:HTTP {}", status.as_u16()));
|
||||
}
|
||||
|
||||
let payload = match parse_sub2api_api_key_usage_payload(&Map::new(), &response_json) {
|
||||
Ok(payload) => payload,
|
||||
Err(message) => return admin_provider_ops_verify_failure(message),
|
||||
};
|
||||
let quota = payload.get("total_available").and_then(Value::as_f64);
|
||||
let extra = payload
|
||||
.get("extra")
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
admin_provider_ops_verify_success(
|
||||
admin_provider_ops_verify_user_payload(
|
||||
Some("Sub2API API Key".to_string()),
|
||||
Some("Sub2API API Key".to_string()),
|
||||
None,
|
||||
quota,
|
||||
Some(extra),
|
||||
),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
// 对齐 Python httpx.AsyncClient(base_url=...) 的行为:
|
||||
// 以 "/" 开头的端点始终相对站点根路径解析,而不是简单字符串拼接。
|
||||
pub(in super::super) fn admin_provider_ops_sub2api_request_url(
|
||||
|
||||
@@ -564,6 +564,29 @@ fn provider_query_build_test_request_body_for_api_format(
|
||||
client_api_format.as_str(),
|
||||
payload,
|
||||
);
|
||||
} else if matches!(
|
||||
client_api_format.as_str(),
|
||||
"openai:responses" | "openai:responses:compact"
|
||||
) && !value_has_non_empty_text(object.get("input"))
|
||||
{
|
||||
if let Some(prompt) = object
|
||||
.remove("prompt")
|
||||
.filter(|value| value_has_non_empty_text(Some(value)))
|
||||
{
|
||||
object.insert("input".to_string(), prompt);
|
||||
}
|
||||
}
|
||||
if matches!(
|
||||
client_api_format.as_str(),
|
||||
"openai:responses" | "openai:responses:compact"
|
||||
) && value_has_non_empty_text(object.get("input"))
|
||||
{
|
||||
object.remove("prompt");
|
||||
}
|
||||
if client_api_format == "openai:responses:compact"
|
||||
&& value_has_non_empty_text(object.get("input"))
|
||||
{
|
||||
object.remove("messages");
|
||||
}
|
||||
}
|
||||
return body;
|
||||
|
||||
@@ -30,7 +30,9 @@ pub(super) fn provider_query_standard_test_client_api_format(
|
||||
provider_api_format: &str,
|
||||
) -> &'static str {
|
||||
let normalized_api_format = crate::ai_serving::normalize_api_format_alias(provider_api_format);
|
||||
if crate::ai_serving::is_embedding_api_format(&normalized_api_format) {
|
||||
if normalized_api_format == "openai:responses:compact" {
|
||||
"openai:responses:compact"
|
||||
} else if crate::ai_serving::is_embedding_api_format(&normalized_api_format) {
|
||||
"openai:embedding"
|
||||
} else if crate::ai_serving::is_rerank_api_format(&normalized_api_format) {
|
||||
"openai:rerank"
|
||||
|
||||
@@ -324,6 +324,117 @@ fn provider_query_responses_test_request_body_defaults_to_responses_input() {
|
||||
assert!(body.get("messages").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_compact_test_request_body_defaults_to_responses_input() {
|
||||
let payload = json!({"message": "hello from compact"});
|
||||
|
||||
let client_api_format =
|
||||
provider_query_standard_test_client_api_format("openai:responses:compact");
|
||||
let body = provider_query_build_test_request_body_for_api_format(
|
||||
&payload,
|
||||
"gpt-5.4-mini",
|
||||
"/api/admin/provider-query/test-model",
|
||||
client_api_format,
|
||||
);
|
||||
|
||||
assert_eq!(client_api_format, "openai:responses:compact");
|
||||
assert_eq!(body["model"], json!("gpt-5.4-mini"));
|
||||
assert_eq!(body["input"], json!("hello from compact"));
|
||||
assert!(body.get("messages").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_compact_test_request_body_promotes_prompt_to_input() {
|
||||
let payload = json!({
|
||||
"request_body": {
|
||||
"model": "custom-model",
|
||||
"prompt": "hello from prompt"
|
||||
}
|
||||
});
|
||||
|
||||
let body = provider_query_build_test_request_body_for_api_format(
|
||||
&payload,
|
||||
"fallback-model",
|
||||
"/api/admin/provider-query/test-model",
|
||||
"openai:responses:compact",
|
||||
);
|
||||
|
||||
assert_eq!(body["model"], json!("custom-model"));
|
||||
assert_eq!(body["input"], json!("hello from prompt"));
|
||||
assert!(body.get("prompt").is_none());
|
||||
assert!(provider_query_request_body_is_openai_responses_shape(&body));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_compact_test_request_body_strips_stale_chat_fields() {
|
||||
let payload = json!({
|
||||
"request_body": {
|
||||
"model": "custom-model",
|
||||
"input": "hello from input",
|
||||
"messages": [{ "role": "user", "content": "stale chat body" }],
|
||||
"prompt": "stale prompt"
|
||||
}
|
||||
});
|
||||
|
||||
let body = provider_query_build_test_request_body_for_api_format(
|
||||
&payload,
|
||||
"fallback-model",
|
||||
"/api/admin/provider-query/test-model",
|
||||
"openai:responses:compact",
|
||||
);
|
||||
|
||||
assert_eq!(body["model"], json!("custom-model"));
|
||||
assert_eq!(body["input"], json!("hello from input"));
|
||||
assert!(body.get("messages").is_none());
|
||||
assert!(body.get("prompt").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_compact_provider_body_builds_without_chat_conversion() {
|
||||
let payload = json!({"message": "hello compact provider"});
|
||||
let client_api_format =
|
||||
provider_query_standard_test_client_api_format("openai:responses:compact");
|
||||
let mut request_body = provider_query_build_test_request_body_for_api_format(
|
||||
&payload,
|
||||
"gpt-5.4-mini",
|
||||
"/api/admin/provider-query/test-model",
|
||||
client_api_format,
|
||||
);
|
||||
if let Some(object) = request_body.as_object_mut() {
|
||||
object.insert("stream".to_string(), serde_json::Value::Bool(false));
|
||||
}
|
||||
|
||||
assert!(provider_query_request_body_is_openai_responses_shape(
|
||||
&request_body
|
||||
));
|
||||
|
||||
let mut provider_request_body = crate::ai_serving::build_local_openai_responses_request_body(
|
||||
&request_body,
|
||||
"upstream-gpt",
|
||||
false,
|
||||
)
|
||||
.expect("compact model test body should build from responses shape");
|
||||
crate::ai_serving::apply_openai_responses_compact_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
"openai:responses:compact",
|
||||
);
|
||||
crate::ai_serving::enforce_request_body_stream_field(
|
||||
&mut provider_request_body,
|
||||
"openai:responses:compact",
|
||||
false,
|
||||
true,
|
||||
);
|
||||
|
||||
assert_eq!(provider_request_body["model"], json!("upstream-gpt"));
|
||||
assert_eq!(
|
||||
provider_request_body["input"],
|
||||
json!("hello compact provider")
|
||||
);
|
||||
assert!(provider_request_body.get("messages").is_none());
|
||||
assert!(provider_request_body.get("stream").is_none());
|
||||
assert!(provider_request_body.get("store").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_standard_test_rejects_gemini_success_without_visible_output() {
|
||||
let result = aether_contracts::ExecutionResult {
|
||||
|
||||
@@ -12,6 +12,13 @@ pub(crate) fn admin_provider_id_for_refresh_quota(request_path: &str) -> Option<
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_id_for_key_balance(request_path: &str) -> Option<String> {
|
||||
request_path
|
||||
.strip_prefix("/api/admin/endpoints/providers/")?
|
||||
.strip_suffix("/key-balance")
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub(crate) fn admin_reveal_key_id(request_path: &str) -> Option<String> {
|
||||
request_path
|
||||
.strip_prefix("/api/admin/endpoints/keys/")?
|
||||
|
||||
@@ -15,9 +15,9 @@ pub(crate) use self::crud::{
|
||||
is_admin_providers_root,
|
||||
};
|
||||
pub(crate) use self::endpoint_keys::{
|
||||
admin_clear_oauth_invalid_key_id, admin_export_key_id, admin_provider_id_for_keys,
|
||||
admin_provider_id_for_refresh_quota, admin_reset_cycle_stats_key_id, admin_reveal_key_id,
|
||||
admin_update_key_id,
|
||||
admin_clear_oauth_invalid_key_id, admin_export_key_id, admin_provider_id_for_key_balance,
|
||||
admin_provider_id_for_keys, admin_provider_id_for_refresh_quota,
|
||||
admin_reset_cycle_stats_key_id, admin_reveal_key_id, admin_update_key_id,
|
||||
};
|
||||
pub(crate) use self::oauth::{
|
||||
admin_provider_oauth_batch_import_provider_id, admin_provider_oauth_batch_import_task_path,
|
||||
|
||||
@@ -31,7 +31,7 @@ pub(crate) async fn build_admin_provider_summary_payload(
|
||||
active_global_model_ids_result,
|
||||
) = tokio::join!(
|
||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
|
||||
state.read_provider_quota_snapshot(provider_id),
|
||||
state.list_provider_model_stats(&provider_ids),
|
||||
state.list_active_global_model_ids_by_provider_ids(&provider_ids),
|
||||
@@ -197,7 +197,7 @@ pub(crate) async fn build_admin_providers_summary_payload(
|
||||
} else {
|
||||
let (endpoints_result, keys_result, model_stats_result, active_global_model_refs_result) = tokio::join!(
|
||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
|
||||
state.list_provider_model_stats(&provider_ids),
|
||||
state.list_active_global_model_ids_by_provider_ids(&provider_ids),
|
||||
);
|
||||
|
||||
@@ -7,7 +7,7 @@ use aether_data_contracts::repository::candidates::{
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Map, Value};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
fn json_truthy(value: &serde_json::Value) -> bool {
|
||||
@@ -27,6 +27,80 @@ fn endpoint_timestamp_or_now(value: Option<u64>, now_unix_secs: u64) -> serde_js
|
||||
.unwrap_or(serde_json::Value::Null)
|
||||
}
|
||||
|
||||
fn finite_json_number(value: Option<&Value>) -> Option<f64> {
|
||||
match value {
|
||||
Some(Value::Number(number)) => number.as_f64().filter(|value| value.is_finite()),
|
||||
Some(Value::String(value)) => value
|
||||
.trim()
|
||||
.parse::<f64>()
|
||||
.ok()
|
||||
.filter(|value| value.is_finite()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn finite_json_u64(value: Option<&Value>) -> Option<u64> {
|
||||
finite_json_number(value).and_then(|value| {
|
||||
if value >= 0.0 {
|
||||
Some(value as u64)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn latest_key_balance_summary(keys: &[StoredProviderCatalogKey]) -> Value {
|
||||
let mut selected: Option<(u64, &StoredProviderCatalogKey, &Map<String, Value>)> = None;
|
||||
|
||||
for key in keys {
|
||||
let Some(balance) = key
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|metadata| metadata.get("balance_query"))
|
||||
.and_then(Value::as_object)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(updated_at) = finite_json_u64(balance.get("updated_at")) else {
|
||||
continue;
|
||||
};
|
||||
let has_balance_value = ["total_available", "total_used", "total_granted"]
|
||||
.into_iter()
|
||||
.any(|field| finite_json_number(balance.get(field)).is_some());
|
||||
if !has_balance_value {
|
||||
continue;
|
||||
}
|
||||
|
||||
if selected
|
||||
.as_ref()
|
||||
.is_none_or(|(selected_updated_at, _, _)| updated_at > *selected_updated_at)
|
||||
{
|
||||
selected = Some((updated_at, key, balance));
|
||||
}
|
||||
}
|
||||
|
||||
let Some((updated_at, key, balance)) = selected else {
|
||||
return Value::Null;
|
||||
};
|
||||
|
||||
json!({
|
||||
"key_id": key.id.clone(),
|
||||
"key_name": key.name.clone(),
|
||||
"updated_at": updated_at,
|
||||
"architecture_id": balance.get("architecture_id").cloned().unwrap_or(Value::Null),
|
||||
"status": balance.get("status").cloned().unwrap_or_else(|| json!("success")),
|
||||
"executed_at": balance.get("executed_at").cloned().unwrap_or(Value::Null),
|
||||
"response_time_ms": balance.get("response_time_ms").cloned().unwrap_or(Value::Null),
|
||||
"total_available": balance.get("total_available").cloned().unwrap_or(Value::Null),
|
||||
"total_used": balance.get("total_used").cloned().unwrap_or(Value::Null),
|
||||
"total_granted": balance.get("total_granted").cloned().unwrap_or(Value::Null),
|
||||
"currency": balance.get("currency").cloned().unwrap_or_else(|| json!("USD")),
|
||||
"plan_name": balance.get("plan_name").cloned().unwrap_or(Value::Null),
|
||||
"query_config": balance.get("query_config").cloned().unwrap_or(Value::Null),
|
||||
"extra": balance.get("extra").cloned().unwrap_or(Value::Null),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn build_admin_provider_summary_value(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
@@ -158,6 +232,7 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
.and_then(|quota| quota.quota_expires_at_unix_secs)
|
||||
.or(provider.quota_expires_at_unix_secs)
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
let key_balance_summary = latest_key_balance_summary(keys);
|
||||
|
||||
json!({
|
||||
"id": provider.id.clone(),
|
||||
@@ -196,6 +271,7 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
"endpoint_health_details": endpoint_health_details,
|
||||
"ops_configured": ops_configured,
|
||||
"ops_architecture_id": ops_architecture_id,
|
||||
"key_balance_summary": key_balance_summary,
|
||||
"kiro_simulated_cache_enabled": kiro_simulated_cache_enabled,
|
||||
"created_at": endpoint_timestamp_or_now(provider.created_at_unix_ms, now_unix_secs),
|
||||
"updated_at": endpoint_timestamp_or_now(provider.updated_at_unix_secs, now_unix_secs),
|
||||
|
||||
@@ -177,6 +177,21 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.update_provider_catalog_key(key).await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key_upstream_metadata(
|
||||
&self,
|
||||
key_id: &str,
|
||||
upstream_metadata: Option<&serde_json::Value>,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.app
|
||||
.update_provider_catalog_key_upstream_metadata(
|
||||
key_id,
|
||||
upstream_metadata,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_provider_catalog_key(
|
||||
&self,
|
||||
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
|
||||
|
||||
@@ -344,6 +344,7 @@ impl<'a> AdminAppState<'a> {
|
||||
"allowed_models".to_string(),
|
||||
json!(key.allowed_models.clone()),
|
||||
),
|
||||
("ip_rules".to_string(), json!(key.ip_rules.clone())),
|
||||
("rate_limit".to_string(), json!(key.rate_limit)),
|
||||
("concurrent_limit".to_string(), json!(key.concurrent_limit)),
|
||||
(
|
||||
|
||||
@@ -14,7 +14,7 @@ use crate::handlers::admin::system::shared::configs::apply_admin_system_config_u
|
||||
use crate::handlers::admin::users::{
|
||||
hash_admin_user_api_key, normalize_admin_feature_settings, normalize_admin_list_policy_mode,
|
||||
normalize_admin_rate_limit_policy_mode, normalize_admin_user_api_formats,
|
||||
normalize_admin_user_string_list,
|
||||
normalize_admin_user_ip_rules, normalize_admin_user_string_list,
|
||||
};
|
||||
use crate::handlers::public::normalize_admin_base_url;
|
||||
use crate::GatewayError;
|
||||
@@ -705,6 +705,27 @@ fn normalize_imported_user_api_formats(
|
||||
)?)
|
||||
}
|
||||
|
||||
fn imported_ip_rules_field<'a>(
|
||||
object: &'a Map<String, Value>,
|
||||
) -> (&'static str, Option<&'a Value>) {
|
||||
if let Some(value) = object.get("ip_rules") {
|
||||
("ip_rules", Some(value))
|
||||
} else {
|
||||
("allowed_ips", object.get("allowed_ips"))
|
||||
}
|
||||
}
|
||||
|
||||
fn imported_ip_rules_present(object: &Map<String, Value>) -> bool {
|
||||
object.contains_key("ip_rules") || object.contains_key("allowed_ips")
|
||||
}
|
||||
|
||||
fn normalize_imported_user_ip_rules(
|
||||
object: &Map<String, Value>,
|
||||
) -> Result<Option<Vec<String>>, String> {
|
||||
let (field_name, value) = imported_ip_rules_field(object);
|
||||
normalize_admin_user_ip_rules(imported_string_list_from_value(value, field_name)?)
|
||||
}
|
||||
|
||||
fn build_imported_user_group_record(
|
||||
group: &Map<String, Value>,
|
||||
field_name: &str,
|
||||
@@ -2493,6 +2514,7 @@ impl<'a> AdminAppState<'a> {
|
||||
));
|
||||
let allowed_models =
|
||||
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
|
||||
let ip_rules = invalid_value!(normalize_imported_user_ip_rules(key));
|
||||
let rate_limit =
|
||||
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
|
||||
.unwrap_or(0);
|
||||
@@ -2558,6 +2580,8 @@ impl<'a> AdminAppState<'a> {
|
||||
} else {
|
||||
None
|
||||
},
|
||||
ip_rules: imported_ip_rules_present(key)
|
||||
.then(|| ip_rules.clone()),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
@@ -2626,6 +2650,7 @@ impl<'a> AdminAppState<'a> {
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
ip_rules,
|
||||
rate_limit,
|
||||
concurrent_limit,
|
||||
force_capabilities,
|
||||
@@ -2710,6 +2735,7 @@ impl<'a> AdminAppState<'a> {
|
||||
));
|
||||
let allowed_models =
|
||||
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
|
||||
let ip_rules = invalid_value!(normalize_imported_user_ip_rules(key));
|
||||
let rate_limit =
|
||||
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
|
||||
.unwrap_or(0);
|
||||
@@ -2781,6 +2807,8 @@ impl<'a> AdminAppState<'a> {
|
||||
allowed_providers: Some(allowed_providers.clone()),
|
||||
allowed_api_formats: Some(allowed_api_formats.clone()),
|
||||
allowed_models: Some(allowed_models.clone()),
|
||||
ip_rules: imported_ip_rules_present(key)
|
||||
.then(|| ip_rules.clone()),
|
||||
expires_at_present: false,
|
||||
expires_at_unix_secs: None,
|
||||
auto_delete_on_expiry_present: false,
|
||||
@@ -2841,6 +2869,7 @@ impl<'a> AdminAppState<'a> {
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
ip_rules,
|
||||
rate_limit: Some(rate_limit),
|
||||
concurrent_limit,
|
||||
force_capabilities,
|
||||
|
||||
@@ -42,6 +42,13 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.list_export_users_page(query).await
|
||||
}
|
||||
|
||||
pub(crate) async fn count_export_users(
|
||||
&self,
|
||||
query: &aether_data::repository::users::UserExportListQuery,
|
||||
) -> Result<u64, GatewayError> {
|
||||
self.app.count_export_users(query).await
|
||||
}
|
||||
|
||||
pub(crate) async fn find_export_user_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
|
||||
@@ -10,7 +10,7 @@ use crate::handlers::admin::system::shared::paths::{
|
||||
admin_management_token_status_id_from_path, is_admin_management_tokens_root,
|
||||
};
|
||||
use crate::handlers::internal::build_management_token_payload;
|
||||
use crate::handlers::shared::generate_gateway_secret_plaintext;
|
||||
use crate::handlers::shared::{generate_gateway_secret_plaintext, parse_json_ip_rules};
|
||||
use crate::{GatewayError, LocalMutationOutcome};
|
||||
use aether_data::repository::management_tokens::{
|
||||
CreateManagementTokenRecord, ManagementTokenListQuery, RegenerateManagementTokenSecret,
|
||||
@@ -97,59 +97,10 @@ fn admin_management_token_prefix(value: &str) -> Option<String> {
|
||||
.then(|| value[..value.len().min(ADMIN_MANAGEMENT_TOKEN_DISPLAY_PREFIX_LEN)].to_string())
|
||||
}
|
||||
|
||||
fn admin_validate_ip_or_cidr(value: &str) -> bool {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if value.parse::<std::net::IpAddr>().is_ok() {
|
||||
return true;
|
||||
}
|
||||
let Some((host, prefix)) = value.split_once('/') else {
|
||||
return false;
|
||||
};
|
||||
let Ok(ip) = host.trim().parse::<std::net::IpAddr>() else {
|
||||
return false;
|
||||
};
|
||||
let Ok(prefix) = prefix.trim().parse::<u8>() else {
|
||||
return false;
|
||||
};
|
||||
match ip {
|
||||
std::net::IpAddr::V4(_) => prefix <= 32,
|
||||
std::net::IpAddr::V6(_) => prefix <= 128,
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_parse_management_token_allowed_ips(
|
||||
value: Option<&serde_json::Value>,
|
||||
) -> Result<Option<serde_json::Value>, String> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
match value {
|
||||
serde_json::Value::Null => Ok(None),
|
||||
serde_json::Value::Array(items) => {
|
||||
if items.is_empty() {
|
||||
return Err("IP 白名单不能为空列表,如需取消限制请不提供此字段".to_string());
|
||||
}
|
||||
let mut normalized = Vec::with_capacity(items.len());
|
||||
for (index, item) in items.iter().enumerate() {
|
||||
let Some(raw) = item.as_str() else {
|
||||
return Err("IP 白名单必须是字符串数组".to_string());
|
||||
};
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err(format!("IP 白名单第 {} 项为空", index + 1));
|
||||
}
|
||||
if !admin_validate_ip_or_cidr(trimmed) {
|
||||
return Err(format!("无效的 IP 地址或 CIDR: {raw}"));
|
||||
}
|
||||
normalized.push(trimmed.to_string());
|
||||
}
|
||||
Ok(Some(json!(normalized)))
|
||||
}
|
||||
_ => Err("IP 白名单必须是字符串数组".to_string()),
|
||||
}
|
||||
parse_json_ip_rules(value)
|
||||
}
|
||||
|
||||
fn admin_parse_management_token_expires_at(
|
||||
|
||||
@@ -46,7 +46,7 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
|
||||
AdminModuleDefinition {
|
||||
name: "management_tokens",
|
||||
display_name: "访问令牌",
|
||||
description: "管理 API 访问令牌,支持细粒度权限控制和 IP 白名单",
|
||||
description: "管理 API 访问令牌,支持细粒度权限控制和 IP 限制",
|
||||
category: "security",
|
||||
env_key: "MANAGEMENT_TOKENS_AVAILABLE",
|
||||
default_available: true,
|
||||
|
||||
@@ -44,6 +44,7 @@ pub(super) fn build_admin_user_api_key_detail_payload(
|
||||
"total_cost_usd": record.total_cost_usd,
|
||||
"rate_limit": record.rate_limit,
|
||||
"concurrent_limit": record.concurrent_limit,
|
||||
"ip_rules": record.ip_rules,
|
||||
"feature_settings": record.feature_settings,
|
||||
"expires_at": format_optional_unix_secs_iso8601(record.expires_at_unix_secs),
|
||||
"last_used_at": format_optional_unix_secs_iso8601(record.last_used_at_unix_secs),
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use super::super::super::{
|
||||
build_admin_users_bad_request_response, build_admin_users_data_unavailable_response,
|
||||
build_admin_users_read_only_response, normalize_admin_feature_settings,
|
||||
AdminCreateUserApiKeyRequest,
|
||||
normalize_admin_user_ip_rules, AdminCreateUserApiKeyRequest,
|
||||
};
|
||||
use super::super::helpers::{
|
||||
attach_audit_response, default_admin_user_api_key_name, format_optional_unix_secs_iso8601,
|
||||
@@ -71,7 +71,7 @@ pub(crate) async fn build_admin_create_user_api_key_response(
|
||||
{
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": "当前仅支持 name、rate_limit、concurrent_limit、allowed_providers 字段" })),
|
||||
Json(json!({ "detail": "当前仅支持 name、rate_limit、concurrent_limit、allowed_providers、ip_rules 字段" })),
|
||||
)
|
||||
.into_response());
|
||||
}
|
||||
@@ -107,6 +107,16 @@ pub(crate) async fn build_admin_create_user_api_key_response(
|
||||
.into_response());
|
||||
}
|
||||
};
|
||||
let ip_rules = match normalize_admin_user_ip_rules(payload.ip_rules) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response());
|
||||
}
|
||||
};
|
||||
let rate_limit = payload.rate_limit.unwrap_or(0);
|
||||
if rate_limit < 0 {
|
||||
return Ok((
|
||||
@@ -146,6 +156,7 @@ pub(crate) async fn build_admin_create_user_api_key_response(
|
||||
allowed_providers: None,
|
||||
allowed_api_formats: None,
|
||||
allowed_models: None,
|
||||
ip_rules,
|
||||
rate_limit,
|
||||
concurrent_limit,
|
||||
force_capabilities: None,
|
||||
@@ -196,6 +207,7 @@ pub(crate) async fn build_admin_create_user_api_key_response(
|
||||
"key_display": masked_user_api_key_display(state, created.key_encrypted.as_deref()),
|
||||
"rate_limit": created.rate_limit,
|
||||
"concurrent_limit": created.concurrent_limit,
|
||||
"ip_rules": created.ip_rules,
|
||||
"expires_at": format_optional_unix_secs_iso8601(created.expires_at_unix_secs),
|
||||
"last_used_at": format_optional_unix_secs_iso8601(created.last_used_at_unix_secs),
|
||||
"created_at": format_optional_unix_secs_iso8601(created.created_at_unix_secs),
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use super::super::super::{
|
||||
build_admin_users_bad_request_response, build_admin_users_read_only_response,
|
||||
normalize_admin_feature_settings, AdminUpdateUserApiKeyRequest,
|
||||
normalize_admin_feature_settings, normalize_admin_user_ip_rules, AdminUpdateUserApiKeyRequest,
|
||||
};
|
||||
use super::super::helpers::{
|
||||
attach_audit_response, build_admin_user_api_key_detail_payload,
|
||||
@@ -94,6 +94,19 @@ pub(crate) async fn build_admin_update_user_api_key_response(
|
||||
.into_response());
|
||||
}
|
||||
};
|
||||
let ip_rules = match payload.ip_rules {
|
||||
Some(value) => match normalize_admin_user_ip_rules(value) {
|
||||
Ok(value) => Some(value),
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response());
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
|
||||
let Some(updated) = state
|
||||
.update_user_api_key_basic(aether_data::repository::auth::UpdateUserApiKeyBasicRecord {
|
||||
@@ -102,6 +115,7 @@ pub(crate) async fn build_admin_update_user_api_key_response(
|
||||
name,
|
||||
rate_limit: payload.rate_limit,
|
||||
concurrent_limit,
|
||||
ip_rules,
|
||||
})
|
||||
.await?
|
||||
else {
|
||||
|
||||
@@ -37,16 +37,20 @@ pub(in super::super) async fn build_admin_list_users_response(
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
let paged_rows = state
|
||||
.list_export_users_page(&aether_data::repository::users::UserExportListQuery {
|
||||
skip,
|
||||
limit,
|
||||
role: role.clone(),
|
||||
is_active,
|
||||
search,
|
||||
group_id,
|
||||
})
|
||||
.await?;
|
||||
let query = aether_data::repository::users::UserExportListQuery {
|
||||
skip,
|
||||
limit,
|
||||
role: role.clone(),
|
||||
is_active,
|
||||
search,
|
||||
group_id,
|
||||
};
|
||||
let (paged_rows_result, total_result) = tokio::join!(
|
||||
state.list_export_users_page(&query),
|
||||
state.count_export_users(&query),
|
||||
);
|
||||
let paged_rows = paged_rows_result?;
|
||||
let total = total_result?;
|
||||
let user_ids = paged_rows
|
||||
.iter()
|
||||
.map(|row| row.id.clone())
|
||||
@@ -116,7 +120,15 @@ pub(in super::super) async fn build_admin_list_users_response(
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Json(payload).into_response())
|
||||
let has_more = (skip as u64).saturating_add(payload.len() as u64) < total;
|
||||
Ok(Json(json!({
|
||||
"items": payload,
|
||||
"total": total,
|
||||
"skip": skip,
|
||||
"limit": limit,
|
||||
"has_more": has_more,
|
||||
}))
|
||||
.into_response())
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_get_user_response(
|
||||
|
||||
@@ -56,7 +56,8 @@ use self::shared::{
|
||||
};
|
||||
pub(crate) use self::shared::{
|
||||
normalize_admin_list_policy_mode, normalize_admin_rate_limit_policy_mode,
|
||||
normalize_admin_user_api_formats, normalize_admin_user_string_list,
|
||||
normalize_admin_user_api_formats, normalize_admin_user_ip_rules,
|
||||
normalize_admin_user_string_list,
|
||||
};
|
||||
pub(crate) use crate::handlers::shared::normalize_feature_settings as normalize_admin_feature_settings;
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use super::ADMIN_USERS_DATA_UNAVAILABLE_DETAIL;
|
||||
use crate::handlers::admin::shared::AdminTypedObjectPatch;
|
||||
use crate::handlers::shared::{deserialize_optional_string_list_patch, normalize_ip_rules};
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
@@ -19,6 +20,8 @@ pub(super) struct AdminCreateUserApiKeyRequest {
|
||||
pub(super) allowed_api_formats: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_models: Option<Vec<String>>,
|
||||
#[serde(default, alias = "allowed_ips")]
|
||||
pub(super) ip_rules: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) rate_limit: Option<i32>,
|
||||
#[serde(default)]
|
||||
@@ -49,6 +52,12 @@ pub(super) struct AdminUpdateUserApiKeyRequest {
|
||||
pub(super) concurrent_limit: Option<i32>,
|
||||
#[serde(default)]
|
||||
pub(super) feature_settings: Option<Option<Value>>,
|
||||
#[serde(
|
||||
default,
|
||||
alias = "allowed_ips",
|
||||
deserialize_with = "deserialize_optional_string_list_patch"
|
||||
)]
|
||||
pub(super) ip_rules: Option<Option<Vec<String>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
@@ -280,6 +289,12 @@ pub(crate) fn normalize_admin_user_api_formats(
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_admin_user_ip_rules(
|
||||
value: Option<Vec<String>>,
|
||||
) -> Result<Option<Vec<String>>, String> {
|
||||
normalize_ip_rules(value)
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_admin_list_policy_mode(value: &str) -> Result<String, String> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"inherit" | "unrestricted" | "specific" | "deny_all" => {
|
||||
@@ -338,7 +353,8 @@ pub(super) fn format_optional_datetime_iso8601(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::normalize_admin_user_api_formats;
|
||||
use super::{normalize_admin_user_api_formats, AdminUpdateUserApiKeyRequest};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn admin_user_api_formats_accept_current_canonical_signatures() {
|
||||
@@ -375,4 +391,31 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_update_api_key_distinguishes_missing_null_and_present_ip_rules() {
|
||||
let missing = serde_json::from_value::<AdminUpdateUserApiKeyRequest>(json!({
|
||||
"name": "unchanged-ip-rules",
|
||||
}))
|
||||
.expect("missing ip_rules should deserialize");
|
||||
assert_eq!(missing.ip_rules, None);
|
||||
|
||||
let cleared = serde_json::from_value::<AdminUpdateUserApiKeyRequest>(json!({
|
||||
"ip_rules": null,
|
||||
}))
|
||||
.expect("null ip_rules should deserialize");
|
||||
assert_eq!(cleared.ip_rules, Some(None));
|
||||
|
||||
let updated = serde_json::from_value::<AdminUpdateUserApiKeyRequest>(json!({
|
||||
"ip_rules": ["203.0.113.10", "10.0.0.0/24"],
|
||||
}))
|
||||
.expect("present ip_rules should deserialize");
|
||||
assert_eq!(
|
||||
updated.ip_rules,
|
||||
Some(Some(vec![
|
||||
"203.0.113.10".to_string(),
|
||||
"10.0.0.0/24".to_string(),
|
||||
])),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,8 +46,8 @@ use crate::frontdoor_loop_guard::{
|
||||
frontdoor_self_loop_public_ai_path, request_has_execution_runtime_loop_guard,
|
||||
};
|
||||
use crate::handlers::shared::{
|
||||
build_admin_proxy_auth_required_response, build_unhandled_admin_proxy_response,
|
||||
local_proxy_route_requires_buffered_body, request_enables_control_execute,
|
||||
build_admin_proxy_auth_required_response, build_unhandled_admin_proxy_response, ip_rules_allow,
|
||||
json_ip_rules_allow, local_proxy_route_requires_buffered_body, request_enables_control_execute,
|
||||
should_strip_forwarded_provider_credential_header, should_strip_forwarded_trusted_admin_header,
|
||||
};
|
||||
use crate::headers::{
|
||||
@@ -197,57 +197,11 @@ fn hash_management_token(value: &str) -> String {
|
||||
}
|
||||
|
||||
fn remote_ip_allowed(allowed_ips: Option<&serde_json::Value>, remote_ip: std::net::IpAddr) -> bool {
|
||||
let Some(allowed_ips) = allowed_ips else {
|
||||
return true;
|
||||
};
|
||||
if allowed_ips.is_null() {
|
||||
return true;
|
||||
}
|
||||
let Some(items) = allowed_ips.as_array() else {
|
||||
return false;
|
||||
};
|
||||
if items.is_empty() {
|
||||
return false;
|
||||
}
|
||||
items
|
||||
.iter()
|
||||
.filter_map(serde_json::Value::as_str)
|
||||
.any(|value| ip_or_cidr_matches(value, remote_ip))
|
||||
json_ip_rules_allow(allowed_ips, remote_ip)
|
||||
}
|
||||
|
||||
fn ip_or_cidr_matches(pattern: &str, remote_ip: std::net::IpAddr) -> bool {
|
||||
let pattern = pattern.trim();
|
||||
if pattern.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if let Ok(ip) = pattern.parse::<std::net::IpAddr>() {
|
||||
return ip == remote_ip;
|
||||
}
|
||||
let Some((network, prefix)) = pattern.split_once('/') else {
|
||||
return false;
|
||||
};
|
||||
let Ok(prefix) = prefix.trim().parse::<u8>() else {
|
||||
return false;
|
||||
};
|
||||
match (network.trim().parse::<std::net::IpAddr>(), remote_ip) {
|
||||
(Ok(std::net::IpAddr::V4(network)), std::net::IpAddr::V4(remote)) if prefix <= 32 => {
|
||||
let mask = if prefix == 0 {
|
||||
0
|
||||
} else {
|
||||
u32::MAX << (32 - prefix)
|
||||
};
|
||||
(u32::from(network) & mask) == (u32::from(remote) & mask)
|
||||
}
|
||||
(Ok(std::net::IpAddr::V6(network)), std::net::IpAddr::V6(remote)) if prefix <= 128 => {
|
||||
let mask = if prefix == 0 {
|
||||
0
|
||||
} else {
|
||||
u128::MAX << (128 - prefix)
|
||||
};
|
||||
(u128::from(network) & mask) == (u128::from(remote) & mask)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
fn api_key_remote_ip_allowed(ip_rules: Option<&[String]>, remote_ip: std::net::IpAddr) -> bool {
|
||||
ip_rules_allow(ip_rules, remote_ip)
|
||||
}
|
||||
|
||||
async fn maybe_promote_management_token_admin_principal(
|
||||
@@ -986,6 +940,31 @@ pub(crate) async fn proxy_request(
|
||||
&mut request_context,
|
||||
)
|
||||
.await?;
|
||||
if let Some(auth_context) = request_context
|
||||
.control_decision
|
||||
.as_ref()
|
||||
.and_then(|decision| decision.auth_context.as_ref())
|
||||
{
|
||||
if !api_key_remote_ip_allowed(auth_context.ip_rules.as_deref(), remote_addr.ip()) {
|
||||
let rejection = crate::control::GatewayLocalAuthRejection::IpNotAllowed {
|
||||
remote_ip: remote_addr.ip().to_string(),
|
||||
};
|
||||
let response = build_local_auth_rejection_response(
|
||||
&trace_id,
|
||||
request_context.control_decision.as_ref(),
|
||||
&rejection,
|
||||
)?;
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
EXECUTION_PATH_LOCAL_AUTH_DENIED,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
}
|
||||
let request_context_ms = request_context_started_at.elapsed().as_millis() as u64;
|
||||
if request_context
|
||||
.control_decision
|
||||
@@ -1999,14 +1978,43 @@ fn local_execution_runtime_miss_route_detail(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
|
||||
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
|
||||
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic,
|
||||
api_key_remote_ip_allowed, diagnostic_is_auth_api_key_concurrency_limited,
|
||||
local_execution_runtime_miss_detail, restore_redacted_stream_execution_response,
|
||||
restore_redacted_sync_execution_response, GatewayControlDecision,
|
||||
LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::http::{header, Response};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn api_key_remote_ip_allows_unrestricted_keys() {
|
||||
let remote_ip = "203.0.113.10".parse().expect("valid ip");
|
||||
assert!(api_key_remote_ip_allowed(None, remote_ip));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_key_remote_ip_applies_ip_rules() {
|
||||
let ip_rules = vec![
|
||||
"198.51.100.1".to_string(),
|
||||
"203.0.113.*".to_string(),
|
||||
"!203.0.113.13".to_string(),
|
||||
];
|
||||
|
||||
assert!(api_key_remote_ip_allowed(
|
||||
Some(&ip_rules),
|
||||
"198.51.100.1".parse().expect("valid ip"),
|
||||
));
|
||||
assert!(api_key_remote_ip_allowed(
|
||||
Some(&ip_rules),
|
||||
"203.0.113.42".parse().expect("valid ip"),
|
||||
));
|
||||
assert!(!api_key_remote_ip_allowed(
|
||||
Some(&ip_rules),
|
||||
"203.0.113.13".parse().expect("valid ip"),
|
||||
));
|
||||
}
|
||||
|
||||
fn redaction_slot_for_email() -> (crate::privacy::RedactionSessionSlot, String) {
|
||||
let masked = crate::privacy::mask_chat_request_json(
|
||||
br#"{"messages":[{"role":"user","content":"Email alice@example.com"}]}"#,
|
||||
|
||||
@@ -11,7 +11,8 @@ use serde_json::json;
|
||||
|
||||
use crate::handlers::shared::{
|
||||
api_key_placeholder_display, deserialize_optional_json_patch,
|
||||
generate_gateway_api_key_plaintext, masked_gateway_api_key_display, normalize_feature_settings,
|
||||
deserialize_optional_string_list_patch, generate_gateway_api_key_plaintext,
|
||||
masked_gateway_api_key_display, normalize_feature_settings, normalize_ip_rules,
|
||||
normalize_optional_api_key_concurrent_limit,
|
||||
};
|
||||
|
||||
@@ -34,6 +35,8 @@ struct UsersMeCreateApiKeyRequest {
|
||||
concurrent_limit: Option<i32>,
|
||||
#[serde(default)]
|
||||
feature_settings: Option<serde_json::Value>,
|
||||
#[serde(default, alias = "allowed_ips")]
|
||||
ip_rules: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -46,6 +49,12 @@ struct UsersMeUpdateApiKeyRequest {
|
||||
concurrent_limit: Option<i32>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_json_patch")]
|
||||
feature_settings: Option<Option<serde_json::Value>>,
|
||||
#[serde(
|
||||
default,
|
||||
alias = "allowed_ips",
|
||||
deserialize_with = "deserialize_optional_string_list_patch"
|
||||
)]
|
||||
ip_rules: Option<Option<Vec<String>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -160,6 +169,7 @@ fn build_users_me_api_key_list_payload(
|
||||
"rate_limit": record.rate_limit,
|
||||
"concurrent_limit": record.concurrent_limit,
|
||||
"allowed_providers": record.allowed_providers,
|
||||
"ip_rules": record.ip_rules,
|
||||
"force_capabilities": record.force_capabilities,
|
||||
"feature_settings": record.feature_settings,
|
||||
})
|
||||
@@ -177,6 +187,7 @@ fn build_users_me_api_key_detail_payload(
|
||||
"is_active": record.is_active,
|
||||
"is_locked": is_locked,
|
||||
"allowed_providers": record.allowed_providers,
|
||||
"ip_rules": record.ip_rules,
|
||||
"force_capabilities": record.force_capabilities,
|
||||
"feature_settings": record.feature_settings,
|
||||
"rate_limit": record.rate_limit,
|
||||
@@ -199,6 +210,10 @@ fn generate_users_me_api_key_plaintext() -> String {
|
||||
generate_gateway_api_key_plaintext()
|
||||
}
|
||||
|
||||
fn normalize_users_me_ip_rules(values: Option<Vec<String>>) -> Result<Option<Vec<String>>, String> {
|
||||
normalize_ip_rules(values)
|
||||
}
|
||||
|
||||
fn hash_users_me_api_key(value: &str) -> String {
|
||||
use sha2::Digest;
|
||||
|
||||
@@ -542,6 +557,12 @@ pub(super) async fn handle_users_me_api_key_create(
|
||||
return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false);
|
||||
}
|
||||
};
|
||||
let ip_rules = match normalize_users_me_ip_rules(payload.ip_rules) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false);
|
||||
}
|
||||
};
|
||||
|
||||
let plaintext_key = generate_users_me_api_key_plaintext();
|
||||
let Some(key_encrypted) = encrypt_catalog_secret_with_fallbacks(state, &plaintext_key) else {
|
||||
@@ -560,6 +581,7 @@ pub(super) async fn handle_users_me_api_key_create(
|
||||
allowed_providers: None,
|
||||
allowed_api_formats: None,
|
||||
allowed_models: None,
|
||||
ip_rules,
|
||||
rate_limit,
|
||||
concurrent_limit,
|
||||
force_capabilities: None,
|
||||
@@ -614,6 +636,7 @@ pub(super) async fn handle_users_me_api_key_create(
|
||||
"is_locked": false,
|
||||
"rate_limit": created.rate_limit,
|
||||
"concurrent_limit": created.concurrent_limit,
|
||||
"ip_rules": created.ip_rules,
|
||||
"feature_settings": created.feature_settings,
|
||||
"last_used_at": format_users_me_optional_unix_secs_iso8601(created.last_used_at_unix_secs),
|
||||
"created_at": format_users_me_optional_unix_secs_iso8601(created.created_at_unix_secs),
|
||||
@@ -695,6 +718,15 @@ pub(super) async fn handle_users_me_api_key_update(
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
let ip_rules = match payload.ip_rules {
|
||||
Some(value) => match normalize_users_me_ip_rules(value) {
|
||||
Ok(value) => Some(value),
|
||||
Err(detail) => {
|
||||
return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false);
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
|
||||
let Some(updated) = (match state
|
||||
.update_user_api_key_basic(aether_data::repository::auth::UpdateUserApiKeyBasicRecord {
|
||||
@@ -703,6 +735,7 @@ pub(super) async fn handle_users_me_api_key_update(
|
||||
name,
|
||||
rate_limit,
|
||||
concurrent_limit,
|
||||
ip_rules,
|
||||
})
|
||||
.await
|
||||
{
|
||||
@@ -1057,3 +1090,58 @@ pub(super) async fn handle_users_me_api_key_capabilities_put(
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{normalize_users_me_ip_rules, UsersMeUpdateApiKeyRequest};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn normalize_ip_rules_trims_ip_and_cidr_values() {
|
||||
let values = normalize_users_me_ip_rules(Some(vec![
|
||||
" 203.0.113.10 ".to_string(),
|
||||
"10.0.0.0/24".to_string(),
|
||||
]))
|
||||
.expect("valid IP rules should normalize");
|
||||
|
||||
assert_eq!(
|
||||
values,
|
||||
Some(vec!["203.0.113.10".to_string(), "10.0.0.0/24".to_string()]),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_ip_rules_rejects_invalid_cidr() {
|
||||
let err = normalize_users_me_ip_rules(Some(vec!["10.0.0.0/99".to_string()]))
|
||||
.expect_err("invalid cidr should fail");
|
||||
|
||||
assert_eq!(err, "无效的 IP 限制规则: 10.0.0.0/99(第 1 项)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_payload_distinguishes_missing_null_and_present_ip_rules() {
|
||||
let missing = serde_json::from_value::<UsersMeUpdateApiKeyRequest>(json!({
|
||||
"name": "unchanged-ip-rules",
|
||||
}))
|
||||
.expect("missing ip_rules should deserialize");
|
||||
assert_eq!(missing.ip_rules, None);
|
||||
|
||||
let cleared = serde_json::from_value::<UsersMeUpdateApiKeyRequest>(json!({
|
||||
"ip_rules": null,
|
||||
}))
|
||||
.expect("null ip_rules should deserialize");
|
||||
assert_eq!(cleared.ip_rules, Some(None));
|
||||
|
||||
let updated = serde_json::from_value::<UsersMeUpdateApiKeyRequest>(json!({
|
||||
"ip_rules": ["203.0.113.10", "10.0.0.0/24"],
|
||||
}))
|
||||
.expect("present ip_rules should deserialize");
|
||||
assert_eq!(
|
||||
updated.ip_rules,
|
||||
Some(Some(vec![
|
||||
"203.0.113.10".to_string(),
|
||||
"10.0.0.0/24".to_string(),
|
||||
])),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,7 +20,7 @@ use super::{
|
||||
GatewayPublicRequestContext,
|
||||
};
|
||||
use crate::control::normalize_assignable_management_token_permissions;
|
||||
use crate::handlers::shared::generate_gateway_secret_plaintext;
|
||||
use crate::handlers::shared::{generate_gateway_secret_plaintext, parse_json_ip_rules};
|
||||
use crate::LocalMutationOutcome;
|
||||
|
||||
const USERS_ME_MANAGEMENT_TOKEN_PREFIX: &str = "ae";
|
||||
@@ -177,59 +177,10 @@ fn users_me_management_token_skip(query: Option<&str>) -> usize {
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn users_me_validate_ip_or_cidr(value: &str) -> bool {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if value.parse::<std::net::IpAddr>().is_ok() {
|
||||
return true;
|
||||
}
|
||||
let Some((host, prefix)) = value.split_once('/') else {
|
||||
return false;
|
||||
};
|
||||
let Ok(ip) = host.trim().parse::<std::net::IpAddr>() else {
|
||||
return false;
|
||||
};
|
||||
let Ok(prefix) = prefix.trim().parse::<u8>() else {
|
||||
return false;
|
||||
};
|
||||
match ip {
|
||||
std::net::IpAddr::V4(_) => prefix <= 32,
|
||||
std::net::IpAddr::V6(_) => prefix <= 128,
|
||||
}
|
||||
}
|
||||
|
||||
fn users_me_parse_management_token_allowed_ips(
|
||||
value: Option<&serde_json::Value>,
|
||||
) -> Result<Option<serde_json::Value>, String> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
match value {
|
||||
serde_json::Value::Null => Ok(None),
|
||||
serde_json::Value::Array(items) => {
|
||||
if items.is_empty() {
|
||||
return Err("IP 白名单不能为空列表,如需取消限制请不提供此字段".to_string());
|
||||
}
|
||||
let mut normalized = Vec::with_capacity(items.len());
|
||||
for (index, item) in items.iter().enumerate() {
|
||||
let Some(raw) = item.as_str() else {
|
||||
return Err("IP 白名单必须是字符串数组".to_string());
|
||||
};
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err(format!("IP 白名单第 {} 项为空", index + 1));
|
||||
}
|
||||
if !users_me_validate_ip_or_cidr(trimmed) {
|
||||
return Err(format!("无效的 IP 地址或 CIDR: {raw}"));
|
||||
}
|
||||
normalized.push(trimmed.to_string());
|
||||
}
|
||||
Ok(Some(json!(normalized)))
|
||||
}
|
||||
_ => Err("IP 白名单必须是字符串数组".to_string()),
|
||||
}
|
||||
parse_json_ip_rules(value)
|
||||
}
|
||||
|
||||
fn users_me_parse_management_token_expires_at(
|
||||
|
||||
@@ -21,6 +21,7 @@ const OAUTH_ACCOUNT_BLOCK_PREFIX: &str = "[ACCOUNT_BLOCK] ";
|
||||
const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] ";
|
||||
const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] ";
|
||||
const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] ";
|
||||
const BALANCE_QUERY_SECRET_CIPHERTEXT_KEY: &str = "secret_ciphertext";
|
||||
|
||||
pub(crate) fn provider_catalog_key_supports_format(
|
||||
key: &StoredProviderCatalogKey,
|
||||
@@ -173,6 +174,31 @@ pub(crate) fn parse_catalog_auth_config_json(
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn sanitized_admin_upstream_metadata(upstream_metadata: Option<&Value>) -> Value {
|
||||
let Some(mut metadata) = upstream_metadata.cloned() else {
|
||||
return Value::Null;
|
||||
};
|
||||
let Some(balance_query) = metadata
|
||||
.as_object_mut()
|
||||
.and_then(|metadata| metadata.get_mut("balance_query"))
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
return metadata;
|
||||
};
|
||||
let has_saved_secret = balance_query
|
||||
.remove(BALANCE_QUERY_SECRET_CIPHERTEXT_KEY)
|
||||
.is_some();
|
||||
if has_saved_secret {
|
||||
let query_config = balance_query
|
||||
.entry("query_config".to_string())
|
||||
.or_insert_with(|| Value::Object(Map::new()));
|
||||
if let Some(query_config) = query_config.as_object_mut() {
|
||||
query_config.insert("has_saved_secret".to_string(), Value::Bool(true));
|
||||
}
|
||||
}
|
||||
metadata
|
||||
}
|
||||
|
||||
pub(crate) fn default_provider_key_status_snapshot() -> serde_json::Value {
|
||||
json!({
|
||||
"oauth": {
|
||||
@@ -1937,7 +1963,7 @@ pub(crate) fn build_admin_provider_key_response(
|
||||
);
|
||||
payload.insert(
|
||||
"upstream_metadata".to_string(),
|
||||
json!(key.upstream_metadata),
|
||||
sanitized_admin_upstream_metadata(key.upstream_metadata.as_ref()),
|
||||
);
|
||||
payload.insert("proxy".to_string(), json!(key.proxy));
|
||||
payload.insert("fingerprint".to_string(), json!(key.fingerprint));
|
||||
|
||||
@@ -36,8 +36,9 @@ pub(crate) use self::email_templates::{
|
||||
};
|
||||
pub(crate) use self::external_models::OFFICIAL_EXTERNAL_MODEL_PROVIDERS;
|
||||
pub(crate) use self::normalize::{
|
||||
deserialize_optional_json_patch, normalize_feature_settings, normalize_json_array,
|
||||
normalize_json_object, normalize_string_list,
|
||||
deserialize_optional_json_patch, deserialize_optional_string_list_patch, ip_rules_allow,
|
||||
json_ip_rules_allow, normalize_feature_settings, normalize_ip_rules, normalize_json_array,
|
||||
normalize_json_object, normalize_string_list, parse_json_ip_rules,
|
||||
};
|
||||
pub(crate) use self::payloads::{
|
||||
InternalGatewayAuthContextRequest, InternalGatewayExecuteRequest,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::collections::BTreeSet;
|
||||
use std::net::IpAddr;
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
@@ -63,6 +64,215 @@ pub(crate) fn normalize_feature_settings(value: Option<Value>) -> Result<Option<
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_ip_rules(
|
||||
values: Option<Vec<String>>,
|
||||
) -> Result<Option<Vec<String>>, String> {
|
||||
let Some(values) = values else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut normalized = Vec::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
for (index, raw) in values.into_iter().enumerate() {
|
||||
let rule = normalize_ip_rule(raw.trim())
|
||||
.map_err(|detail| format!("{detail}(第 {} 项)", index + 1))?;
|
||||
if seen.insert(rule.clone()) {
|
||||
normalized.push(rule);
|
||||
}
|
||||
}
|
||||
Ok((!normalized.is_empty()).then_some(normalized))
|
||||
}
|
||||
|
||||
pub(crate) fn parse_json_ip_rules(value: Option<&Value>) -> Result<Option<Value>, String> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
match value {
|
||||
Value::Null => Ok(None),
|
||||
Value::Array(items) => {
|
||||
let mut values = Vec::with_capacity(items.len());
|
||||
for item in items {
|
||||
let Some(value) = item.as_str() else {
|
||||
return Err("IP 限制规则必须是字符串数组".to_string());
|
||||
};
|
||||
values.push(value.to_string());
|
||||
}
|
||||
Ok(normalize_ip_rules(Some(values))?.map(|rules| serde_json::json!(rules)))
|
||||
}
|
||||
_ => Err("IP 限制规则必须是字符串数组".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn ip_rules_allow(rules: Option<&[String]>, remote_ip: IpAddr) -> bool {
|
||||
let Some(rules) = rules else {
|
||||
return true;
|
||||
};
|
||||
if rules.is_empty() {
|
||||
return true;
|
||||
}
|
||||
|
||||
let mut has_allow_rule = false;
|
||||
let mut matched_allow_rule = false;
|
||||
for raw in rules {
|
||||
let rule = raw.trim();
|
||||
if rule.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let (deny, pattern) = match rule.strip_prefix('!') {
|
||||
Some(pattern) => (true, pattern.trim()),
|
||||
None => (false, rule),
|
||||
};
|
||||
let matched = ip_rule_pattern_matches(pattern, remote_ip);
|
||||
if deny && matched {
|
||||
return false;
|
||||
}
|
||||
if !deny {
|
||||
has_allow_rule = true;
|
||||
if matched {
|
||||
matched_allow_rule = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if has_allow_rule {
|
||||
matched_allow_rule
|
||||
} else {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn json_ip_rules_allow(value: Option<&Value>, remote_ip: IpAddr) -> bool {
|
||||
let Some(value) = value else {
|
||||
return true;
|
||||
};
|
||||
if value.is_null() {
|
||||
return true;
|
||||
}
|
||||
let Some(items) = value.as_array() else {
|
||||
return false;
|
||||
};
|
||||
let mut rules = Vec::with_capacity(items.len());
|
||||
for item in items {
|
||||
let Some(rule) = item.as_str() else {
|
||||
return false;
|
||||
};
|
||||
rules.push(rule.to_string());
|
||||
}
|
||||
ip_rules_allow(Some(&rules), remote_ip)
|
||||
}
|
||||
|
||||
fn normalize_ip_rule(raw: &str) -> Result<String, String> {
|
||||
if raw.is_empty() {
|
||||
return Err("IP 限制规则不能为空".to_string());
|
||||
}
|
||||
let (deny, pattern) = match raw.strip_prefix('!') {
|
||||
Some(pattern) => (true, pattern.trim()),
|
||||
None => (false, raw),
|
||||
};
|
||||
if pattern.is_empty() {
|
||||
return Err("IP 限制规则不能为空".to_string());
|
||||
}
|
||||
if !valid_ip_rule_pattern(pattern) {
|
||||
return Err(format!("无效的 IP 限制规则: {raw}"));
|
||||
}
|
||||
if deny {
|
||||
Ok(format!("!{pattern}"))
|
||||
} else {
|
||||
Ok(pattern.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn valid_ip_rule_pattern(pattern: &str) -> bool {
|
||||
if pattern == "*" {
|
||||
return true;
|
||||
}
|
||||
if pattern.parse::<IpAddr>().is_ok() {
|
||||
return true;
|
||||
}
|
||||
if valid_cidr_pattern(pattern) {
|
||||
return true;
|
||||
}
|
||||
valid_ipv4_wildcard_pattern(pattern)
|
||||
}
|
||||
|
||||
fn valid_cidr_pattern(pattern: &str) -> bool {
|
||||
let Some((host, prefix)) = pattern.split_once('/') else {
|
||||
return false;
|
||||
};
|
||||
let Ok(ip) = host.trim().parse::<IpAddr>() else {
|
||||
return false;
|
||||
};
|
||||
let Ok(prefix) = prefix.trim().parse::<u8>() else {
|
||||
return false;
|
||||
};
|
||||
match ip {
|
||||
IpAddr::V4(_) => prefix <= 32,
|
||||
IpAddr::V6(_) => prefix <= 128,
|
||||
}
|
||||
}
|
||||
|
||||
fn valid_ipv4_wildcard_pattern(pattern: &str) -> bool {
|
||||
if !pattern.contains('*') {
|
||||
return false;
|
||||
}
|
||||
let parts = pattern.split('.').collect::<Vec<_>>();
|
||||
parts.len() == 4
|
||||
&& parts
|
||||
.iter()
|
||||
.all(|part| *part == "*" || part.parse::<u8>().is_ok())
|
||||
}
|
||||
|
||||
fn ip_rule_pattern_matches(pattern: &str, remote_ip: IpAddr) -> bool {
|
||||
if pattern == "*" {
|
||||
return true;
|
||||
}
|
||||
if let Ok(ip) = pattern.parse::<IpAddr>() {
|
||||
return ip == remote_ip;
|
||||
}
|
||||
if ipv4_wildcard_matches(pattern, remote_ip) {
|
||||
return true;
|
||||
}
|
||||
let Some((network, prefix)) = pattern.split_once('/') else {
|
||||
return false;
|
||||
};
|
||||
let Ok(prefix) = prefix.trim().parse::<u8>() else {
|
||||
return false;
|
||||
};
|
||||
match (network.trim().parse::<IpAddr>(), remote_ip) {
|
||||
(Ok(IpAddr::V4(network)), IpAddr::V4(remote)) if prefix <= 32 => {
|
||||
let mask = if prefix == 0 {
|
||||
0
|
||||
} else {
|
||||
u32::MAX << (32 - prefix)
|
||||
};
|
||||
(u32::from(network) & mask) == (u32::from(remote) & mask)
|
||||
}
|
||||
(Ok(IpAddr::V6(network)), IpAddr::V6(remote)) if prefix <= 128 => {
|
||||
let mask = if prefix == 0 {
|
||||
0
|
||||
} else {
|
||||
u128::MAX << (128 - prefix)
|
||||
};
|
||||
(u128::from(network) & mask) == (u128::from(remote) & mask)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn ipv4_wildcard_matches(pattern: &str, remote_ip: IpAddr) -> bool {
|
||||
let IpAddr::V4(remote_ip) = remote_ip else {
|
||||
return false;
|
||||
};
|
||||
if !valid_ipv4_wildcard_pattern(pattern) {
|
||||
return false;
|
||||
}
|
||||
pattern
|
||||
.split('.')
|
||||
.zip(remote_ip.octets())
|
||||
.all(|(pattern_part, remote_part)| {
|
||||
pattern_part == "*" || pattern_part.parse::<u8>() == Ok(remote_part)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn deserialize_optional_json_patch<'de, D>(
|
||||
deserializer: D,
|
||||
) -> Result<Option<Option<Value>>, D::Error>
|
||||
@@ -72,6 +282,15 @@ where
|
||||
<Option<Value> as serde::Deserialize>::deserialize(deserializer).map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn deserialize_optional_string_list_patch<'de, D>(
|
||||
deserializer: D,
|
||||
) -> Result<Option<Option<Vec<String>>>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
<Option<Vec<String>> as serde::Deserialize>::deserialize(deserializer).map(Some)
|
||||
}
|
||||
|
||||
fn normalize_chat_pii_redaction_feature_settings(
|
||||
settings: &mut Map<String, Value>,
|
||||
) -> Result<(), String> {
|
||||
@@ -106,3 +325,80 @@ fn normalize_chat_pii_redaction_feature_object(
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{ip_rules_allow, json_ip_rules_allow, normalize_ip_rules, parse_json_ip_rules};
|
||||
use serde_json::json;
|
||||
use std::net::{IpAddr, Ipv4Addr};
|
||||
|
||||
fn v4(a: u8, b: u8, c: u8, d: u8) -> IpAddr {
|
||||
IpAddr::V4(Ipv4Addr::new(a, b, c, d))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_ip_rules_accepts_ip_cidr_wildcard_and_deny_rules() {
|
||||
let rules = normalize_ip_rules(Some(vec![
|
||||
" 203.0.113.10 ".to_string(),
|
||||
"10.0.0.0/24".to_string(),
|
||||
"192.168.*.*".to_string(),
|
||||
"! 10.0.0.13 ".to_string(),
|
||||
"203.0.113.10".to_string(),
|
||||
]))
|
||||
.expect("valid IP rules should normalize");
|
||||
|
||||
assert_eq!(
|
||||
rules,
|
||||
Some(vec![
|
||||
"203.0.113.10".to_string(),
|
||||
"10.0.0.0/24".to_string(),
|
||||
"192.168.*.*".to_string(),
|
||||
"!10.0.0.13".to_string(),
|
||||
]),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ip_rules_allow_applies_allow_rules_and_deny_overrides() {
|
||||
let rules = vec![
|
||||
"10.0.0.0/24".to_string(),
|
||||
"192.168.*.*".to_string(),
|
||||
"!10.0.0.13".to_string(),
|
||||
];
|
||||
|
||||
assert!(ip_rules_allow(Some(&rules), v4(10, 0, 0, 12)));
|
||||
assert!(ip_rules_allow(Some(&rules), v4(192, 168, 2, 3)));
|
||||
assert!(!ip_rules_allow(Some(&rules), v4(10, 0, 0, 13)));
|
||||
assert!(!ip_rules_allow(Some(&rules), v4(203, 0, 113, 10)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ip_rules_allow_defaults_to_allow_when_only_deny_rules_exist() {
|
||||
let rules = vec!["!10.0.*.*".to_string()];
|
||||
|
||||
assert!(!ip_rules_allow(Some(&rules), v4(10, 0, 0, 13)));
|
||||
assert!(ip_rules_allow(Some(&rules), v4(203, 0, 113, 10)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_json_ip_rules_normalizes_empty_and_string_arrays() {
|
||||
assert_eq!(
|
||||
parse_json_ip_rules(Some(&json!([" 203.0.113.10 ", "!10.0.0.13"])))
|
||||
.expect("valid JSON IP rules should parse"),
|
||||
Some(json!(["203.0.113.10", "!10.0.0.13"])),
|
||||
);
|
||||
assert_eq!(
|
||||
parse_json_ip_rules(Some(&json!([]))).expect("empty rules should parse"),
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_ip_rules_allow_rejects_invalid_stored_shape() {
|
||||
assert!(!json_ip_rules_allow(
|
||||
Some(&json!({"bad": true})),
|
||||
v4(10, 0, 0, 1)
|
||||
));
|
||||
assert!(!json_ip_rules_allow(Some(&json!([123])), v4(10, 0, 0, 1)));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -221,6 +221,7 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("create_endpoint"))
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("batch_delete_keys"))
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("refresh_quota"))
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("query_key_balance"))
|
||||
| (Some("endpoints_manage"), http::Method::PUT, Some("update_key"))
|
||||
| (Some("endpoints_manage"), http::Method::PUT, Some("update_endpoint"))
|
||||
| (Some("modules_manage"), http::Method::PUT, Some("set_enabled"))
|
||||
|
||||
@@ -498,6 +498,7 @@ mod tests {
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: None,
|
||||
allowed_models: None,
|
||||
ip_rules: None,
|
||||
});
|
||||
let state = AppState::new().expect("state should build for tests");
|
||||
|
||||
@@ -537,6 +538,7 @@ mod tests {
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: None,
|
||||
allowed_models: None,
|
||||
ip_rules: None,
|
||||
});
|
||||
let state = AppState::new().expect("state should build for tests");
|
||||
|
||||
@@ -576,6 +578,7 @@ mod tests {
|
||||
admin_bypass_limits: true,
|
||||
local_rejection: None,
|
||||
allowed_models: None,
|
||||
ip_rules: None,
|
||||
});
|
||||
let state = AppState::new().expect("state should build for tests");
|
||||
|
||||
@@ -619,6 +622,7 @@ mod tests {
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: None,
|
||||
allowed_models: None,
|
||||
ip_rules: None,
|
||||
});
|
||||
let state = AppState::new().expect("state should build for tests");
|
||||
|
||||
|
||||
@@ -129,6 +129,7 @@ pub(super) fn sample_auth_snapshot(api_key_id: &str) -> GatewayAuthApiKeySnapsho
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
api_key_ip_rules: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -169,6 +169,16 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn count_export_users(
|
||||
&self,
|
||||
query: &aether_data::repository::users::UserExportListQuery,
|
||||
) -> Result<u64, GatewayError> {
|
||||
self.data
|
||||
.count_export_users(query)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn find_export_user_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
|
||||
@@ -927,6 +927,7 @@ mod tests {
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: None,
|
||||
allowed_models: Some(vec!["gpt-4.1".to_string()]),
|
||||
ip_rules: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1107,7 +1107,8 @@ async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_pla
|
||||
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
|
||||
assert_eq!(seen_plan.body_json["stream"], true);
|
||||
assert_eq!(seen_plan.body_json["input"], "Draw a mountain observatory");
|
||||
assert!(seen_plan.body_json.get("tools").is_none());
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
|
||||
@@ -1490,7 +1490,24 @@ async fn gateway_verifies_admin_provider_ops_locally_for_new_api_with_trusted_ad
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("session=foo")
|
||||
);
|
||||
assert!(headers.contains_key("sec-ch-ua"));
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(axum::http::header::USER_AGENT)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("cc-switch/1.0")
|
||||
);
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(axum::http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("application/json")
|
||||
);
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(axum::http::header::ACCEPT)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("application/json")
|
||||
);
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
|
||||
encoder
|
||||
.write_all(
|
||||
|
||||
@@ -300,7 +300,11 @@ async fn gateway_handles_admin_users_root_locally_with_trusted_admin_principal()
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload.as_array().expect("list payload should be array");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["skip"], 0);
|
||||
assert_eq!(payload["limit"], 20);
|
||||
assert_eq!(payload["has_more"], false);
|
||||
let items = payload["items"].as_array().expect("items should be array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "user-1");
|
||||
assert_eq!(items[0]["email"], "alice@example.com");
|
||||
@@ -332,9 +336,10 @@ async fn gateway_handles_admin_users_root_locally_with_trusted_admin_principal()
|
||||
.json()
|
||||
.await
|
||||
.expect("search json body should parse");
|
||||
let search_items = search_payload
|
||||
assert_eq!(search_payload["total"], 1);
|
||||
let search_items = search_payload["items"]
|
||||
.as_array()
|
||||
.expect("search list payload should be array");
|
||||
.expect("search items should be array");
|
||||
assert_eq!(search_items.len(), 1);
|
||||
assert_eq!(search_items[0]["id"], "user-3");
|
||||
assert_eq!(search_items[0]["email"], "carol@example.com");
|
||||
@@ -355,9 +360,10 @@ async fn gateway_handles_admin_users_root_locally_with_trusted_admin_principal()
|
||||
.json()
|
||||
.await
|
||||
.expect("id search json body should parse");
|
||||
let id_search_items = id_search_payload
|
||||
assert_eq!(id_search_payload["total"], 1);
|
||||
let id_search_items = id_search_payload["items"]
|
||||
.as_array()
|
||||
.expect("id search list payload should be array");
|
||||
.expect("id search items should be array");
|
||||
assert_eq!(id_search_items.len(), 1);
|
||||
assert_eq!(id_search_items[0]["id"], "user-3");
|
||||
|
||||
@@ -377,9 +383,11 @@ async fn gateway_handles_admin_users_root_locally_with_trusted_admin_principal()
|
||||
.json()
|
||||
.await
|
||||
.expect("limited search json body should parse");
|
||||
let limited_search_items = limited_search_payload
|
||||
assert_eq!(limited_search_payload["total"], 3);
|
||||
assert_eq!(limited_search_payload["has_more"], true);
|
||||
let limited_search_items = limited_search_payload["items"]
|
||||
.as_array()
|
||||
.expect("limited search list payload should be array");
|
||||
.expect("limited search items should be array");
|
||||
assert_eq!(limited_search_items.len(), 2);
|
||||
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
@@ -1057,7 +1065,8 @@ async fn gateway_handles_admin_users_root_locally_with_bearer_admin_session() {
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload.as_array().expect("list payload should be array");
|
||||
assert_eq!(payload["total"], 1);
|
||||
let items = payload["items"].as_array().expect("items should be array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "user-1");
|
||||
assert_eq!(items[0]["email"], "alice@example.com");
|
||||
|
||||
@@ -4378,7 +4378,7 @@ async fn gateway_handles_wallet_balance_locally_without_proxying_upstream() {
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
||||
let auth_now = Utc::now();
|
||||
let usage_now = auth_now - chrono::Duration::minutes(30);
|
||||
let usage_now = auth_now;
|
||||
let user = sample_auth_user(auth_now);
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
|
||||
@@ -54,6 +54,7 @@ pub(super) fn sample_auth_context() -> GatewayControlAuthContext {
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: None,
|
||||
allowed_models: None,
|
||||
ip_rules: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -307,6 +307,7 @@ mod tests {
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
api_key_ip_rules: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user