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
@@ -7,11 +7,11 @@ use std::time::{Duration, Instant};
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_gateway::tunnel_protocol as protocol;
use aether_testkit::{
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for, run_http_load_probe,
BenchmarkRuntimeSnapshot, ExecutionRuntimeHarness, ExecutionRuntimeHarnessConfig,
GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
HttpLoadProbeResult, SpawnedServer, TunnelHarness, TunnelHarnessConfig,
GATEWAY_HARNESS_API_KEY,
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for,
insert_tunnel_harness_auth_headers, run_http_load_probe, BenchmarkRuntimeSnapshot,
ExecutionRuntimeHarness, ExecutionRuntimeHarnessConfig, GatewayHarness, GatewayHarnessConfig,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, SpawnedServer,
TunnelHarness, TunnelHarnessConfig, GATEWAY_HARNESS_API_KEY, TUNNEL_HARNESS_NODE_ID,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::StatusCode;
@@ -308,6 +308,7 @@ async fn run_tunnel_curve(
for limit in &config.points {
let relay_concurrency = (*limit).saturating_sub(1).max(1);
let tunnel = TunnelHarness::start(TunnelHarnessConfig {
node_id: TUNNEL_HARNESS_NODE_ID.to_string(),
max_streams: (*limit).max(128),
ping_interval: Duration::from_secs(15),
outbound_queue_capacity: 128,
@@ -657,9 +658,7 @@ async fn connect_protocol_peer(
);
let request = ws_url.into_client_request()?;
let mut request = request;
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static("node-baseline"));
insert_tunnel_harness_auth_headers(request.headers_mut(), TUNNEL_HARNESS_NODE_ID)?;
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_static(
@@ -11,8 +11,9 @@ use aether_data::driver::postgres::{
use aether_data::{DataLayerError, PostgresBackend};
use aether_runtime_state::{RedisClientConfig, RedisLockRunner, RedisLockRunnerConfig};
use aether_testkit::{
init_test_runtime_for, reserve_local_port, BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot,
ManagedPostgresServer, ManagedRedisServer, TunnelHarness, TunnelHarnessConfig,
init_test_runtime_for, insert_tunnel_harness_auth_headers, reserve_local_port,
BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot, ManagedPostgresServer, ManagedRedisServer,
TunnelHarness, TunnelHarnessConfig, TUNNEL_HARNESS_NODE_ID,
};
use futures_util::{FutureExt, StreamExt};
use serde::Serialize;
@@ -508,12 +509,8 @@ async fn benchmark_tunnel_restart_recovery(
.clone()
.into_client_request()
.map_err(|err| format!("failed to build websocket request: {err}"))?;
request.headers_mut().insert(
"x-node-id",
format!("recovery-node-{worker_index}-{current}")
.parse()
.map_err(|err| format!("failed to build x-node-id header: {err}"))?,
);
insert_tunnel_harness_auth_headers(request.headers_mut(), TUNNEL_HARNESS_NODE_ID)
.map_err(|err| format!("failed to build tunnel auth headers: {err}"))?;
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION_STR
@@ -1,7 +1,11 @@
// Gateway-backed benchmark scenarios live outside the reusable testkit.
use std::env;
#[cfg(unix)]
use std::fs;
use std::path::PathBuf;
use std::io;
#[cfg(unix)]
use std::io::Write;
use std::path::{Path, PathBuf};
use aether_data::repository::auth::CreateStandaloneApiKeyRecord;
use aether_data::repository::wallet::WalletLookupKey;
@@ -522,7 +526,11 @@ async fn seed_api_key(
.update_standalone_api_key_basic(
aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord {
api_key_id: api_key_id.clone(),
key_encrypted: None,
key_encrypted_present: false,
name: Some(format!("Local pressure API key {}", key_index + 1)),
name_present: true,
force_capabilities: None,
rate_limit_present: true,
rate_limit: Some(0),
concurrent_limit_present: true,
@@ -655,22 +663,18 @@ async fn verify_candidate_selection(
}
fn write_outputs(config: &Config) -> Result<(), Box<dyn std::error::Error>> {
if let Some(parent) = config.output_env_path.parent() {
fs::create_dir_all(parent)?;
}
if let Some(parent) = config.output_key_path.parent() {
fs::create_dir_all(parent)?;
}
if let Some(parent) = config.output_key_list_path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(&config.output_key_path, format!("{}\n", config.api_key))?;
write_private_output(
&config.output_key_path,
format!("{}\n", config.api_key).as_bytes(),
)?;
let key_list = (0..config.api_key_count)
.map(|index| pressure_api_key_value(config, index))
.collect::<Vec<_>>()
.join("\n");
fs::write(&config.output_key_list_path, format!("{key_list}\n"))?;
write_private_output(
&config.output_key_list_path,
format!("{key_list}\n").as_bytes(),
)?;
let env_content = format!(
concat!(
"export AETHER_API_KEY_FILE={key_path}\n",
@@ -689,11 +693,110 @@ fn write_outputs(config: &Config) -> Result<(), Box<dyn std::error::Error>> {
model = shell_escape(&config.model),
mock_upstream_base_url = shell_escape(&config.mock_upstream_base_url),
);
fs::write(&config.output_env_path, env_content)?;
write_private_output(&config.output_env_path, env_content.as_bytes())?;
Ok(())
}
fn write_private_output(path: &Path, contents: &[u8]) -> io::Result<()> {
#[cfg(not(unix))]
{
let _ = (path, contents);
return Err(io::Error::new(
io::ErrorKind::Unsupported,
"private benchmark credential outputs currently require Unix filesystem checks",
));
}
#[cfg(unix)]
{
use std::os::unix::fs::{MetadataExt, OpenOptionsExt, PermissionsExt};
let file_name = path.file_name().ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"private output path must name a file",
)
})?;
let input_parent = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."));
let parent = fs::canonicalize(input_parent)?;
let parent_metadata = fs::symlink_metadata(&parent)?;
if !parent_metadata.is_dir() || parent_metadata.file_type().is_symlink() {
return Err(io::Error::other(
"private output parent must be a real directory",
));
}
let target = parent.join(file_name);
let temporary = parent.join(format!(
".aether-pressure-output-{}.tmp",
uuid::Uuid::new_v4()
));
let mut file = fs::OpenOptions::new()
.write(true)
.create_new(true)
.mode(0o600)
.open(&temporary)?;
let result = (|| -> io::Result<()> {
let owner_uid = file.metadata()?.uid();
validate_private_output_directory(&parent, owner_uid)?;
match fs::symlink_metadata(&target) {
Ok(metadata)
if metadata.is_file()
&& !metadata.file_type().is_symlink()
&& metadata.uid() == owner_uid
&& metadata.nlink() == 1 => {}
Ok(_) => {
return Err(io::Error::other(
"refusing to replace a symlink, special file, hard link, or foreign-owned private output",
));
}
Err(error) if error.kind() == io::ErrorKind::NotFound => {}
Err(error) => return Err(error),
}
file.set_permissions(fs::Permissions::from_mode(0o600))?;
file.write_all(contents)?;
file.sync_all()?;
drop(file);
fs::rename(&temporary, &target)?;
fs::File::open(&parent)?.sync_all()
})();
if result.is_err() {
let _ = fs::remove_file(&temporary);
}
result
}
}
#[cfg(unix)]
fn validate_private_output_directory(directory: &Path, owner_uid: u32) -> io::Result<()> {
use std::os::unix::fs::MetadataExt;
let mut ancestor = Some(directory);
while let Some(path) = ancestor {
let metadata = fs::symlink_metadata(path)?;
let mode = metadata.mode();
if !metadata.is_dir()
|| metadata.file_type().is_symlink()
|| (metadata.uid() != owner_uid && metadata.uid() != 0)
|| (mode & 0o022 != 0 && mode & 0o1000 == 0)
{
return Err(io::Error::other(format!(
"private output directory '{}' has unsafe ownership or permissions",
path.display()
)));
}
ancestor = path.parent();
}
Ok(())
}
fn sha256_hex(value: &str) -> String {
let mut hasher = sha2::Sha256::new();
hasher.update(value.as_bytes());
@@ -789,7 +892,7 @@ Options:\n\
mod tests {
use serde_json::json;
use super::pressure_provider_transport_config;
use super::{pressure_provider_transport_config, write_private_output};
#[test]
fn pressure_provider_transport_config_enables_h2c_prior_knowledge() {
@@ -811,4 +914,38 @@ mod tests {
fn pressure_provider_transport_config_is_absent_by_default() {
assert_eq!(pressure_provider_transport_config(false), None);
}
#[cfg(unix)]
#[test]
fn private_outputs_are_atomic_private_and_refuse_link_targets() {
use std::os::unix::fs::{symlink, MetadataExt, PermissionsExt};
let root = std::env::temp_dir().join(format!(
"aether-pressure-private-output-test-{}",
uuid::Uuid::new_v4()
));
std::fs::create_dir(&root).unwrap();
std::fs::set_permissions(&root, std::fs::Permissions::from_mode(0o700)).unwrap();
let output = root.join("api-key");
write_private_output(&output, b"secret\n").unwrap();
let metadata = std::fs::symlink_metadata(&output).unwrap();
assert_eq!(std::fs::read(&output).unwrap(), b"secret\n");
assert_eq!(metadata.mode() & 0o777, 0o600);
assert_eq!(metadata.nlink(), 1);
let victim = root.join("victim");
std::fs::write(&victim, b"known-good").unwrap();
std::fs::remove_file(&output).unwrap();
symlink(&victim, &output).unwrap();
assert!(write_private_output(&output, b"replacement\n").is_err());
assert_eq!(std::fs::read(&victim).unwrap(), b"known-good");
std::fs::remove_file(&output).unwrap();
std::fs::hard_link(&victim, &output).unwrap();
assert!(write_private_output(&output, b"replacement\n").is_err());
assert_eq!(std::fs::read(&victim).unwrap(), b"known-good");
std::fs::remove_dir_all(root).unwrap();
}
}
@@ -5,9 +5,10 @@ use std::time::Duration;
use aether_gateway::tunnel_protocol as protocol;
use aether_testkit::{
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for, run_http_load_probe,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, PrometheusSample,
TunnelHarness, TunnelHarnessConfig,
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for,
insert_tunnel_harness_auth_headers, run_http_load_probe, HttpLoadProbeConfig,
HttpLoadProbeResponseMode, HttpLoadProbeResult, PrometheusSample, TunnelHarness,
TunnelHarnessConfig, TUNNEL_HARNESS_NODE_ID,
};
use futures_util::{SinkExt, StreamExt};
use reqwest::Method;
@@ -242,9 +243,7 @@ async fn connect_protocol_peer(
);
let request = ws_url.into_client_request()?;
let mut request = request;
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static("node-baseline"));
insert_tunnel_harness_auth_headers(request.headers_mut(), TUNNEL_HARNESS_NODE_ID)?;
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_static(
@@ -8,7 +8,8 @@ use std::time::{Duration, Instant};
use aether_gateway::tunnel_protocol as protocol;
use aether_testkit::{
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for,
BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot, TunnelHarness, TunnelHarnessConfig,
insert_tunnel_harness_auth_headers, BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot,
TunnelHarness, TunnelHarnessConfig,
};
use futures_util::stream::SplitSink;
use futures_util::{SinkExt, StreamExt};
@@ -341,6 +342,7 @@ async fn run_suite(
let chunks_per_stream = config.effective_chunks_per_stream();
let config = Arc::new(config);
let tunnel = TunnelHarness::start(TunnelHarnessConfig {
node_id: config.node_id.clone(),
max_streams: config.tunnel_max_streams,
ping_interval: config.ping_interval,
outbound_queue_capacity: config.outbound_queue_capacity,
@@ -633,10 +635,7 @@ async fn connect_protocol_peer(
);
let request = ws_url.into_client_request()?;
let mut request = request;
request.headers_mut().insert(
"x-node-id",
http::HeaderValue::from_str(config.node_id.as_str())?,
);
insert_tunnel_harness_auth_headers(request.headers_mut(), config.node_id.as_str())?;
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_static(
@@ -9,11 +9,11 @@ use aether_runtime_state::{
RedisClientConfig, RuntimeSemaphore, RuntimeSemaphoreConfig, RuntimeState,
};
use aether_testkit::{
init_test_runtime_for, run_multi_url_http_load_probe, BenchmarkRuntimeSampler,
BenchmarkRuntimeSnapshot, ExecutionRuntimeHarness, ExecutionRuntimeHarnessConfig,
GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
ManagedRedisServer, MultiUrlHttpLoadProbeResult, SpawnedServer, TunnelHarness,
TunnelHarnessConfig, GATEWAY_HARNESS_API_KEY,
init_test_runtime_for, insert_tunnel_harness_auth_headers, run_multi_url_http_load_probe,
BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot, ExecutionRuntimeHarness,
ExecutionRuntimeHarnessConfig, GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig,
HttpLoadProbeResponseMode, ManagedRedisServer, MultiUrlHttpLoadProbeResult, SpawnedServer,
TunnelHarness, TunnelHarnessConfig, GATEWAY_HARNESS_API_KEY, TUNNEL_HARNESS_NODE_ID,
};
use axum::body::to_bytes;
use axum::extract::Request;
@@ -493,12 +493,8 @@ async fn run_tunnel_proxy_connection_probe(
let mut request = url
.into_client_request()
.map_err(|err| format!("failed to build websocket request: {err}"))?;
request.headers_mut().insert(
"x-node-id",
format!("baseline-node-{worker_index}-{current}")
.parse()
.map_err(|err| format!("failed to build x-node-id header: {err}"))?,
);
insert_tunnel_harness_auth_headers(request.headers_mut(), TUNNEL_HARNESS_NODE_ID)
.map_err(|err| format!("failed to build tunnel auth headers: {err}"))?;
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION_STR
@@ -5,13 +5,15 @@ use std::time::Duration;
use aether_gateway::tunnel_protocol as protocol;
use aether_gateway::GatewayDataConfig;
use aether_testkit::{
init_test_runtime_for, prepare_aether_postgres_schema, reserve_local_port, run_http_load_probe,
wait_until, GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig,
HttpLoadProbeResponseMode, HttpLoadProbeResult, ManagedPostgresServer, ManagedRedisServer,
init_test_runtime_for, insert_tunnel_harness_auth_headers, prepare_aether_postgres_schema,
reserve_local_port, run_http_load_probe, wait_until, GatewayHarness, GatewayHarnessConfig,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, ManagedPostgresServer,
ManagedRedisServer, TUNNEL_HARNESS_GENERATION, TUNNEL_HARNESS_MANAGEMENT_TOKEN,
};
use futures_util::{SinkExt, StreamExt};
use reqwest::Method;
use serde::Serialize;
use sha2::{Digest, Sha256};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message;
@@ -114,6 +116,7 @@ async fn run_suite(
.expect("postgres url should be resolved");
prepare_aether_postgres_schema(&postgres_url).await?;
seed_tunnel_auth(&postgres_url).await?;
let key_prefix = format!("aether-owner-relay-baseline-{}", std::process::id());
let shared_data = GatewayDataConfig::from_postgres_url(postgres_url.clone(), false)
@@ -287,9 +290,7 @@ async fn connect_protocol_peer(
);
let request = ws_url.into_client_request()?;
let mut request = request;
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static(NODE_ID));
insert_tunnel_harness_auth_headers(request.headers_mut(), NODE_ID)?;
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_static(
@@ -333,6 +334,98 @@ async fn connect_protocol_peer(
}))
}
async fn seed_tunnel_auth(postgres_url: &str) -> Result<(), Box<dyn std::error::Error>> {
const USER_ID: &str = "user-owner-relay-baseline";
const TOKEN_ID: &str = "token-owner-relay-baseline";
let pool = sqlx::postgres::PgPoolOptions::new()
.max_connections(1)
.connect(postgres_url)
.await?;
let mut transaction = pool.begin().await?;
sqlx::query(
r#"
INSERT INTO users (
id, email, username, role, auth_source, email_verified, is_active, is_deleted,
created_at, updated_at
) VALUES ($1, $2, $3, 'admin', 'local', TRUE, TRUE, FALSE, 1, 1)
ON CONFLICT (id) DO UPDATE SET
email = EXCLUDED.email,
username = EXCLUDED.username,
role = 'admin',
auth_source = 'local',
email_verified = TRUE,
is_active = TRUE,
is_deleted = FALSE,
updated_at = EXCLUDED.updated_at
"#,
)
.bind(USER_ID)
.bind("owner-relay-baseline@example.com")
.bind("owner_relay_baseline_admin")
.execute(&mut *transaction)
.await?;
let token_hash = format!(
"{:x}",
Sha256::digest(TUNNEL_HARNESS_MANAGEMENT_TOKEN.as_bytes())
);
sqlx::query(
r#"
INSERT INTO management_tokens (
id, user_id, name, token_hash, token_prefix, permissions, usage_count,
is_active, created_at, updated_at
) VALUES ($1, $2, $3, $4, 'ae-tunnel-harness', $5, 0, TRUE, 1, 1)
ON CONFLICT (id) DO UPDATE SET
user_id = EXCLUDED.user_id,
name = EXCLUDED.name,
token_hash = EXCLUDED.token_hash,
token_prefix = EXCLUDED.token_prefix,
permissions = EXCLUDED.permissions,
is_active = TRUE,
updated_at = EXCLUDED.updated_at
"#,
)
.bind(TOKEN_ID)
.bind(USER_ID)
.bind("owner relay tunnel token")
.bind(token_hash)
.bind(serde_json::json!(["admin:proxy_nodes:admin"]))
.execute(&mut *transaction)
.await?;
sqlx::query(
r#"
INSERT INTO proxy_nodes (
id, tunnel_generation, name, ip, port, status, heartbeat_interval,
active_connections, total_requests, is_manual, created_at, updated_at,
config_version, tunnel_mode, tunnel_connected, failed_requests, dns_failures,
stream_errors
) VALUES (
$1, $2, 'owner relay baseline node', '127.0.0.1', 0, 'offline', 30,
0, 0, FALSE, 1, 1, 0, TRUE, FALSE, 0, 0, 0
)
ON CONFLICT (id) DO UPDATE SET
tunnel_generation = EXCLUDED.tunnel_generation,
name = EXCLUDED.name,
ip = EXCLUDED.ip,
port = EXCLUDED.port,
status = 'offline',
active_connections = 0,
tunnel_mode = TRUE,
tunnel_connected = FALSE,
updated_at = EXCLUDED.updated_at
"#,
)
.bind(NODE_ID)
.bind(TUNNEL_HARNESS_GENERATION)
.execute(&mut *transaction)
.await?;
transaction.commit().await?;
Ok(())
}
async fn handle_binary_frame<S>(
sink: &mut S,
data: Vec<u8>,
@@ -17,6 +17,7 @@ use sqlx::{PgPool, Row};
use tokio::sync::Mutex;
const PROXY_NODE_ID: &str = "proxy-node-hotspot";
const PROXY_NODE_TUNNEL_GENERATION: &str = "usage-aux-hotspot-generation";
const MANAGEMENT_TOKEN_ID: &str = "management-token-hotspot";
const API_KEY_ID: &str = "api-key-last-used-hotspot";
const USER_ID: &str = "usage-aux-hotspot-user";
@@ -324,6 +325,7 @@ async fn enqueue_aux_counter_deltas(
fn proxy_delta_for_index(index: usize) -> ProxyNodeCounterDelta {
ProxyNodeCounterDelta {
node_id: PROXY_NODE_ID.to_string(),
expected_tunnel_generation: Some(PROXY_NODE_TUNNEL_GENERATION.to_string()),
total_requests_delta: 1,
failed_requests_delta: if index.is_multiple_of(10) { 1 } else { 0 },
dns_failures_delta: if index.is_multiple_of(25) { 1 } else { 0 },
@@ -457,10 +459,13 @@ ON CONFLICT (id) DO UPDATE SET
sqlx::query(
r#"
INSERT INTO proxy_nodes (
id, name, ip, port, status, total_requests, failed_requests,
id, tunnel_generation, name, ip, port, status, total_requests, failed_requests,
dns_failures, stream_errors
)
VALUES ($1, 'usage aux hotspot proxy', '127.0.0.1', 8080, 'online', 0, 0, 0, 0)
VALUES (
$1, 'usage-aux-hotspot-generation', 'usage aux hotspot proxy',
'127.0.0.1', 8080, 'online', 0, 0, 0, 0
)
ON CONFLICT (id) DO UPDATE SET
total_requests = 0,
failed_requests = 0,
@@ -15,7 +15,7 @@ use std::sync::Arc;
use std::time::Duration;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::CreateStandaloneApiKeyRecord;
use aether_data::repository::auth::{CreateStandaloneApiKeyRecord, CreateUserApiKeyRecord};
use aether_data::repository::wallet::WalletLookupKey;
use aether_data::{
DataBackends, DataLayerConfig, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig,
@@ -162,6 +162,38 @@ async fn continuation_reuses_one_upstream_connection_and_bills_both_turns() -> R
Ok(())
}
#[tokio::test]
async fn weekly_plan_request_limit_counts_turns_but_not_the_websocket_upgrade(
) -> Result<(), BoxError> {
let harness = Harness::start_with_weekly_request_limit(1).await?;
let mut client = harness.connect().await?;
client
.send(response_create(json!({"input": "allowed turn"})))
.await?;
receive_event(&mut client, "response.completed").await?;
client
.send(response_create(json!({"input": "rejected turn"})))
.await?;
let rejected = receive_error_or_close(&mut client)
.await?
.ok_or("gateway closed without a plan-limit error event")?;
assert_eq!(rejected["status"], json!(429));
assert_eq!(
rejected.pointer("/error/code"),
Some(&json!("plan_usage_limit_exceeded"))
);
assert_eq!(
harness.upstream.observed_events().await.len(),
1,
"the rejected logical turn must never reach the upstream"
);
client.close(None).await?;
Ok(())
}
#[tokio::test]
async fn persisted_previous_response_can_continue_on_a_new_client_connection(
) -> Result<(), BoxError> {
@@ -793,6 +825,28 @@ impl Harness {
.await
}
async fn start_with_weekly_request_limit(limit: u64) -> Result<Self, BoxError> {
let mut harness = Self::start(UpstreamBehavior::CompleteEveryTurn).await?;
seed_weekly_request_limit(&harness.database.config, limit).await?;
// The gateway may have cached the pre-entitlement auth context during
// startup, so restart it after seeding the user-owned key and plan.
let data_config = GatewayDataConfig::from_database_config(harness.database.config.clone())
.with_encryption_key(DEVELOPMENT_ENCRYPTION_KEY);
let state = AppState::new()?
.with_data_config_and_background_isolation(data_config, false)?
.with_usage_runtime_config(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
})?;
let gateway_server = SpawnedServer::start(build_router_with_state(state)).await?;
harness.websocket_url = format!(
"{}/v1/responses",
gateway_server.base_url().replacen("http://", "ws://", 1)
);
harness._gateway_server = gateway_server;
Ok(harness)
}
async fn start_with(
behavior: UpstreamBehavior,
fixture: ProviderFixture,
@@ -1617,6 +1671,145 @@ async fn seed_client_api_key(backends: &DataBackends, user_id: &str) -> Result<(
Ok(())
}
async fn seed_weekly_request_limit(
database: &SqlDatabaseConfig,
limit: u64,
) -> Result<(), BoxError> {
let backends = DataBackends::from_config(DataLayerConfig::from_database(database.clone()))?;
let user_id = backends
.read()
.users()
.ok_or("user reader unavailable")?
.find_user_auth_by_username("responses-ws-e2e")
.await?
.ok_or("seeded E2E user unavailable")?
.id;
backends
.write()
.auth_api_keys()
.ok_or("auth API key writer unavailable")?
.delete_standalone_api_key(API_KEY_ID)
.await?;
backends
.write()
.auth_api_keys()
.ok_or("auth API key writer unavailable")?
.create_user_api_key(CreateUserApiKeyRecord {
user_id: user_id.clone(),
api_key_id: API_KEY_ID.to_string(),
key_hash: sha256_hex(CLIENT_API_KEY),
key_encrypted: Some(CLIENT_API_KEY.to_string()),
name: Some("Responses WebSocket plan-policy E2E".to_string()),
allowed_providers: Some(vec![PROVIDER_ID.to_string()]),
allowed_api_formats: Some(vec!["openai:responses".to_string()]),
allowed_models: Some(vec![PUBLIC_MODEL.to_string()]),
ip_rules: None,
rate_limit: 0,
concurrent_limit: None,
force_capabilities: None,
feature_settings: None,
is_active: true,
expires_at_unix_secs: None,
auto_delete_on_expiry: false,
total_requests: 0,
total_tokens: 0,
total_cost_usd: 0.0,
})
.await?;
let wallet_id = backends
.read()
.wallets()
.ok_or("wallet reader unavailable")?
.find(WalletLookupKey::UserId(&user_id))
.await?
.ok_or("seeded E2E user wallet unavailable")?
.id;
let pool = backends
.sqlite()
.ok_or("SQLite backend unavailable")?
.pool();
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)?
.as_secs()
.min(i64::MAX as u64) as i64;
sqlx::query(
r#"
INSERT INTO billing_plans (
id, title, description, price_amount, price_currency, duration_unit,
duration_value, enabled, sort_order, max_active_per_user,
purchase_limit_scope, entitlements_json, created_at, updated_at
) VALUES (?, 'WS weekly policy', NULL, 0, 'USD', 'month', 1, 1, 0, 1,
'active_period', ?, ?, ?)
"#,
)
.bind("plan-responses-ws-weekly")
.bind(
json!([{
"type": "usage_policy",
"rules": [{
"metric": "request_count",
"window": {"kind": "calendar_week"},
"limit": limit
}]
}])
.to_string(),
)
.bind(now)
.bind(now)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO payment_orders (
id, order_no, wallet_id, user_id, amount_usd, pay_currency, status,
payment_method, created_at, paid_at, credited_at, expires_at
) VALUES (?, ?, ?, ?, 0, 'USD', 'paid', 'test', ?, ?, ?, ?)
"#,
)
.bind("order-responses-ws-weekly")
.bind("order-no-responses-ws-weekly")
.bind(&wallet_id)
.bind(&user_id)
.bind(now)
.bind(now)
.bind(now)
.bind(now + 86_400)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO user_plan_entitlements (
id, user_id, plan_id, payment_order_id, status, starts_at, expires_at,
entitlements_snapshot, created_at, updated_at
) VALUES (?, ?, ?, ?, 'active', ?, ?, ?, ?, ?)
"#,
)
.bind("entitlement-responses-ws-weekly")
.bind(&user_id)
.bind("plan-responses-ws-weekly")
.bind("order-responses-ws-weekly")
.bind(now - 1)
.bind(now + 86_400)
.bind(
json!([{
"type": "usage_policy",
"rules": [{
"metric": "request_count",
"window": {"kind": "calendar_week"},
"limit": limit
}]
}])
.to_string(),
)
.bind(now)
.bind(now)
.execute(pool)
.await?;
Ok(())
}
/// 打开 chat PII 脱敏:系统模块开关 + 这把 client key 的 feature 开关。
///
/// 规则集刻意不写:缺省即内置规则(含 email 规则),和生产上「只打开开关」的最小