mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Enforce API key IP restrictions in proxy auth
This commit is contained in:
@@ -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(
|
||||
@@ -566,6 +567,7 @@ mod tests {
|
||||
admin_bypass_limits: false,
|
||||
local_rejection: None,
|
||||
allowed_models: Some(allowed_models),
|
||||
allowed_ips: 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) 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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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"}]}"#,
|
||||
|
||||
Reference in New Issue
Block a user