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

feat(security): harden gateway request and runtime controls

Browse Source
This commit is contained in:
elky
2026-07-12 14:10:54 +08:00
parent bc1da3bf3f
commit 7f61bb43c7
19 changed files with 975 additions and 340 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/state/app.rs
+36 -1
View File
@@ -10,7 +10,7 @@ use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_runtime::ConcurrencyGate;
use aether_runtime_state::{RuntimeSemaphore, RuntimeState};
use dashmap::DashMap;
use tokio::sync::Mutex as TokioMutex;
use tokio::sync::{Mutex as TokioMutex, Semaphore};
use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
use super::super::cache::{
@@ -36,6 +36,11 @@ const DEFAULT_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 120_000;
const MIN_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 1_000;
const MAX_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 600_000;
const REQUEST_BODY_READ_TIMEOUT_MS_ENV: &str = "AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS";
const DEFAULT_REQUEST_BODY_BUFFER_BUDGET_MB: usize = 256;
const MIN_REQUEST_BODY_BUFFER_BUDGET_MB: usize = 64;
const MAX_REQUEST_BODY_BUFFER_BUDGET_MB: usize = 16 * 1024;
const REQUEST_BODY_BUFFER_BUDGET_MB_ENV: &str = "AETHER_GATEWAY_REQUEST_BODY_BUFFER_BUDGET_MB";
pub(crate) const REQUEST_BODY_BUFFER_PERMIT_BYTES: usize = 64 * 1024;
const DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 30_000;
const MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 500;
@@ -90,6 +95,8 @@ impl std::fmt::Debug for TestExecutionRuntimeSyncOverride {
#[derive(Debug, Clone)]
pub(crate) struct FrontdoorRuntimeGuardConfig {
pub(crate) request_body_read_timeout: Duration,
pub(crate) request_body_buffer_budget_bytes: usize,
pub(crate) request_body_buffer_budget_permits: usize,
pub(crate) local_execution_planning_timeout: Duration,
pub(crate) internal_gate_queue_budget: Duration,
pub(crate) auth_capacity_cache_ttl: Duration,
@@ -108,6 +115,8 @@ impl FrontdoorRuntimeGuardConfig {
MIN_REQUEST_BODY_READ_TIMEOUT_MS,
MAX_REQUEST_BODY_READ_TIMEOUT_MS,
),
request_body_buffer_budget_bytes: request_body_buffer_budget_bytes_from_env(),
request_body_buffer_budget_permits: request_body_buffer_budget_permits_from_env(),
local_execution_planning_timeout: env_duration_ms(
LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV,
DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
@@ -140,6 +149,8 @@ impl FrontdoorRuntimeGuardConfig {
) -> Self {
Self {
request_body_read_timeout,
request_body_buffer_budget_bytes: DEFAULT_REQUEST_BODY_BUFFER_BUDGET_MB * 1024 * 1024,
request_body_buffer_budget_permits: DEFAULT_REQUEST_BODY_BUFFER_BUDGET_MB * 16,
local_execution_planning_timeout,
internal_gate_queue_budget: Duration::from_millis(
DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS,
@@ -153,6 +164,27 @@ impl FrontdoorRuntimeGuardConfig {
}
}
fn request_body_buffer_budget_mb_from_env() -> usize {
std::env::var(REQUEST_BODY_BUFFER_BUDGET_MB_ENV)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_REQUEST_BODY_BUFFER_BUDGET_MB)
.clamp(
MIN_REQUEST_BODY_BUFFER_BUDGET_MB,
MAX_REQUEST_BODY_BUFFER_BUDGET_MB,
)
}
fn request_body_buffer_budget_bytes_from_env() -> usize {
request_body_buffer_budget_mb_from_env().saturating_mul(1024 * 1024)
}
fn request_body_buffer_budget_permits_from_env() -> usize {
request_body_buffer_budget_bytes_from_env().saturating_add(REQUEST_BODY_BUFFER_PERMIT_BYTES - 1)
/ REQUEST_BODY_BUFFER_PERMIT_BYTES
}
fn env_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Duration {
let ms = std::env::var(key)
.ok()
@@ -337,6 +369,7 @@ pub struct AppState {
pub(crate) video_tasks: Arc<VideoTaskService>,
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
pub(crate) frontdoor_runtime_guards: Arc<FrontdoorRuntimeGuardConfig>,
pub(crate) request_body_buffer_budget: Arc<Semaphore>,
pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) auth_snapshot_load_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) candidate_planning_gate: Option<Arc<ConcurrencyGate>>,
@@ -346,6 +379,8 @@ pub struct AppState {
pub(crate) client: reqwest::Client,
pub(crate) auth_context_cache: Arc<AuthContextCache>,
pub(crate) auth_snapshot_cache: Arc<AuthSnapshotCache>,
pub(crate) admin_security_blacklist_cache: Arc<ValueCache<String, bool>>,
pub(crate) admin_security_whitelist_cache: Arc<ValueCache<String, Vec<String>>>,
pub(crate) user_model_capability_settings_cache: Arc<JsonValueCache<String>>,
pub(crate) user_feature_settings_cache: Arc<JsonValueCache<String>>,
pub(crate) auth_api_key_force_capabilities_cache:
apps/aether-gateway/src/state/core.rs
+56
View File
@@ -266,6 +266,9 @@ impl AppState {
)),
video_task_poller: None,
frontdoor_runtime_guards: Arc::clone(&frontdoor_runtime_guards),
request_body_buffer_budget: Arc::new(tokio::sync::Semaphore::new(
frontdoor_runtime_guards.request_body_buffer_budget_permits,
)),
request_gate: None,
auth_snapshot_load_gate: frontdoor_runtime_guards
.auth_snapshot_load_gate_limit
@@ -286,6 +289,8 @@ impl AppState {
client,
auth_context_cache: Arc::new(AuthContextCache::default()),
auth_snapshot_cache: Arc::new(AuthSnapshotCache::default()),
admin_security_blacklist_cache: Arc::new(ValueCache::default()),
admin_security_whitelist_cache: Arc::new(ValueCache::default()),
user_model_capability_settings_cache: Arc::new(JsonValueCache::default()),
user_feature_settings_cache: Arc::new(JsonValueCache::default()),
auth_api_key_force_capabilities_cache: Arc::new(JsonValueCache::default()),
@@ -480,6 +485,8 @@ impl AppState {
pub fn with_runtime_state(mut self, runtime_state: Arc<RuntimeState>) -> Self {
self.runtime_state = runtime_state;
self.admin_security_blacklist_cache.clear();
self.admin_security_whitelist_cache.clear();
self.data = Arc::new(
(*self.data)
.clone()
@@ -1062,6 +1069,38 @@ impl AppState {
pub(crate) async fn metric_samples(&self) -> Vec<MetricSample> {
let mut samples = vec![service_up_sample("aether-gateway")];
let request_body_buffer_budget_bytes = self
.frontdoor_runtime_guards
.request_body_buffer_budget_bytes;
let request_body_buffer_available_bytes = self
.request_body_buffer_budget
.available_permits()
.saturating_mul(super::REQUEST_BODY_BUFFER_PERMIT_BYTES)
.min(request_body_buffer_budget_bytes);
samples.extend([
MetricSample::new(
"request_body_buffer_budget_bytes",
"Configured weighted request body buffering budget in bytes.",
MetricKind::Gauge,
u64::try_from(request_body_buffer_budget_bytes).unwrap_or(u64::MAX),
),
MetricSample::new(
"request_body_buffer_available_bytes",
"Currently available weighted request body buffering budget in bytes.",
MetricKind::Gauge,
u64::try_from(request_body_buffer_available_bytes).unwrap_or(u64::MAX),
),
MetricSample::new(
"request_body_buffer_in_use_bytes",
"Currently reserved weighted request body buffering budget in bytes.",
MetricKind::Gauge,
u64::try_from(
request_body_buffer_budget_bytes
.saturating_sub(request_body_buffer_available_bytes),
)
.unwrap_or(u64::MAX),
),
]);
if let Some(snapshot) = self.request_concurrency_snapshot() {
samples.extend(snapshot.to_metric_samples("gateway_requests"));
}
@@ -2935,6 +2974,23 @@ mod tests {
}));
}
#[tokio::test]
async fn metric_samples_include_request_body_buffer_budget_usage() {
let state = AppState::new().expect("app state should build");
let _permit = Arc::clone(&state.request_body_buffer_budget)
.acquire_many_owned(2)
.await
.expect("request body budget should be open");
let samples = state.metric_samples().await;
assert!(samples.iter().any(|sample| {
sample.name == "request_body_buffer_in_use_bytes"
&& sample.value
== u64::try_from(2 * crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES)
.unwrap_or(u64::MAX)
}));
}
#[test]
fn scheduler_affinity_epoch_blocks_stale_rewarm_after_invalidation() {
let state = AppState::new().expect("app state should build");
apps/aether-gateway/src/state/mod.rs
+1 -1
View File
@@ -29,7 +29,7 @@ pub(crate) use self::admin_types::{
pub use self::app::AppState;
pub(crate) use self::app::{
upstream_target_gate_auto_limit, upstream_target_gate_limit_from_env,
FrontdoorRuntimeGuardConfig,
FrontdoorRuntimeGuardConfig, REQUEST_BODY_BUFFER_PERMIT_BYTES,
};
pub(crate) use self::cache::{
CachedProviderTransportSnapshot, AUTH_API_KEY_LAST_USED_MAX_ENTRIES,
apps/aether-gateway/src/state/runtime/security.rs
+77 -34
View File
@@ -1,31 +1,68 @@
use crate::state::AdminSecurityBlacklistEntry;
use crate::{AppState, GatewayError};
use std::net::IpAddr;
use std::sync::LazyLock;
use std::time::Duration;
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
const ADMIN_SECURITY_CACHE_TTL_MS_ENV: &str = "AETHER_GATEWAY_SECURITY_CACHE_TTL_MS";
const DEFAULT_ADMIN_SECURITY_CACHE_TTL_MS: u64 = 1_000;
const MAX_ADMIN_SECURITY_CACHE_TTL_MS: u64 = 30_000;
const ADMIN_SECURITY_WHITELIST_CACHE_KEY: &str = "rules";
static ADMIN_SECURITY_CACHE_TTL: LazyLock<Duration> = LazyLock::new(|| {
let ttl_ms = std::env::var(ADMIN_SECURITY_CACHE_TTL_MS_ENV)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(DEFAULT_ADMIN_SECURITY_CACHE_TTL_MS)
.min(MAX_ADMIN_SECURITY_CACHE_TTL_MS);
Duration::from_millis(ttl_ms)
});
fn admin_security_cache_ttl() -> Duration {
*ADMIN_SECURITY_CACHE_TTL
}
impl AppState {
pub(crate) async fn admin_security_ip_blacklisted(
&self,
ip_address: IpAddr,
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
self.runtime_state
.kv_exists(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{ip_address}"))
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
let cache_key = ip_address.to_string();
let runtime_key = format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{cache_key}");
Ok(self
.admin_security_blacklist_cache
.get_or_load_once(cache_key, admin_security_cache_ttl(), || async {
self.runtime_state
.kv_exists(&runtime_key)
.await
.map(Some)
.map_err(|err| GatewayError::Internal(err.to_string()))
})
.await?
.unwrap_or(false))
}
pub(crate) async fn admin_security_ip_whitelisted(
&self,
ip_address: IpAddr,
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
let rules = self
.runtime_state
.set_members(ADMIN_SECURITY_WHITELIST_KEY)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
.admin_security_whitelist_cache
.get_or_load_once(
ADMIN_SECURITY_WHITELIST_CACHE_KEY.to_string(),
admin_security_cache_ttl(),
|| async {
self.runtime_state
.set_members(ADMIN_SECURITY_WHITELIST_KEY)
.await
.map(Some)
.map_err(|err| GatewayError::Internal(err.to_string()))
},
)
.await?
.unwrap_or_default();
Ok(rules
.iter()
.any(|rule| crate::handlers::shared::ip_rule_pattern_matches(rule.trim(), ip_address)))
@@ -37,8 +74,6 @@ impl AppState {
reason: &str,
ttl_seconds: Option<u64>,
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
let key = format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{ip_address}");
self.runtime_state
.kv_set(
@@ -47,28 +82,40 @@ impl AppState {
ttl_seconds.map(std::time::Duration::from_secs),
)
.await
.map(|_| true)
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if let Ok(ip_address) = ip_address.parse::<IpAddr>() {
self.admin_security_blacklist_cache.insert(
ip_address.to_string(),
Some(true),
admin_security_cache_ttl(),
);
}
Ok(true)
}
pub(crate) async fn remove_admin_security_blacklist(
&self,
ip_address: &str,
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
let key = format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{ip_address}");
self.runtime_state
let removed = self
.runtime_state
.kv_delete(&key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if let Ok(ip_address) = ip_address.parse::<IpAddr>() {
self.admin_security_blacklist_cache.insert(
ip_address.to_string(),
Some(false),
admin_security_cache_ttl(),
);
}
Ok(removed)
}
pub(crate) async fn admin_security_blacklist_stats(
&self,
) -> Result<(bool, usize, Option<String>), GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
let total = self
.runtime_state
.scan_keys(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}*"), 100)
@@ -81,8 +128,6 @@ impl AppState {
pub(crate) async fn list_admin_security_blacklist(
&self,
) -> Result<Vec<AdminSecurityBlacklistEntry>, GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
let keys = self
.runtime_state
.scan_keys(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}*"), 100)
@@ -123,30 +168,28 @@ impl AppState {
&self,
ip_address: &str,
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
self.runtime_state
.set_add(ADMIN_SECURITY_WHITELIST_KEY, ip_address)
.await
.map(|_| true)
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.admin_security_whitelist_cache.clear();
Ok(true)
}
pub(crate) async fn remove_admin_security_whitelist(
&self,
ip_address: &str,
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
self.runtime_state
let removed = self
.runtime_state
.set_remove(ADMIN_SECURITY_WHITELIST_KEY, ip_address)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.admin_security_whitelist_cache.clear();
Ok(removed)
}
pub(crate) async fn list_admin_security_whitelist(&self) -> Result<Vec<String>, GatewayError> {
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
self.runtime_state
.set_members(ADMIN_SECURITY_WHITELIST_KEY)
.await