feat(security): harden gateway boundaries and usage policies

Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
+1
View File
@@ -11,6 +11,7 @@ async-stream.workspace = true
axum = { version = "0.8" }
chrono.workspace = true
futures-util.workspace = true
libc = "0.2"
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
+53 -19
View File
@@ -2,6 +2,7 @@ use async_stream::stream;
use axum::body::Body;
use axum::http::Response;
use futures_util::StreamExt;
use std::sync::Arc;
use std::time::Duration;
use crate::concurrency::ConcurrencyPermit;
@@ -10,24 +11,31 @@ const ADMISSION_HEALTH_POLL_INTERVAL: Duration = Duration::from_secs(1);
pub trait AdmissionPermitHealth: Send + Sync {
fn is_healthy(&self) -> bool;
fn requires_health_poll(&self) -> bool {
true
}
}
impl AdmissionPermitHealth for ConcurrencyPermit {
fn is_healthy(&self) -> bool {
true
}
fn requires_health_poll(&self) -> bool {
false
}
}
#[derive(Clone)]
pub struct AdmissionPermit {
_local: Option<ConcurrencyPermit>,
_distributed: Option<Box<dyn AdmissionPermitHealth>>,
_permits: Vec<Arc<dyn AdmissionPermitHealth>>,
}
impl std::fmt::Debug for AdmissionPermit {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AdmissionPermit")
.field("has_local", &self._local.is_some())
.field("has_distributed", &self._distributed.is_some())
.field("permit_count", &self._permits.len())
.finish()
}
}
@@ -37,34 +45,39 @@ impl AdmissionPermit {
local: Option<ConcurrencyPermit>,
distributed: Option<D>,
) -> Option<Self> {
if local.is_none() && distributed.is_none() {
None
} else {
Some(Self {
_local: local,
_distributed: distributed
.map(|permit| Box::new(permit) as Box<dyn AdmissionPermitHealth>),
})
let mut permits = Vec::<Arc<dyn AdmissionPermitHealth>>::new();
if let Some(local) = local {
permits.push(Arc::new(local));
}
if let Some(distributed) = distributed {
permits.push(Arc::new(distributed));
}
(!permits.is_empty()).then_some(Self { _permits: permits })
}
pub fn combine(permits: impl IntoIterator<Item = AdmissionPermit>) -> Option<Self> {
let permits = permits
.into_iter()
.flat_map(|permit| permit._permits)
.collect::<Vec<_>>();
(!permits.is_empty()).then_some(Self { _permits: permits })
}
pub fn is_healthy(&self) -> bool {
self._distributed
.as_ref()
.map(|permit| permit.is_healthy())
.unwrap_or(true)
self._permits.iter().all(|permit| permit.is_healthy())
}
fn requires_health_poll(&self) -> bool {
self._distributed.is_some()
self._permits
.iter()
.any(|permit| permit.requires_health_poll())
}
}
impl From<ConcurrencyPermit> for AdmissionPermit {
fn from(value: ConcurrencyPermit) -> Self {
Self {
_local: Some(value),
_distributed: None,
_permits: vec![Arc::new(value)],
}
}
}
@@ -210,6 +223,27 @@ mod tests {
assert_eq!(gate.snapshot().in_flight, 0);
}
#[test]
fn cloned_permit_releases_capacity_only_after_last_clone_drops() {
let gate = ConcurrencyGate::new("test", 1);
let permit = AdmissionPermit::from(gate.try_acquire().expect("first permit"));
let cloned = permit.clone();
drop(permit);
assert_eq!(gate.snapshot().in_flight, 1);
assert!(
gate.try_acquire().is_err(),
"capacity should remain held by the clone"
);
drop(cloned);
assert_eq!(gate.snapshot().in_flight, 0);
assert!(
gate.try_acquire().is_ok(),
"last clone should release capacity"
);
}
#[tokio::test]
async fn holds_combined_local_and_distributed_permit_until_future_finishes() {
let local_gate = ConcurrencyGate::new("local", 1);
+102 -7
View File
@@ -716,10 +716,68 @@ impl RollingFileSink {
}
fn open_bucketed_log_file(dir: &Path, service_name: &str, bucket: &str) -> io::Result<File> {
OpenOptions::new()
.create(true)
.append(true)
.open(bucketed_log_path(dir, service_name, bucket))
let path = bucketed_log_path(dir, service_name, bucket);
let mut options = OpenOptions::new();
options.create(true).append(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt as _;
options
.mode(0o600)
.custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW);
}
let file = options.open(&path)?;
validate_open_log_file(&file, &path)?;
Ok(file)
}
#[cfg(unix)]
fn validate_open_log_file(file: &File, path: &Path) -> io::Result<()> {
use std::os::unix::fs::{MetadataExt as _, PermissionsExt as _};
let metadata = file.metadata()?;
if !metadata.file_type().is_file() {
return Err(io::Error::new(
io::ErrorKind::PermissionDenied,
format!("log destination is not a regular file: {}", path.display()),
));
}
if metadata.uid() != unsafe { libc::geteuid() } {
return Err(io::Error::new(
io::ErrorKind::PermissionDenied,
format!(
"log destination is owned by another user: {}",
path.display()
),
));
}
if metadata.nlink() != 1 {
return Err(io::Error::new(
io::ErrorKind::PermissionDenied,
format!(
"log destination has multiple hard links: {}",
path.display()
),
));
}
// `mode(0o600)` only affects newly-created files. Tighten an existing
// bucket as well so a historical permissive umask cannot keep exposing
// request and operational data after an upgrade.
file.set_permissions(fs::Permissions::from_mode(0o600))?;
Ok(())
}
#[cfg(not(unix))]
fn validate_open_log_file(file: &File, path: &Path) -> io::Result<()> {
let metadata = file.metadata()?;
if !metadata.file_type().is_file() {
return Err(io::Error::new(
io::ErrorKind::PermissionDenied,
format!("log destination is not a regular file: {}", path.display()),
));
}
Ok(())
}
fn bucketed_log_path(dir: &Path, service_name: &str, bucket: &str) -> PathBuf {
@@ -838,9 +896,9 @@ fn select_log_files_for_cleanup(
mod tests {
use super::{
bucketed_log_path, cleanup_log_files, format_target_cell, log_bucket_key,
select_log_files_for_cleanup, FileLoggingConfig, JsonRuntimeEventFormatter,
LogFileCandidate, LogRotation, PrettyRuntimeEventFormatter, RollingFileSink,
RuntimeLogIdentity,
open_bucketed_log_file, select_log_files_for_cleanup, FileLoggingConfig,
JsonRuntimeEventFormatter, LogFileCandidate, LogRotation, PrettyRuntimeEventFormatter,
RollingFileSink, RuntimeLogIdentity,
};
use chrono::{Local, TimeZone};
use std::fs;
@@ -962,6 +1020,43 @@ mod tests {
fs::remove_dir_all(&dir).expect("temp dir should be removable");
}
#[cfg(unix)]
#[test]
fn rolling_log_file_is_private_and_rejects_symlink_destination() {
use std::os::unix::fs::{symlink, PermissionsExt as _};
let dir = std::env::temp_dir().join(format!("aether-runtime-logs-{}", Uuid::new_v4()));
fs::create_dir_all(&dir).expect("temp dir should exist");
let private = open_bucketed_log_file(&dir, "runtime-test", "private")
.expect("private log should open");
assert_eq!(
private
.metadata()
.expect("log metadata")
.permissions()
.mode()
& 0o777,
0o600
);
drop(private);
let victim = dir.join("victim.txt");
fs::write(&victim, b"unchanged").expect("victim should exist");
let symlink_path = bucketed_log_path(&dir, "runtime-test", "symlink");
symlink(&victim, &symlink_path).expect("symlink should exist");
assert!(
open_bucketed_log_file(&dir, "runtime-test", "symlink").is_err(),
"rolling logs must not follow a pre-created symlink"
);
assert_eq!(
fs::read(&victim).expect("victim should remain readable"),
b"unchanged"
);
fs::remove_dir_all(&dir).expect("temp dir should be removable");
}
#[test]
fn rolling_file_sink_treats_startup_cleanup_failure_as_non_fatal() {
fn fail_cleanup(_: &str, _: &FileLoggingConfig) -> std::io::Result<usize> {
File diff suppressed because it is too large Load Diff
+426
View File
@@ -6,20 +6,30 @@ use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use crate::UsageLimitCheck;
use crate::{DataLayerError, RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats};
const MEMORY_RATE_LIMIT_COUNTER_SHARD_COUNT: usize = 64;
const MEMORY_RATE_LIMIT_COUNTER_PRUNE_INTERVAL: u64 = 256;
const MEMORY_USAGE_LIMIT_PRUNE_INTERVAL: u64 = 256;
const DEFAULT_MAX_USAGE_LIMIT_WINDOWS: usize = 10_000;
const DEFAULT_MAX_USAGE_LIMIT_EVENTS: usize = 100_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MemoryRuntimeStateConfig {
pub max_kv_entries: usize,
/// Maximum number of active sliding-window keys retained by the memory backend.
pub max_usage_limit_windows: usize,
/// Maximum number of event identities retained across all usage-limit windows.
pub max_usage_limit_events: usize,
}
impl Default for MemoryRuntimeStateConfig {
fn default() -> Self {
Self {
max_kv_entries: 10_000,
max_usage_limit_windows: DEFAULT_MAX_USAGE_LIMIT_WINDOWS,
max_usage_limit_events: DEFAULT_MAX_USAGE_LIMIT_EVENTS,
}
}
}
@@ -42,6 +52,7 @@ pub(crate) struct MemoryRuntimeBackend {
config: MemoryRuntimeStateConfig,
kv: Mutex<HashMap<String, MemoryKvEntry>>,
counters: MemoryRateLimitCounters,
usage_limits: Mutex<MemoryUsageLimitState>,
sets: Mutex<HashMap<String, MemorySetEntry>>,
scores: Mutex<HashMap<String, MemoryScoreEntry>>,
queues: Mutex<HashMap<String, MemoryQueueStream>>,
@@ -51,6 +62,111 @@ pub(crate) struct MemoryRuntimeBackend {
semaphores: Mutex<HashMap<String, BTreeMap<String, u64>>>,
}
#[derive(Debug, Clone)]
struct MemoryUsageLimitWindow {
window_ms: u64,
expires_at_unix_ms: u64,
events: HashMap<String, u64>,
}
#[derive(Debug, Default)]
struct MemoryUsageLimitState {
windows: HashMap<String, MemoryUsageLimitWindow>,
total_events: usize,
operations_since_prune: u64,
next_expiry_unix_ms: Option<u64>,
}
impl MemoryUsageLimitState {
fn amortized_prune(&mut self, now_unix_ms: u64) {
self.operations_since_prune = self.operations_since_prune.saturating_add(1);
if self.operations_since_prune < MEMORY_USAGE_LIMIT_PRUNE_INTERVAL {
return;
}
self.operations_since_prune = 0;
if self
.next_expiry_unix_ms
.is_some_and(|expires_at| expires_at <= now_unix_ms)
{
self.prune_all(now_unix_ms);
}
}
fn prune_all(&mut self, now_unix_ms: u64) {
self.operations_since_prune = 0;
let mut total_events = 0_usize;
let mut next_expiry_unix_ms = None;
self.windows.retain(|_, window| {
if window.expires_at_unix_ms <= now_unix_ms {
return false;
}
prune_usage_limit_events(&mut window.events, now_unix_ms, window.window_ms);
if window.events.is_empty() {
return false;
}
total_events = total_events.saturating_add(window.events.len());
update_earliest_expiry(&mut next_expiry_unix_ms, window.expires_at_unix_ms);
for timestamp in window.events.values() {
update_earliest_expiry(
&mut next_expiry_unix_ms,
timestamp.saturating_add(window.window_ms),
);
}
true
});
self.total_events = total_events;
self.next_expiry_unix_ms = next_expiry_unix_ms;
}
fn prune_rule_window(&mut self, key: &str, now_unix_ms: u64, window_ms: u64) {
if self
.windows
.get(key)
.is_some_and(|window| window.expires_at_unix_ms <= now_unix_ms)
{
if let Some(window) = self.windows.remove(key) {
self.total_events = self.total_events.saturating_sub(window.events.len());
}
return;
}
let Some(window) = self.windows.get_mut(key) else {
return;
};
let before = window.events.len();
window.window_ms = window_ms;
prune_usage_limit_events(&mut window.events, now_unix_ms, window_ms);
self.total_events = self
.total_events
.saturating_sub(before.saturating_sub(window.events.len()));
update_earliest_expiry(&mut self.next_expiry_unix_ms, window.expires_at_unix_ms);
for timestamp in window.events.values() {
update_earliest_expiry(
&mut self.next_expiry_unix_ms,
timestamp.saturating_add(window_ms),
);
}
if window.events.is_empty() {
self.windows.remove(key);
}
}
fn additions_for(&self, input: crate::UsageLimitInput<'_>) -> (usize, usize) {
input.rules.iter().fold(
(0_usize, 0_usize),
|(additional_windows, additional_events), rule| match self.windows.get(rule.key) {
Some(window) if window.events.contains_key(input.event_id) => {
(additional_windows, additional_events)
}
Some(_) => (additional_windows, additional_events.saturating_add(1)),
None => (
additional_windows.saturating_add(1),
additional_events.saturating_add(1),
),
},
)
}
}
#[derive(Debug, Clone)]
struct MemoryCounterEntry {
value: u32,
@@ -206,6 +322,34 @@ impl MemoryRuntimeBackend {
);
}
pub(crate) async fn kv_set_if_absent(&self, key: &str, value: String, ttl: Duration) -> bool {
let mut kv = self.kv.lock().await;
let now = Instant::now();
prune_kv(&mut kv, now);
if kv.contains_key(key) {
return false;
}
while kv.len() >= self.config.max_kv_entries.max(1) {
let Some(oldest_key) = kv
.iter()
.min_by_key(|(_, entry)| entry.inserted_at)
.map(|(key, _)| key.clone())
else {
break;
};
kv.remove(&oldest_key);
}
kv.insert(
key.to_string(),
MemoryKvEntry {
value,
inserted_at: now,
expires_at: Some(now + ttl),
},
);
true
}
pub(crate) fn kv_set_nowait(&self, key: &str, value: String, ttl: Option<Duration>) -> bool {
let Ok(mut kv) = self.kv.try_lock() else {
return false;
@@ -502,6 +646,143 @@ impl MemoryRuntimeBackend {
})
}
pub(crate) async fn check_and_consume_usage_limits(
&self,
input: crate::UsageLimitInput<'_>,
) -> Result<UsageLimitCheck, DataLayerError> {
let mut state = self.usage_limits.lock().await;
state.amortized_prune(input.now_unix_ms);
for (index, rule) in input.rules.iter().enumerate() {
let window_ms = rule.window_seconds.saturating_mul(1_000);
state.prune_rule_window(rule.key, input.now_unix_ms, window_ms);
let Some(window) = state.windows.get(rule.key) else {
continue;
};
if window.events.contains_key(input.event_id) {
continue;
}
if window.events.len() as u64 >= rule.limit {
let earliest = window
.events
.values()
.copied()
.min()
.unwrap_or(input.now_unix_ms);
let retry_after_ms = earliest
.saturating_add(window_ms)
.saturating_sub(input.now_unix_ms);
return Ok(UsageLimitCheck::Rejected {
rule_index: index,
limit: rule.limit,
retry_after: retry_after_ms.saturating_add(999) / 1_000,
});
}
}
let (mut additional_windows, mut additional_events) = state.additions_for(input);
if state.windows.len().saturating_add(additional_windows)
> self.config.max_usage_limit_windows
|| state.total_events.saturating_add(additional_events)
> self.config.max_usage_limit_events
{
// Redis drops an idle sorted-set key after its retention TTL. Force the equivalent full
// cleanup before rejecting capacity so stale high-cardinality keys cannot pin memory.
if state
.next_expiry_unix_ms
.is_some_and(|expires_at| expires_at <= input.now_unix_ms)
{
state.prune_all(input.now_unix_ms);
}
(additional_windows, additional_events) = state.additions_for(input);
}
if state.windows.len().saturating_add(additional_windows)
> self.config.max_usage_limit_windows
|| state.total_events.saturating_add(additional_events)
> self.config.max_usage_limit_events
{
return Err(DataLayerError::UnexpectedValue(format!(
"runtime memory usage-limit capacity exhausted (windows {}/{}, events {}/{})",
state.windows.len(),
self.config.max_usage_limit_windows,
state.total_events,
self.config.max_usage_limit_events,
)));
}
for rule in input.rules {
let window_ms = rule.window_seconds.saturating_mul(1_000);
let expires_at_unix_ms = input
.now_unix_ms
.saturating_add(rule.retention_seconds.saturating_mul(1_000));
let inserted = match state.windows.entry(rule.key.to_string()) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
let window = entry.get_mut();
window.window_ms = window_ms;
window.expires_at_unix_ms = expires_at_unix_ms;
match window.events.entry(input.event_id.to_string()) {
std::collections::hash_map::Entry::Occupied(_) => false,
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(input.now_unix_ms);
true
}
}
}
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(MemoryUsageLimitWindow {
window_ms,
expires_at_unix_ms,
events: HashMap::from([(input.event_id.to_string(), input.now_unix_ms)]),
});
true
}
};
update_earliest_expiry(&mut state.next_expiry_unix_ms, expires_at_unix_ms);
if inserted {
state.total_events = state.total_events.saturating_add(1);
update_earliest_expiry(
&mut state.next_expiry_unix_ms,
input.now_unix_ms.saturating_add(window_ms),
);
}
}
Ok(UsageLimitCheck::Allowed)
}
pub(crate) async fn release_usage_limits(
&self,
input: crate::UsageLimitReleaseInput<'_>,
) -> Result<(), DataLayerError> {
let mut state = self.usage_limits.lock().await;
for rule in input.rules {
let mut remove_window = false;
let mut removed_event = false;
if let Some(window) = state.windows.get_mut(rule.key) {
removed_event = window.events.remove(input.event_id).is_some();
remove_window = window.events.is_empty();
}
if removed_event {
state.total_events = state.total_events.saturating_sub(1);
}
if remove_window {
state.windows.remove(rule.key);
}
}
state.next_expiry_unix_ms = state
.windows
.values()
.flat_map(|window| {
std::iter::once(window.expires_at_unix_ms).chain(
window
.events
.values()
.map(|timestamp| timestamp.saturating_add(window.window_ms)),
)
})
.min();
Ok(())
}
pub(crate) fn rate_limit_count(&self, key: &str, bucket: u64) -> Result<u32, DataLayerError> {
let now = Instant::now();
let mut total = 0_u32;
@@ -1013,6 +1294,17 @@ impl MemoryRuntimeBackend {
}
}
fn prune_usage_limit_events(events: &mut HashMap<String, u64>, now_unix_ms: u64, window_ms: u64) {
let Some(cutoff) = now_unix_ms.checked_sub(window_ms) else {
return;
};
events.retain(|_, timestamp| *timestamp > cutoff);
}
fn update_earliest_expiry(current: &mut Option<u64>, candidate: u64) {
*current = Some(current.map_or(candidate, |existing| existing.min(candidate)));
}
fn get_fresh_locked(
kv: &mut HashMap<String, MemoryKvEntry>,
key: &str,
@@ -1197,4 +1489,138 @@ mod tests {
assert!(!shard.entries.contains_key("expired-unrelated-key"));
assert_eq!(shard.operations_since_prune, 0);
}
#[tokio::test]
async fn usage_limit_capacity_is_atomic_and_fail_closed() {
let backend = MemoryRuntimeBackend::new(MemoryRuntimeStateConfig {
max_usage_limit_windows: 2,
max_usage_limit_events: 2,
..MemoryRuntimeStateConfig::default()
});
let first = [crate::UsageLimitRule {
key: "usage:{user-1}:one",
limit: 10,
window_seconds: 60,
retention_seconds: 60,
}];
backend
.check_and_consume_usage_limits(crate::UsageLimitInput {
rules: &first,
event_id: "event-1",
now_unix_ms: 1_000,
})
.await
.expect("first event");
let two_new_windows = [
crate::UsageLimitRule {
key: "usage:{user-1}:two",
limit: 10,
window_seconds: 60,
retention_seconds: 60,
},
crate::UsageLimitRule {
key: "usage:{user-1}:three",
limit: 10,
window_seconds: 60,
retention_seconds: 60,
},
];
let error = backend
.check_and_consume_usage_limits(crate::UsageLimitInput {
rules: &two_new_windows,
event_id: "event-2",
now_unix_ms: 2_000,
})
.await
.expect_err("capacity must fail closed");
assert!(error.to_string().contains("capacity exhausted"));
let state = backend.usage_limits.lock().await;
assert_eq!(state.windows.len(), 1);
assert_eq!(state.total_events, 1);
assert!(!state.windows.contains_key(two_new_windows[0].key));
assert!(!state.windows.contains_key(two_new_windows[1].key));
}
#[tokio::test]
async fn usage_limit_capacity_reclaims_expired_windows_before_rejecting() {
let backend = MemoryRuntimeBackend::new(MemoryRuntimeStateConfig {
max_usage_limit_windows: 1,
max_usage_limit_events: 1,
..MemoryRuntimeStateConfig::default()
});
let old = [crate::UsageLimitRule {
key: "usage:{user-1}:old",
limit: 1,
window_seconds: 1,
retention_seconds: 1,
}];
backend
.check_and_consume_usage_limits(crate::UsageLimitInput {
rules: &old,
event_id: "event-old",
now_unix_ms: 1_000,
})
.await
.expect("old event");
let current = [crate::UsageLimitRule {
key: "usage:{user-1}:current",
limit: 1,
window_seconds: 1,
retention_seconds: 1,
}];
assert_eq!(
backend
.check_and_consume_usage_limits(crate::UsageLimitInput {
rules: &current,
event_id: "event-current",
now_unix_ms: 2_000,
})
.await
.expect("expired capacity should be reclaimed"),
UsageLimitCheck::Allowed
);
let state = backend.usage_limits.lock().await;
assert_eq!(state.windows.len(), 1);
assert_eq!(state.total_events, 1);
assert!(state.windows.contains_key(current[0].key));
}
#[tokio::test]
async fn usage_limit_idempotent_replay_does_not_consume_event_capacity() {
let backend = MemoryRuntimeBackend::new(MemoryRuntimeStateConfig {
max_usage_limit_windows: 1,
max_usage_limit_events: 1,
..MemoryRuntimeStateConfig::default()
});
let rules = [crate::UsageLimitRule {
key: "usage:{user-1}:idempotent",
limit: 10,
window_seconds: 60,
retention_seconds: 60,
}];
for now_unix_ms in [1_000, 2_000] {
assert_eq!(
backend
.check_and_consume_usage_limits(crate::UsageLimitInput {
rules: &rules,
event_id: "same-event",
now_unix_ms,
})
.await
.expect("idempotent replay"),
UsageLimitCheck::Allowed
);
}
let state = backend.usage_limits.lock().await;
assert_eq!(state.total_events, 1);
assert_eq!(
state.windows[rules[0].key].events["same-event"], 1_000,
"idempotent replay must preserve the original Redis ZADD NX timestamp"
);
}
}
@@ -17,12 +17,46 @@ pub(crate) const REDIS_COMMAND_LATENCY_BUCKETS_MS: [u64; 12] =
[1, 5, 10, 25, 50, 100, 250, 500, 1_000, 2_500, 5_000, 10_000];
const REDIS_COMMAND_LATENCY_BUCKET_COUNT: usize = REDIS_COMMAND_LATENCY_BUCKETS_MS.len() + 1;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
#[derive(Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct RedisClientConfig {
pub url: String,
pub key_prefix: Option<String>,
}
impl std::fmt::Debug for RedisClientConfig {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("RedisClientConfig")
.field("url", &redact_redis_url_for_debug(&self.url))
.field("key_prefix_len", &self.key_prefix.as_ref().map(String::len))
.finish()
}
}
fn redact_redis_url_for_debug(raw: &str) -> String {
const MAX_DEBUG_URL_CHARS: usize = 512;
let raw = raw.trim();
let Ok(mut url) = url::Url::parse(raw) else {
return format!("[invalid-redis-url len={}]", raw.len());
};
let _ = url.set_username("");
let _ = url.set_password(None);
url.set_query(None);
url.set_fragment(None);
let rendered = url.to_string();
if rendered.chars().count() <= MAX_DEBUG_URL_CHARS {
rendered
} else {
format!(
"{}...",
rendered
.chars()
.take(MAX_DEBUG_URL_CHARS.saturating_sub(3))
.collect::<String>()
)
}
}
impl RedisClientConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
let raw = self.url.trim();
@@ -466,6 +500,25 @@ mod tests {
.expect("lazy redis client should build");
}
#[test]
fn redis_config_debug_redacts_url_credentials_and_query() {
let config = RedisClientConfig {
url: "redis://redis-user:redis-password@redis.example/0?token=redis-secret".into(),
key_prefix: Some("tenant-secret".into()),
};
let debug = format!("{config:?}");
for secret in [
"redis-user",
"redis-password",
"redis-secret",
"tenant-secret",
] {
assert!(!debug.contains(secret), "debug leaked {secret}: {debug}");
}
assert!(debug.contains("redis.example"));
assert!(debug.contains("key_prefix_len"));
}
#[test]
fn blocking_stream_lane_count_uses_requested_as_floor() {
let default_lanes = default_blocking_stream_lane_count();
@@ -7,6 +7,7 @@ use crate::redis::{
};
use crate::{
DataLayerError, RateLimitCheck, RateLimitInput, RateLimitScope, RuntimeSemaphoreError,
UsageLimitCheck, UsageLimitInput, UsageLimitReleaseInput,
};
const RATE_LIMIT_CHECK_AND_CONSUME_SCRIPT: &str = r#"
@@ -51,6 +52,50 @@ end
return {1, 0, 0, remaining}
"#;
const USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT: &str = r#"
local count = #KEYS
local now = tonumber(ARGV[1])
local event_id = ARGV[2]
for i = 1, count do
local window_ms = tonumber(ARGV[(i - 1) * 3 + 4]) * 1000
local cutoff = now - window_ms
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', cutoff)
end
for i = 1, count do
local limit = tonumber(ARGV[(i - 1) * 3 + 3])
local window_ms = tonumber(ARGV[(i - 1) * 3 + 4]) * 1000
local already_consumed = redis.call('ZSCORE', KEYS[i], event_id)
if not already_consumed then
local current = redis.call('ZCARD', KEYS[i])
if current >= limit then
local earliest = redis.call('ZRANGE', KEYS[i], 0, 0, 'WITHSCORES')
local retry_after = 1
if #earliest >= 2 then
local retry_after_ms = math.max(1, tonumber(earliest[2]) + window_ms - now)
retry_after = math.ceil(retry_after_ms / 1000)
end
return {0, i, limit, retry_after}
end
end
end
for i = 1, count do
local retention = tonumber(ARGV[(i - 1) * 3 + 5])
redis.call('ZADD', KEYS[i], 'NX', now, event_id)
redis.call('EXPIRE', KEYS[i], retention + 1)
end
return {1, 0, 0, 0}
"#;
const USAGE_LIMIT_RELEASE_SCRIPT: &str = r#"
for i = 1, #KEYS do
redis.call('ZREM', KEYS[i], ARGV[1])
end
return 1
"#;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RedisRuntimeDiagnostics {
pub connected_clients: Option<u64>,
@@ -147,6 +192,37 @@ impl RedisRuntimeRunner {
Ok(())
}
pub(crate) async fn kv_set_if_absent(
&self,
key: &str,
value: String,
ttl: Duration,
) -> Result<bool, DataLayerError> {
let namespaced_key = self.keyspace.key(key);
let ttl_ms = u64::try_from(ttl.as_millis().max(1)).unwrap_or(u64::MAX);
let mut command = cmd("SET");
command
.arg(namespaced_key)
.arg(value)
.arg("NX")
.arg("PX")
.arg(ttl_ms);
let response = self
.query::<Option<String>>(
RedisConnectionLane::Fast,
"runtime kv set if absent",
command,
)
.await?;
match response.as_deref() {
Some("OK") => Ok(true),
None => Ok(false),
Some(value) => Err(DataLayerError::UnexpectedValue(format!(
"unexpected runtime kv set if absent response {value}"
))),
}
}
pub(crate) async fn kv_get_many(
&self,
keys: &[String],
@@ -282,6 +358,105 @@ impl RedisRuntimeRunner {
Ok(RateLimitCheck::Rejected { scope, limit })
}
pub(crate) async fn check_and_consume_usage_limits(
&self,
input: UsageLimitInput<'_>,
) -> Result<UsageLimitCheck, DataLayerError> {
let script = script(USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT);
let mut invocation = script.prepare_invoke();
for rule in input.rules {
invocation.key(self.keyspace.key(rule.key));
}
invocation.arg(input.now_unix_ms as i64);
invocation.arg(input.event_id);
for rule in input.rules {
invocation.arg(rule.limit as i64);
invocation.arg(rule.window_seconds as i64);
invocation.arg(rule.retention_seconds as i64);
}
let raw = run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
self.command_timeout_ms,
"runtime usage limit check",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
invocation
.invoke_async::<Vec<i64>>(&mut connection)
.await
.map_redis_err()
},
)
.await?;
match raw.first().copied() {
Some(1) if raw.len() >= 4 => return Ok(UsageLimitCheck::Allowed),
Some(0) if raw.len() >= 4 => {}
_ => {
return Err(DataLayerError::UnexpectedValue(
"runtime usage limit script returned an invalid response".to_string(),
));
}
}
let rule_index = raw
.get(1)
.copied()
.and_then(|value| usize::try_from(value.saturating_sub(1)).ok())
.filter(|index| *index < input.rules.len())
.ok_or_else(|| {
DataLayerError::UnexpectedValue(
"runtime usage limit script returned an invalid rule index".to_string(),
)
})?;
let limit = raw
.get(2)
.copied()
.and_then(|value| u64::try_from(value).ok())
.filter(|limit| *limit > 0)
.ok_or_else(|| {
DataLayerError::UnexpectedValue(
"runtime usage limit script returned an invalid limit".to_string(),
)
})?;
let retry_after = raw
.get(3)
.copied()
.and_then(|value| u64::try_from(value).ok())
.unwrap_or(1)
.max(1);
Ok(UsageLimitCheck::Rejected {
rule_index,
limit,
retry_after,
})
}
pub(crate) async fn release_usage_limits(
&self,
input: UsageLimitReleaseInput<'_>,
) -> Result<(), DataLayerError> {
let script = script(USAGE_LIMIT_RELEASE_SCRIPT);
let mut invocation = script.prepare_invoke();
for rule in input.rules {
invocation.key(self.keyspace.key(rule.key));
}
invocation.arg(input.event_id);
run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
self.command_timeout_ms,
"runtime usage limit release",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
invocation
.invoke_async::<i64>(&mut connection)
.await
.map_redis_err()
},
)
.await?;
Ok(())
}
pub(crate) async fn set_add(&self, key: &str, member: &str) -> Result<bool, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("SADD");