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

refactor: extract runtime state backends

Browse Source
This commit is contained in:
fawney19
2026-05-08 00:18:12 +08:00
parent 6f620d92be
commit 6247ac3edc
111 changed files with 4358 additions and 3203 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
apps/aether-gateway/src/state/app.rs
View File

@@ -2,7 +2,8 @@ use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use aether_runtime::{ConcurrencyGate, DistributedConcurrencyGate};
use aether_runtime::ConcurrencyGate;
use aether_runtime_state::{RuntimeSemaphore, RuntimeState};
use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
use super::super::cache::{
@@ -47,11 +48,12 @@ pub struct AppState {
#[cfg(test)]
pub(crate) execution_runtime_sync_override: Option<TestExecutionRuntimeSyncOverride>,
pub(crate) data: Arc<GatewayDataState>,
pub(crate) runtime_state: Arc<RuntimeState>,
pub(crate) usage_runtime: Arc<usage::UsageRuntime>,
pub(crate) video_tasks: Arc<VideoTaskService>,
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
pub(crate) client: reqwest::Client,
pub(crate) auth_context_cache: Arc<AuthContextCache>,
pub(crate) auth_api_key_last_used_cache: Arc<AuthApiKeyLastUsedCache>,

206
apps/aether-gateway/src/state/core.rs
View File

@@ -9,9 +9,12 @@ use aether_data::repository::proxy_nodes::{
};
use aether_http::{build_http_client, HttpClientConfig};
use aether_runtime::{
service_up_sample, AdmissionPermit, ConcurrencyGate, ConcurrencySnapshot,
DistributedConcurrencyError, DistributedConcurrencyGate, DistributedConcurrencySnapshot,
MetricKind, MetricLabel, MetricSample,
service_up_sample, AdmissionPermit, ConcurrencyGate, ConcurrencySnapshot, MetricKind,
MetricLabel, MetricSample,
};
use aether_runtime_state::{
MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeSemaphore, RuntimeSemaphoreError,
RuntimeSemaphoreSnapshot, RuntimeState,
};
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
use tokio::task::JoinHandle;
@@ -53,21 +56,32 @@ use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(3);
impl AppState {
fn usage_worker_queue_for(
runtime_state: &Arc<RuntimeState>,
) -> Option<Arc<dyn RuntimeQueueStore>> {
if runtime_state.is_redis() {
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
Some(queue)
} else {
None
}
}
fn spawn_scheduler_affinity_redis_write(
&self,
cache_key: &str,
target: &SchedulerAffinityTarget,
ttl: Duration,
) {
let Some(runner) = self.redis_kv_runner() else {
if self.runtime_state.is_memory() {
return;
};
}
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
let cache_key = cache_key.to_string();
let namespaced_cache_key = runner.keyspace().key(&cache_key);
let runtime_state = self.runtime_state.clone();
let provider_id = target.provider_id.clone();
let endpoint_id = target.endpoint_id.clone();
let key_id = target.key_id.clone();
@@ -76,56 +90,55 @@ impl AppState {
let expire_at = now_unix_secs.saturating_add(ttl_seconds);
handle.spawn(async move {
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await
else {
return;
};
let script = r#"
local existing = redis.call('GET', KEYS[1])
local request_count = 0
local created_at = tonumber(ARGV[4])
if existing then
local ok, payload = pcall(cjson.decode, existing)
if ok and type(payload) == 'table' then
if type(payload['request_count']) == 'number' then
request_count = payload['request_count']
end
if type(payload['created_at']) == 'number' then
created_at = payload['created_at']
end
end
end
request_count = request_count + 1
local payload = {
provider_id = ARGV[1],
endpoint_id = ARGV[2],
key_id = ARGV[3],
created_at = created_at,
expire_at = tonumber(ARGV[5]),
request_count = request_count
}
redis.call('SETEX', KEYS[1], tonumber(ARGV[6]), cjson.encode(payload))
return request_count
"#;
let _ = redis::cmd("EVAL")
.arg(script)
.arg(1)
.arg(&namespaced_cache_key)
.arg(&provider_id)
.arg(&endpoint_id)
.arg(&key_id)
.arg(now_unix_secs)
.arg(expire_at)
.arg(ttl_seconds)
.query_async::<i64>(&mut connection)
.await;
let existing = runtime_state
.kv_get(&cache_key)
.await
.ok()
.flatten()
.and_then(|raw| serde_json::from_str::<serde_json::Value>(&raw).ok());
let request_count = existing
.as_ref()
.and_then(|value| value.get("request_count"))
.and_then(serde_json::Value::as_u64)
.unwrap_or_default()
.saturating_add(1);
let created_at = existing
.as_ref()
.and_then(|value| value.get("created_at"))
.and_then(serde_json::Value::as_u64)
.unwrap_or(now_unix_secs);
let payload = serde_json::json!({
"provider_id": provider_id,
"endpoint_id": endpoint_id,
"key_id": key_id,
"created_at": created_at,
"expire_at": expire_at,
"request_count": request_count,
});
if let Ok(serialized) = serde_json::to_string(&payload) {
let _ = runtime_state
.kv_set(
&cache_key,
serialized,
Some(Duration::from_secs(ttl_seconds)),
)
.await;
}
});
}
pub(crate) fn replace_data_state(&mut self, data: Arc<GatewayDataState>) {
self.clear_provider_transport_snapshot_cache();
self.system_config_cache.clear();
self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data(Arc::clone(&data));
let data = Arc::new(
(*data)
.clone()
.with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)),
);
self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data_and_runtime_state(
Arc::clone(&data),
self.runtime_state.clone(),
);
self.data = data;
}
@@ -153,7 +166,11 @@ return request_count
}
fn build(execution_runtime_override_base_url: Option<String>) -> Result<Self, reqwest::Error> {
let data = Arc::new(GatewayDataState::disabled());
let runtime_state = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let data = Arc::new(
GatewayDataState::disabled()
.with_usage_worker_queue(Self::usage_worker_queue_for(&runtime_state)),
);
let client = build_http_client(&HttpClientConfig {
connect_timeout_ms: Some(10_000),
request_timeout_ms: Some(300_000),
@@ -168,6 +185,7 @@ return request_count
#[cfg(test)]
execution_runtime_sync_override: None,
data: Arc::clone(&data),
runtime_state: runtime_state.clone(),
usage_runtime: Arc::new(usage::UsageRuntime::disabled()),
video_tasks: Arc::new(VideoTaskService::new(
VideoTaskTruthSourceMode::PythonSyncReport,
@@ -188,7 +206,10 @@ return request_count
frontdoor_user_rpm: Arc::new(FrontdoorUserRpmLimiter::new(
FrontdoorUserRpmConfig::default(),
)),
tunnel: crate::tunnel::EmbeddedTunnelState::with_data(data),
tunnel: crate::tunnel::EmbeddedTunnelState::with_data_and_runtime_state(
data,
runtime_state.clone(),
),
provider_transport_snapshot_cache: Arc::new(StdMutex::new(HashMap::new())),
provider_key_rpm_resets: Arc::new(StdMutex::new(HashMap::new())),
local_execution_runtime_miss_diagnostics: Arc::new(StdMutex::new(HashMap::new())),
@@ -261,11 +282,12 @@ return request_count
instance_id: impl Into<String>,
relay_base_url: Option<impl Into<String>>,
) -> Self {
self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data_and_identity(
self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data_identity_and_runtime_state(
Arc::clone(&self.data),
instance_id,
relay_base_url,
90,
self.runtime_state.clone(),
);
self
}
@@ -334,10 +356,21 @@ return request_count
self
}
pub fn with_distributed_request_concurrency_gate(
mut self,
gate: DistributedConcurrencyGate,
) -> Self {
pub fn with_runtime_state(mut self, runtime_state: Arc<RuntimeState>) -> Self {
self.runtime_state = runtime_state;
self.data = Arc::new(
(*self.data)
.clone()
.with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)),
);
self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data_and_runtime_state(
Arc::clone(&self.data),
self.runtime_state.clone(),
);
self
}
pub fn with_distributed_request_concurrency_gate(mut self, gate: RuntimeSemaphore) -> Self {
self.distributed_request_gate = Some(Arc::new(gate));
self
}
@@ -624,7 +657,7 @@ return request_count
pub(crate) async fn distributed_request_concurrency_snapshot(
&self,
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
) -> Result<Option<RuntimeSemaphoreSnapshot>, RuntimeSemaphoreError> {
match self.distributed_request_gate.as_ref() {
Some(gate) => gate.snapshot().await.map(Some),
None => Ok(None),
@@ -767,11 +800,62 @@ return request_count
}
pub fn has_redis_data_backend(&self) -> bool {
self.data.has_redis_backend()
self.runtime_state.is_redis()
}
pub(crate) fn redis_kv_runner(&self) -> Option<aether_data::driver::redis::RedisKvRunner> {
self.data.kv_runner()
pub(crate) fn runtime_state_backend(&self) -> &'static str {
self.runtime_state.backend_kind().as_str()
}
pub fn runtime_state(&self) -> &RuntimeState {
self.runtime_state.as_ref()
}
pub(crate) async fn runtime_kv_setex(
&self,
key: &str,
value: &str,
ttl_seconds: u64,
) -> Result<(), GatewayError> {
self.runtime_state
.kv_set(
key,
value.to_string(),
Some(Duration::from_secs(ttl_seconds)),
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn runtime_kv_get(&self, key: &str) -> Result<Option<String>, GatewayError> {
self.runtime_state
.kv_get(key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn runtime_kv_getdel(
&self,
key: &str,
) -> Result<Option<String>, GatewayError> {
self.runtime_state
.kv_take(key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn runtime_kv_del(&self, key: &str) -> Result<bool, GatewayError> {
self.runtime_state
.kv_delete(key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn runtime_kv_exists(&self, key: &str) -> Result<bool, GatewayError> {
self.runtime_state
.kv_exists(key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) fn remove_scheduler_affinity_cache_entry(&self, cache_key: &str) -> bool {

14
apps/aether-gateway/src/state/integrations.rs
View File

@@ -159,19 +159,19 @@ impl ModelFetchRuntimeState for AppState {
key_id: &str,
cached_models: &[Value],
) {
let Some(runner) = AppState::redis_kv_runner(self) else {
return;
};
let Ok(serialized) = serde_json::to_string(&aggregate_models_for_cache(cached_models))
else {
return;
};
let cache_key = format!("upstream_models:{provider_id}:{key_id}");
if let Err(err) = runner
.setex(
if let Err(err) = self
.runtime_state
.kv_set(
&cache_key,
&serialized,
Some(model_fetch_interval_minutes().saturating_mul(60)),
serialized,
Some(std::time::Duration::from_secs(
model_fetch_interval_minutes().saturating_mul(60),
)),
)
.await
{

8
apps/aether-gateway/src/state/oauth.rs
View File

@@ -833,7 +833,7 @@ impl AppState {
&self,
transport: &provider_transport::GatewayProviderTransportSnapshot,
) -> Result<Option<provider_transport::LocalResolvedOAuthRequestAuth>, GatewayError> {
let distributed_lock = self.data.oauth_refresh_lock_runner();
let distributed_lock = self.runtime_state.as_ref();
let lock_owner = format!("aether-gateway-{}", std::process::id());
let mut current_transport = transport.clone();
let executor = GatewayLocalOAuthHttpExecutor { state: self };
@@ -844,7 +844,7 @@ impl AppState {
.resolve_with_result(
&executor,
&current_transport,
distributed_lock.as_ref(),
Some(distributed_lock),
Some(lock_owner.as_str()),
)
.await
@@ -929,7 +929,7 @@ impl AppState {
Option<provider_transport::CachedOAuthEntry>,
provider_transport::LocalOAuthRefreshError,
> {
let distributed_lock = self.data.oauth_refresh_lock_runner();
let distributed_lock = self.runtime_state.as_ref();
let lock_owner = format!("aether-gateway-admin-{}", std::process::id());
let mut current_transport = transport.clone();
current_transport.key.decrypted_api_key = "__placeholder__".to_string();
@@ -958,7 +958,7 @@ impl AppState {
.force_refresh_with_result(
&executor,
&current_transport,
distributed_lock.as_ref(),
Some(distributed_lock),
Some(lock_owner.as_str()),
)
.await?;

13
apps/aether-gateway/src/state/runtime/gemini_files.rs
View File

@@ -77,16 +77,21 @@ impl AppState {
value: &str,
ttl_seconds: u64,
) -> Result<(), GatewayError> {
self.data
.cache_set_string_with_ttl(key, value, ttl_seconds)
self.runtime_state
.kv_set(
key,
value.to_string(),
Some(std::time::Duration::from_secs(ttl_seconds)),
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn cache_delete_key(&self, key: &str) -> Result<(), GatewayError> {
self.data
.cache_delete_key(key)
self.runtime_state
.kv_delete(key)
.await
.map(|_| ())
.map_err(|err| GatewayError::Internal(err.to_string()))
}
}

2
apps/aether-gateway/src/state/runtime/mod.rs
View File

@@ -67,7 +67,7 @@ impl AppState {
}
pub fn has_usage_worker_backend(&self) -> bool {
self.data.has_usage_worker_runner()
self.data.has_usage_worker_queue()
}
pub fn has_wallet_data_reader(&self) -> bool {

334
apps/aether-gateway/src/state/runtime/security.rs
View File

@@ -10,41 +10,16 @@ impl AppState {
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
if let Some(runner) = self.redis_kv_runner() {
let mut connection = match runner.client().get_multiplexed_async_connection().await {
Ok(value) => value,
Err(_) => return Ok(false),
};
let key = runner
.keyspace()
.key(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{ip_address}"));
let result = if let Some(ttl_seconds) = ttl_seconds {
redis::cmd("SETEX")
.arg(&key)
.arg(ttl_seconds)
.arg(reason)
.query_async::<String>(&mut connection)
.await
} else {
redis::cmd("SET")
.arg(&key)
.arg(reason)
.query_async::<String>(&mut connection)
.await
};
return Ok(result.is_ok());
}
#[cfg(test)]
if let Some(store) = self.admin_security_blacklist_store.as_ref() {
store
.lock()
.expect("admin security blacklist store should lock")
.insert(ip_address.to_string(), reason.to_string());
return Ok(true);
}
Ok(false)
let key = format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{ip_address}");
self.runtime_state
.kv_set(
&key,
reason.to_string(),
ttl_seconds.map(std::time::Duration::from_secs),
)
.await
.map(|_| true)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn remove_admin_security_blacklist(
@@ -53,36 +28,11 @@ impl AppState {
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
if let Some(runner) = self.redis_kv_runner() {
let mut connection = match runner.client().get_multiplexed_async_connection().await {
Ok(value) => value,
Err(_) => return Ok(false),
};
let key = runner
.keyspace()
.key(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{ip_address}"));
let deleted = match redis::cmd("DEL")
.arg(&key)
.query_async::<i64>(&mut connection)
.await
{
Ok(value) => value,
Err(_) => return Ok(false),
};
return Ok(deleted > 0);
}
#[cfg(test)]
if let Some(store) = self.admin_security_blacklist_store.as_ref() {
let removed = store
.lock()
.expect("admin security blacklist store should lock")
.remove(ip_address)
.is_some();
return Ok(removed);
}
Ok(false)
let key = format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}{ip_address}");
self.runtime_state
.kv_delete(&key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn admin_security_blacklist_stats(
@@ -90,48 +40,13 @@ impl AppState {
) -> Result<(bool, usize, Option<String>), GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
if let Some(runner) = self.redis_kv_runner() {
let mut connection = match runner.client().get_multiplexed_async_connection().await {
Ok(value) => value,
Err(_) => return Ok((false, 0, Some("Redis 不可用".to_string()))),
};
let pattern = runner
.keyspace()
.key(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}*"));
let mut cursor = 0u64;
let mut total = 0usize;
loop {
let (next_cursor, keys) = match redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(100)
.query_async::<(u64, Vec<String>)>(&mut connection)
.await
{
Ok(value) => value,
Err(err) => return Ok((false, 0, Some(err.to_string()))),
};
total += keys.len();
if next_cursor == 0 {
break;
}
cursor = next_cursor;
}
return Ok((true, total, None));
}
#[cfg(test)]
if let Some(store) = self.admin_security_blacklist_store.as_ref() {
let total = store
.lock()
.expect("admin security blacklist store should lock")
.len();
return Ok((true, total, None));
}
Ok((false, 0, Some("Redis 不可用".to_string())))
let total = self
.runtime_state
.scan_keys(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}*"), 100)
.await
.map(|keys| keys.len())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok((true, total, None))
}
pub(crate) async fn list_admin_security_blacklist(
@@ -139,83 +54,40 @@ impl AppState {
) -> Result<Vec<AdminSecurityBlacklistEntry>, GatewayError> {
const ADMIN_SECURITY_BLACKLIST_PREFIX: &str = "ip:blacklist:";
if let Some(runner) = self.redis_kv_runner() {
let mut connection = match runner.client().get_multiplexed_async_connection().await {
Ok(value) => value,
Err(_) => return Ok(Vec::new()),
let keys = self
.runtime_state
.scan_keys(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}*"), 100)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let mut entries = Vec::new();
for full_key in keys {
let raw_key = self.runtime_state.strip_namespace(&full_key);
let ip_address = raw_key
.strip_prefix(ADMIN_SECURITY_BLACKLIST_PREFIX)
.unwrap_or(raw_key)
.to_string();
let Some(reason) = self
.runtime_state
.kv_get(raw_key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
else {
continue;
};
let pattern = runner
.keyspace()
.key(&format!("{ADMIN_SECURITY_BLACKLIST_PREFIX}*"));
let prefix = runner.keyspace().key(ADMIN_SECURITY_BLACKLIST_PREFIX);
let mut cursor = 0u64;
let mut entries = Vec::new();
loop {
let (next_cursor, keys) = match redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(100)
.query_async::<(u64, Vec<String>)>(&mut connection)
.await
{
Ok(value) => value,
Err(_) => break,
};
for full_key in keys {
let ip_address = full_key
.strip_prefix(prefix.as_str())
.map(|value| value.to_string())
.unwrap_or_else(|| full_key.clone());
let reason: Result<String, _> = redis::cmd("GET")
.arg(&full_key)
.query_async(&mut connection)
.await;
let reason = match reason {
Ok(value) => value,
Err(_) => continue,
};
let ttl = match redis::cmd("TTL")
.arg(&full_key)
.query_async::<i64>(&mut connection)
.await
{
Ok(value) if value >= 0 => Some(value),
_ => None,
};
entries.push(AdminSecurityBlacklistEntry {
ip_address,
reason,
ttl_seconds: ttl,
});
}
if next_cursor == 0 {
break;
}
cursor = next_cursor;
}
entries.sort_by(|a, b| a.ip_address.cmp(&b.ip_address));
return Ok(entries);
let ttl_seconds = self
.runtime_state
.kv_ttl_seconds(raw_key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
.filter(|ttl| *ttl >= 0);
entries.push(AdminSecurityBlacklistEntry {
ip_address,
reason,
ttl_seconds,
});
}
#[cfg(test)]
if let Some(store) = self.admin_security_blacklist_store.as_ref() {
let mut entries = store
.lock()
.expect("admin security blacklist store should lock")
.iter()
.map(|(ip, reason)| AdminSecurityBlacklistEntry {
ip_address: ip.clone(),
reason: reason.clone(),
ttl_seconds: None,
})
.collect::<Vec<_>>();
entries.sort_by(|a, b| a.ip_address.cmp(&b.ip_address));
return Ok(entries);
}
Ok(Vec::new())
entries.sort_by(|a, b| a.ip_address.cmp(&b.ip_address));
Ok(entries)
}
pub(crate) async fn add_admin_security_whitelist(
@@ -224,34 +96,11 @@ impl AppState {
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
if let Some(runner) = self.redis_kv_runner() {
let mut connection = match runner.client().get_multiplexed_async_connection().await {
Ok(value) => value,
Err(_) => return Ok(false),
};
let key = runner.keyspace().key(ADMIN_SECURITY_WHITELIST_KEY);
let added = match redis::cmd("SADD")
.arg(&key)
.arg(ip_address)
.query_async::<i64>(&mut connection)
.await
{
Ok(value) => value,
Err(_) => return Ok(false),
};
return Ok(added >= 0);
}
#[cfg(test)]
if let Some(store) = self.admin_security_whitelist_store.as_ref() {
store
.lock()
.expect("admin security whitelist store should lock")
.insert(ip_address.to_string());
return Ok(true);
}
Ok(false)
self.runtime_state
.set_add(ADMIN_SECURITY_WHITELIST_KEY, ip_address)
.await
.map(|_| true)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn remove_admin_security_whitelist(
@@ -260,67 +109,18 @@ impl AppState {
) -> Result<bool, GatewayError> {
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
if let Some(runner) = self.redis_kv_runner() {
let mut connection = match runner.client().get_multiplexed_async_connection().await {
Ok(value) => value,
Err(_) => return Ok(false),
};
let key = runner.keyspace().key(ADMIN_SECURITY_WHITELIST_KEY);
let removed = match redis::cmd("SREM")
.arg(&key)
.arg(ip_address)
.query_async::<i64>(&mut connection)
.await
{
Ok(value) => value,
Err(_) => return Ok(false),
};
return Ok(removed > 0);
}
#[cfg(test)]
if let Some(store) = self.admin_security_whitelist_store.as_ref() {
let removed = store
.lock()
.expect("admin security whitelist store should lock")
.remove(ip_address);
return Ok(removed);
}
Ok(false)
self.runtime_state
.set_remove(ADMIN_SECURITY_WHITELIST_KEY, ip_address)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_admin_security_whitelist(&self) -> Result<Vec<String>, GatewayError> {
const ADMIN_SECURITY_WHITELIST_KEY: &str = "ip:whitelist";
if let Some(runner) = self.redis_kv_runner() {
let mut connection = match runner.client().get_multiplexed_async_connection().await {
Ok(value) => value,
Err(_) => return Ok(Vec::new()),
};
let key = runner.keyspace().key(ADMIN_SECURITY_WHITELIST_KEY);
let mut whitelist = match redis::cmd("SMEMBERS")
.arg(&key)
.query_async::<Vec<String>>(&mut connection)
.await
{
Ok(value) => value,
Err(_) => return Ok(Vec::new()),
};
whitelist.sort();
return Ok(whitelist);
}
#[cfg(test)]
if let Some(store) = self.admin_security_whitelist_store.as_ref() {
return Ok(store
.lock()
.expect("admin security whitelist store should lock")
.iter()
.cloned()
.collect());
}
Ok(Vec::new())
self.runtime_state
.set_members(ADMIN_SECURITY_WHITELIST_KEY)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
}

41
apps/aether-gateway/src/state/testing.rs
View File

@@ -1,5 +1,6 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex as StdMutex};
use std::time::Duration;
use aether_contracts::{ExecutionPlan, ExecutionResult};
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
@@ -207,6 +208,13 @@ impl AppState {
.lock()
.expect("provider oauth state store should lock")
.insert(format!("provider_oauth_state:{nonce}"), payload.to_string());
self.runtime_state.kv_set_local_nowait(
&format!("provider_oauth_state:{nonce}"),
payload.to_string(),
Some(Duration::from_secs(
aether_data::repository::provider_oauth::PROVIDER_OAUTH_STATE_TTL_SECS,
)),
);
self
}
@@ -225,6 +233,11 @@ impl AppState {
format!("device_auth_session:{session_id}"),
payload.to_string(),
);
self.runtime_state.kv_set_local_nowait(
&format!("device_auth_session:{session_id}"),
payload.to_string(),
Some(Duration::from_secs(3600)),
);
self
}
@@ -243,6 +256,13 @@ impl AppState {
format!("provider_oauth_batch_task:{task_id}"),
payload.to_string(),
);
self.runtime_state.kv_set_local_nowait(
&format!("provider_oauth_batch_task:{task_id}"),
payload.to_string(),
Some(Duration::from_secs(
aether_data::repository::provider_oauth::PROVIDER_OAUTH_BATCH_TASK_TTL_SECS,
)),
);
self
}
@@ -417,6 +437,11 @@ impl AppState {
.lock()
.expect("admin security blacklist store should lock");
for (ip_address, reason) in entries {
self.runtime_state.kv_set_local_nowait(
&format!("ip:blacklist:{ip_address}"),
reason.clone(),
None,
);
guard.insert(ip_address, reason);
}
drop(guard);
@@ -434,6 +459,8 @@ impl AppState {
.lock()
.expect("admin security whitelist store should lock");
for ip_address in entries {
self.runtime_state
.set_add_local_nowait("ip:whitelist", &ip_address);
guard.insert(ip_address);
}
drop(guard);
@@ -552,6 +579,15 @@ impl AppState {
})
.to_string(),
);
self.runtime_state.kv_set_local_nowait(
&format!("email:verification:{}", email.trim().to_ascii_lowercase()),
json!({
"code": code,
"created_at": created_at.to_rfc3339(),
})
.to_string(),
Some(Duration::from_secs(600)),
);
self
}
@@ -566,6 +602,11 @@ impl AppState {
format!("email:verified:{}", email.trim().to_ascii_lowercase()),
"verified".to_string(),
);
self.runtime_state.kv_set_local_nowait(
&format!("email:verified:{}", email.trim().to_ascii_lowercase()),
"verified".to_string(),
Some(Duration::from_secs(3600)),
);
self
}