i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Enforce API key IP restrictions in proxy auth

Browse Source
This commit is contained in:
RWDai
2026-05-18 20:47:28 +08:00
parent 276d19b63c
commit fc12cc8a36
4 changed files with 81 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
apps/aether-gateway/src/api/response.rs
View File

@@ -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}"),
),
}
}

2
apps/aether-gateway/src/control/auth/gate.rs
View File

@@ -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(
@@ -566,6 +567,7 @@ mod tests {
admin_bypass_limits: false,
local_rejection: None,
allowed_models: Some(allowed_models),
allowed_ips: None,
});
decision
}

8
apps/aether-gateway/src/control/auth/resolution.rs
View File

@@ -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) allowed_ips: 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,
allowed_ips: None,
}));
};
@@ -638,6 +644,7 @@ async fn resolve_trusted_auth_context(
admin_bypass_limits: false,
local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey),
allowed_models: None,
allowed_ips: None,
}));
};
@@ -728,6 +735,7 @@ async fn build_data_backed_auth_context(
&& !snapshot.api_key_is_standalone,
local_rejection,
allowed_models,
allowed_ips: snapshot.api_key_allowed_ips,
}
}

68
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -151,6 +151,18 @@ fn remote_ip_allowed(allowed_ips: Option<&serde_json::Value>, remote_ip: std::ne
.any(|value| ip_or_cidr_matches(value, remote_ip))
}
fn api_key_remote_ip_allowed(allowed_ips: Option<&[String]>, remote_ip: std::net::IpAddr) -> bool {
let Some(allowed_ips) = allowed_ips else {
return true;
};
if allowed_ips.is_empty() {
return false;
}
allowed_ips
.iter()
.any(|value| ip_or_cidr_matches(value, remote_ip))
}
fn ip_or_cidr_matches(pattern: &str, remote_ip: std::net::IpAddr) -> bool {
let pattern = pattern.trim();
if pattern.is_empty() {
@@ -917,6 +929,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.allowed_ips.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
@@ -1909,14 +1946,39 @@ 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_matches_exact_ip_and_cidr() {
let allowed_ips = vec!["198.51.100.1".to_string(), "203.0.113.0/24".to_string()];
assert!(api_key_remote_ip_allowed(
Some(&allowed_ips),
"198.51.100.1".parse().expect("valid ip"),
));
assert!(api_key_remote_ip_allowed(
Some(&allowed_ips),
"203.0.113.42".parse().expect("valid ip"),
));
assert!(!api_key_remote_ip_allowed(
Some(&allowed_ips),
"203.0.114.42".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"}]}"#,