mirror of
https://github.com/telemt/telemt.git
synced 2026-11-05 23:32:00 +03:00
Rustfmt
This commit is contained in:
@@ -467,3 +467,6 @@ mod decoy_dns_tests;
|
||||
// Exact host groups retain the source shape and process-deferred limit checks.
|
||||
#[cfg(test)]
|
||||
mod host_list_tests;
|
||||
|
||||
#[cfg(test)]
|
||||
mod quota_period_tests;
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_periods_do_not_enable_access_mutation_through_generic_config_patch() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("config.toml");
|
||||
let body = "[access.users]\nalice = '0123456789abcdef0123456789abcdef'\n";
|
||||
std::fs::write(&path, body).unwrap();
|
||||
let error = apply_patch_to_path(
|
||||
&path,
|
||||
&serde_json::json!({"access":{"user_quota_periods":{"alice":"daily"}}}),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.unwrap();
|
||||
assert_eq!(error.code, "access_not_editable");
|
||||
assert_eq!(std::fs::read_to_string(path).unwrap(), body);
|
||||
}
|
||||
@@ -29,6 +29,8 @@ pub(super) enum AccessSection {
|
||||
UserMaxTcpConns,
|
||||
UserExpirations,
|
||||
UserDataQuota,
|
||||
/// Per-user recurring UTC calendar settings.
|
||||
UserQuotaPeriods,
|
||||
UserRateLimits,
|
||||
UserMaxUniqueIps,
|
||||
}
|
||||
@@ -42,6 +44,7 @@ impl AccessSection {
|
||||
Self::UserMaxTcpConns => "access.user_max_tcp_conns",
|
||||
Self::UserExpirations => "access.user_expirations",
|
||||
Self::UserDataQuota => "access.user_data_quota",
|
||||
Self::UserQuotaPeriods => "access.user_quota_periods",
|
||||
Self::UserRateLimits => "access.user_rate_limits",
|
||||
Self::UserMaxUniqueIps => "access.user_max_unique_ips",
|
||||
}
|
||||
|
||||
@@ -231,6 +231,10 @@ pub(super) fn render_access_section(
|
||||
.collect();
|
||||
serialize_table_body(&rows)?
|
||||
}
|
||||
AccessSection::UserQuotaPeriods => {
|
||||
let rows: BTreeMap<_, _> = cfg.access.user_quota_periods.iter().collect();
|
||||
serialize_table_body(&rows)?
|
||||
}
|
||||
AccessSection::UserRateLimits => {
|
||||
let rows: BTreeMap<String, RateLimitBps> = cfg
|
||||
.access
|
||||
@@ -269,6 +273,7 @@ fn access_section_is_empty(cfg: &ProxyConfig, section: AccessSection) -> bool {
|
||||
AccessSection::UserMaxTcpConns => cfg.access.user_max_tcp_conns.is_empty(),
|
||||
AccessSection::UserExpirations => cfg.access.user_expirations.is_empty(),
|
||||
AccessSection::UserDataQuota => cfg.access.user_data_quota.is_empty(),
|
||||
AccessSection::UserQuotaPeriods => cfg.access.user_quota_periods.is_empty(),
|
||||
AccessSection::UserRateLimits => cfg.access.user_rate_limits.is_empty(),
|
||||
AccessSection::UserMaxUniqueIps => cfg.access.user_max_unique_ips.is_empty(),
|
||||
}
|
||||
|
||||
@@ -200,7 +200,13 @@ pub(super) async fn handle(
|
||||
("GET", "/v1/stats/users/quota") => {
|
||||
let revision = current_revision(&shared.config_path).await?;
|
||||
let disk_cfg = load_config_from_disk(&shared.config_path).await?;
|
||||
let data = build_user_quota_list(&disk_cfg, shared.stats.as_ref());
|
||||
let _quota_guard = shared.quota_state.lock().await;
|
||||
let mut data = build_user_quota_list(&disk_cfg, shared.stats.as_ref());
|
||||
for entry in &mut data.users {
|
||||
entry.schedule = shared
|
||||
.quota_state
|
||||
.schedule_status(&entry.username, entry.quota_period);
|
||||
}
|
||||
Ok(success_response(StatusCode::OK, data, revision))
|
||||
}
|
||||
|
||||
|
||||
@@ -23,6 +23,8 @@ pub(in crate::api) struct UserInfo {
|
||||
pub(in crate::api) max_tcp_conns: Option<usize>,
|
||||
pub(in crate::api) expiration_rfc3339: Option<String>,
|
||||
pub(in crate::api) data_quota_bytes: Option<u64>,
|
||||
/// Desired UTC period; runtime activation is reported by the quota endpoint.
|
||||
pub(in crate::api) quota_period: Option<crate::config::UserQuotaPeriod>,
|
||||
pub(in crate::api) rate_limit_up_bps: Option<u64>,
|
||||
pub(in crate::api) rate_limit_down_bps: Option<u64>,
|
||||
pub(in crate::api) max_unique_ips: Option<usize>,
|
||||
@@ -69,6 +71,11 @@ pub(in crate::api) struct UserQuotaListData {
|
||||
pub(in crate::api) struct UserQuotaEntry {
|
||||
pub(in crate::api) username: String,
|
||||
pub(in crate::api) data_quota_bytes: u64,
|
||||
/// Configured period may precede asynchronous runtime activation.
|
||||
pub(in crate::api) quota_period: Option<crate::config::UserQuotaPeriod>,
|
||||
/// Effective schedule comes from the process authority, never desired TOML.
|
||||
#[serde(flatten)]
|
||||
pub(in crate::api) schedule: crate::stats::QuotaScheduleStatus,
|
||||
pub(in crate::api) used_bytes: u64,
|
||||
pub(in crate::api) last_reset_epoch_secs: u64,
|
||||
}
|
||||
@@ -81,6 +88,8 @@ pub(in crate::api) struct CreateUserRequest {
|
||||
pub(in crate::api) max_tcp_conns: Option<usize>,
|
||||
pub(in crate::api) expiration_rfc3339: Option<String>,
|
||||
pub(in crate::api) data_quota_bytes: Option<u64>,
|
||||
/// Desired UTC period; runtime activation is reported by the quota endpoint.
|
||||
pub(in crate::api) quota_period: Option<crate::config::UserQuotaPeriod>,
|
||||
pub(in crate::api) rate_limit_up_bps: Option<u64>,
|
||||
pub(in crate::api) rate_limit_down_bps: Option<u64>,
|
||||
pub(in crate::api) max_unique_ips: Option<usize>,
|
||||
@@ -98,6 +107,9 @@ pub(in crate::api) struct PatchUserRequest {
|
||||
pub(in crate::api) expiration_rfc3339: Patch<String>,
|
||||
#[serde(default, deserialize_with = "patch_field")]
|
||||
pub(in crate::api) data_quota_bytes: Patch<u64>,
|
||||
/// Omitted preserves the period; null removes it.
|
||||
#[serde(default, deserialize_with = "patch_field")]
|
||||
pub(in crate::api) quota_period: Patch<crate::config::UserQuotaPeriod>,
|
||||
#[serde(default, deserialize_with = "patch_field")]
|
||||
pub(in crate::api) rate_limit_up_bps: Patch<u64>,
|
||||
#[serde(default, deserialize_with = "patch_field")]
|
||||
|
||||
@@ -35,3 +35,6 @@ pub(super) use view::{build_user_quota_list, users_from_config};
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
#[cfg(test)]
|
||||
mod period_tests;
|
||||
|
||||
@@ -58,6 +58,7 @@ async fn create_user_to_completion(
|
||||
let credential_id = credential_id_from_hex(&secret)
|
||||
.ok_or_else(|| ApiFailure::internal("validated user secret could not be decoded"))?;
|
||||
let _guard = shared.mutation_lock.lock().await;
|
||||
let quota_guard = shared.quota_state.lock().await;
|
||||
let (mut cfg, base_revision) =
|
||||
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
|
||||
|
||||
@@ -112,10 +113,18 @@ async fn create_user_to_completion(
|
||||
cfg.access.user_enabled.insert(body.username.clone(), false);
|
||||
}
|
||||
|
||||
if let Some(period) = body.quota_period {
|
||||
cfg.access
|
||||
.user_quota_periods
|
||||
.insert(body.username.clone(), period);
|
||||
}
|
||||
cfg.validate()
|
||||
.map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?;
|
||||
|
||||
let mut touched_sections = vec![AccessSection::Users];
|
||||
if body.quota_period.is_some() {
|
||||
touched_sections.push(AccessSection::UserQuotaPeriods);
|
||||
}
|
||||
if touches_user_ad_tags {
|
||||
touched_sections.push(AccessSection::UserAdTags);
|
||||
}
|
||||
@@ -151,6 +160,14 @@ async fn create_user_to_completion(
|
||||
cfg.access.is_user_enabled(&body.username),
|
||||
);
|
||||
|
||||
shared.quota_state.fence_fresh_schedule(
|
||||
&body.username,
|
||||
&shared.active_runtime.load_full().config(),
|
||||
credential_id,
|
||||
cfg.access.user_quota_periods.get(&body.username).copied(),
|
||||
cfg.access.user_data_quota.get(&body.username).copied(),
|
||||
);
|
||||
drop(quota_guard);
|
||||
if let Some(limit) = updated_limit {
|
||||
shared
|
||||
.ip_tracker
|
||||
@@ -187,6 +204,7 @@ async fn create_user_to_completion(
|
||||
.then_some(cfg.access.user_max_tcp_conns_global_each)),
|
||||
expiration_rfc3339: None,
|
||||
data_quota_bytes: None,
|
||||
quota_period: body.quota_period,
|
||||
rate_limit_up_bps: body.rate_limit_up_bps.filter(|limit| *limit > 0),
|
||||
rate_limit_down_bps: body.rate_limit_down_bps.filter(|limit| *limit > 0),
|
||||
max_unique_ips: updated_limit,
|
||||
|
||||
@@ -33,6 +33,7 @@ async fn rotate_secret_to_completion(
|
||||
.ok_or_else(|| ApiFailure::internal("validated user secret could not be decoded"))?;
|
||||
|
||||
let _guard = shared.mutation_lock.lock().await;
|
||||
let quota_guard = shared.quota_state.lock().await;
|
||||
let (mut cfg, base_revision) =
|
||||
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
|
||||
|
||||
@@ -59,6 +60,8 @@ async fn rotate_secret_to_completion(
|
||||
credential_id,
|
||||
cfg.access.is_user_enabled(user),
|
||||
);
|
||||
shared.quota_state.refresh_fresh_schedule(user, &cfg);
|
||||
drop(quota_guard);
|
||||
drop(_guard);
|
||||
|
||||
let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips();
|
||||
@@ -106,6 +109,7 @@ async fn delete_user_to_completion(
|
||||
shared: &ApiShared,
|
||||
) -> Result<(String, String), ApiFailure> {
|
||||
let _guard = shared.mutation_lock.lock().await;
|
||||
let quota_guard = shared.quota_state.lock().await;
|
||||
let (mut cfg, base_revision) =
|
||||
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
|
||||
|
||||
@@ -138,6 +142,9 @@ async fn delete_user_to_completion(
|
||||
if cfg.access.user_expirations.remove(user).is_some() {
|
||||
touched_sections.push(AccessSection::UserExpirations);
|
||||
}
|
||||
if cfg.access.user_quota_periods.remove(user).is_some() {
|
||||
touched_sections.push(AccessSection::UserQuotaPeriods);
|
||||
}
|
||||
if cfg.access.user_data_quota.remove(user).is_some() {
|
||||
touched_sections.push(AccessSection::UserDataQuota);
|
||||
}
|
||||
@@ -161,7 +168,7 @@ async fn delete_user_to_completion(
|
||||
let configured_users = cfg.access.users.keys().cloned().collect();
|
||||
if let Err(error) = shared
|
||||
.quota_state
|
||||
.remove_user(&configured_users, user)
|
||||
.remove_user_locked(&configured_users, user, quota_guard)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
|
||||
@@ -0,0 +1,445 @@
|
||||
use super::*;
|
||||
use crate::api::{ApiRuntimeState, events::ApiEventStore};
|
||||
use crate::config::UserQuotaPeriod;
|
||||
use crate::maestro::{
|
||||
control_plane::ProcessControlPlane, generation::test_runtime_generation, reload::ReloadControl,
|
||||
};
|
||||
use crate::proxy::{shared_state::ProxySharedState, user_admission::UserAdmissionAuthority};
|
||||
use crate::quota_state::QuotaStateOwner;
|
||||
use crate::stats::QuotaStore;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64};
|
||||
use tokio::sync::{Mutex, watch};
|
||||
|
||||
struct Fixture {
|
||||
_dir: tempfile::TempDir,
|
||||
shared: ApiShared,
|
||||
store: Arc<QuotaStore>,
|
||||
config_tx: watch::Sender<Arc<ProxyConfig>>,
|
||||
}
|
||||
|
||||
impl Fixture {
|
||||
fn new() -> Self {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let config_path = dir.path().join("config.toml");
|
||||
std::fs::write(&config_path, "# preserve this operator comment\n[access.users]\nadmin = '00000000000000000000000000000000'\n").unwrap();
|
||||
let cfg = ProxyConfig::load(&config_path).unwrap();
|
||||
let store = Arc::new(QuotaStore::default());
|
||||
let stats = Arc::new(Stats::with_quota_store(store.clone()));
|
||||
let mut runtime = test_runtime_generation(1, cfg.clone());
|
||||
let authority = UserAdmissionAuthority::new_with_quota_store(store.clone());
|
||||
let proxy_shared = ProxySharedState::new_with_direct_buffer_budget_and_user_admission(
|
||||
runtime.proxy_shared.direct_buffer_budget.clone(),
|
||||
authority,
|
||||
);
|
||||
proxy_shared.activate_user_config_source(
|
||||
1,
|
||||
None,
|
||||
&cfg.access.users,
|
||||
&cfg.access.user_enabled,
|
||||
);
|
||||
let (config_tx, config_rx) = watch::channel(Arc::new(cfg.clone()));
|
||||
let mutable = Arc::get_mut(&mut runtime).unwrap();
|
||||
mutable.stats = stats.clone();
|
||||
mutable.proxy_shared = proxy_shared.clone();
|
||||
mutable.config_rx = config_rx;
|
||||
let shared = ApiShared {
|
||||
stats,
|
||||
ip_tracker: runtime.ip_tracker.clone(),
|
||||
me_pool: runtime.me_pool_runtime.clone(),
|
||||
upstream_manager: runtime.upstream_manager.clone(),
|
||||
config_path,
|
||||
quota_state: QuotaStateOwner::new(dir.path().join("quota.json"), store.clone()),
|
||||
detected_ips_rx: watch::channel((None, None)).1,
|
||||
mutation_lock: Arc::new(Mutex::new(())),
|
||||
minimal_cache: Arc::new(Mutex::new(None)),
|
||||
runtime_edge_connections_cache: Arc::new(Mutex::new(None)),
|
||||
runtime_edge_recompute_lock: Arc::new(Mutex::new(())),
|
||||
cache_generation: Arc::new(AtomicU64::new(1)),
|
||||
runtime_events: Arc::new(ApiEventStore::new(64)),
|
||||
request_id: Arc::new(AtomicU64::new(1)),
|
||||
runtime_state: Arc::new(ApiRuntimeState {
|
||||
process_started_at_epoch_secs: 0,
|
||||
config_reload_count: AtomicU64::new(0),
|
||||
last_config_reload_epoch_secs: AtomicU64::new(0),
|
||||
admission_open: AtomicBool::new(true),
|
||||
}),
|
||||
startup_tracker: Arc::new(crate::startup::StartupTracker::new(0)),
|
||||
route_runtime: runtime.route_runtime.clone(),
|
||||
proxy_shared,
|
||||
reload_control: ReloadControl::channel(1).0,
|
||||
active_runtime: Arc::new(arc_swap::ArcSwap::from(runtime)),
|
||||
web_trace: crate::web::trace::WebTraceStore::new(
|
||||
cfg.web.debug.clone(),
|
||||
&cfg.web.limits,
|
||||
),
|
||||
web_runtime_rx: crate::web::control::WebRuntimeControl::new().subscribe(),
|
||||
control_plane: ProcessControlPlane::new(),
|
||||
};
|
||||
Self {
|
||||
_dir: dir,
|
||||
shared,
|
||||
store,
|
||||
config_tx,
|
||||
}
|
||||
}
|
||||
|
||||
async fn publish(&self) {
|
||||
let cfg = Arc::new(ProxyConfig::load(&self.shared.config_path).unwrap());
|
||||
let guard = self.shared.quota_state.lock().await;
|
||||
self.config_tx.send_replace(cfg.clone());
|
||||
assert!(
|
||||
self.shared
|
||||
.proxy_shared
|
||||
.apply_user_config_from_source(1, &cfg.access.users, &cfg.access.user_enabled)
|
||||
.is_some()
|
||||
);
|
||||
self.shared.proxy_shared.publish_quota_config(1, &cfg, true);
|
||||
drop(guard);
|
||||
self.shared
|
||||
.quota_state
|
||||
.process_schedules(crate::quota_state::now_epoch_secs(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn finish(&self) {
|
||||
assert!(
|
||||
self.shared
|
||||
.control_plane
|
||||
.shutdown(std::time::Duration::from_secs(1))
|
||||
.await
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn patch(value: serde_json::Value) -> PatchUserRequest {
|
||||
serde_json::from_value(value).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_period_api_updates_included_owner_and_preserves_root_bytes() {
|
||||
let fixture = Fixture::new();
|
||||
let root = "# root remains unchanged\ninclude = \"users.toml\"\n";
|
||||
let included = fixture.shared.config_path.with_file_name("users.toml");
|
||||
std::fs::write(&included, "[access.users]\nadmin = '00000000000000000000000000000000'\n[access.user_data_quota]\nadmin = 1024\n[access.user_quota_periods]\nadmin = 'daily'\n").unwrap();
|
||||
std::fs::write(&fixture.shared.config_path, root).unwrap();
|
||||
patch_user(
|
||||
"admin",
|
||||
patch(serde_json::json!({"quota_period":"monthly"})),
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&fixture.shared.config_path).unwrap(),
|
||||
root
|
||||
);
|
||||
let cfg = ProxyConfig::load(&fixture.shared.config_path).unwrap();
|
||||
assert_eq!(
|
||||
cfg.access.user_quota_periods["admin"],
|
||||
UserQuotaPeriod::Monthly
|
||||
);
|
||||
assert!(
|
||||
std::fs::read_to_string(included)
|
||||
.unwrap()
|
||||
.contains("monthly")
|
||||
);
|
||||
fixture.finish().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_period_prepared_before_recreation_cannot_attach_retired_schedule() {
|
||||
let fixture = Fixture::new();
|
||||
let create = |period| {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"username":"alice", "secret":"0123456789abcdef0123456789abcdef",
|
||||
"data_quota_bytes":1024, "quota_period":period
|
||||
}))
|
||||
.unwrap()
|
||||
};
|
||||
create_user(create("daily"), None, &fixture.shared)
|
||||
.await
|
||||
.unwrap();
|
||||
fixture.publish().await;
|
||||
// Capture a real source before recreation; its asynchronous preparation finishes later.
|
||||
let mut stale = ProxyConfig::parse_source(&fixture.shared.config_path).unwrap();
|
||||
stale.config.general.ad_tag = Some("11111111111111111111111111111111".into());
|
||||
delete_user("alice", None, &fixture.shared).await.unwrap();
|
||||
create_user(create("monthly"), None, &fixture.shared)
|
||||
.await
|
||||
.unwrap();
|
||||
let stale = Arc::new(stale.prepare().await.unwrap().config);
|
||||
let guard = fixture.shared.quota_state.lock().await;
|
||||
fixture.config_tx.send_replace(stale.clone());
|
||||
fixture.shared.proxy_shared.apply_user_config_from_source(
|
||||
1,
|
||||
&stale.access.users,
|
||||
&stale.access.user_enabled,
|
||||
);
|
||||
fixture
|
||||
.shared
|
||||
.proxy_shared
|
||||
.publish_quota_config(1, &stale, true);
|
||||
drop(guard);
|
||||
fixture
|
||||
.shared
|
||||
.quota_state
|
||||
.process_schedules(crate::quota_state::now_epoch_secs(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
fixture
|
||||
.store
|
||||
.current_handle("alice")
|
||||
.unwrap()
|
||||
.schedule_checkpoint()
|
||||
.quota_period,
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
fixture
|
||||
.shared
|
||||
.quota_state
|
||||
.schedule_status("alice", Some(UserQuotaPeriod::Monthly))
|
||||
.quota_schedule_state,
|
||||
"pending"
|
||||
);
|
||||
fixture.publish().await;
|
||||
assert_eq!(
|
||||
fixture
|
||||
.store
|
||||
.current_handle("alice")
|
||||
.unwrap()
|
||||
.schedule_checkpoint()
|
||||
.quota_period,
|
||||
Some(UserQuotaPeriod::Monthly)
|
||||
);
|
||||
fixture.finish().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_period_fresh_patch_and_rotation_accept_matching_stale_epoch_candidate() {
|
||||
let fixture = Fixture::new();
|
||||
create_user(
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"username":"alice", "data_quota_bytes":1024, "quota_period":"daily"
|
||||
}))
|
||||
.unwrap(),
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
patch_user(
|
||||
"alice",
|
||||
patch(serde_json::json!({"quota_period":"monthly"})),
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
rotate_secret(
|
||||
"alice",
|
||||
RotateSecretRequest {
|
||||
secret: Some("11111111111111111111111111111111".into()),
|
||||
},
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let cfg = Arc::new(ProxyConfig::load(&fixture.shared.config_path).unwrap());
|
||||
let guard = fixture.shared.quota_state.lock().await;
|
||||
fixture
|
||||
.shared
|
||||
.proxy_shared
|
||||
.publish_quota_config(1, &cfg, false);
|
||||
drop(guard);
|
||||
fixture
|
||||
.shared
|
||||
.quota_state
|
||||
.process_schedules(crate::quota_state::now_epoch_secs(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
fixture
|
||||
.store
|
||||
.current_handle("alice")
|
||||
.unwrap()
|
||||
.schedule_checkpoint()
|
||||
.quota_period,
|
||||
Some(UserQuotaPeriod::Monthly)
|
||||
);
|
||||
fixture.finish().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_period_api_create_patch_remove_delete_and_effective_status() {
|
||||
let fixture = Fixture::new();
|
||||
let created = create_user(serde_json::from_value(serde_json::json!({
|
||||
"username": "alice", "secret": "0123456789abcdef0123456789abcdef", "data_quota_bytes": 1024, "quota_period": "daily"
|
||||
})).unwrap(), None, &fixture.shared).await.unwrap().0;
|
||||
assert_eq!(created.user.quota_period, Some(UserQuotaPeriod::Daily));
|
||||
assert_eq!(
|
||||
fixture
|
||||
.shared
|
||||
.quota_state
|
||||
.schedule_status("alice", Some(UserQuotaPeriod::Daily))
|
||||
.quota_schedule_state,
|
||||
"pending"
|
||||
);
|
||||
fixture.store.current_handle("alice").unwrap().charge(300);
|
||||
fixture.publish().await;
|
||||
assert_eq!(fixture.store.used("alice"), 300);
|
||||
let before = fixture
|
||||
.store
|
||||
.current_handle("alice")
|
||||
.unwrap()
|
||||
.schedule_checkpoint();
|
||||
patch_user(
|
||||
"alice",
|
||||
patch(serde_json::json!({"quota_period":"monthly"})),
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let status = fixture
|
||||
.shared
|
||||
.quota_state
|
||||
.schedule_status("alice", Some(UserQuotaPeriod::Monthly));
|
||||
assert_eq!(status.active_quota_period, Some(UserQuotaPeriod::Daily));
|
||||
assert_eq!(status.quota_schedule_state, "pending");
|
||||
fixture.publish().await;
|
||||
assert_eq!(fixture.store.used("alice"), 300);
|
||||
assert!(
|
||||
fixture
|
||||
.store
|
||||
.current_handle("alice")
|
||||
.unwrap()
|
||||
.schedule_checkpoint()
|
||||
.period_cursor_epoch_secs
|
||||
>= before.period_cursor_epoch_secs
|
||||
);
|
||||
let bytes = std::fs::read(&fixture.shared.config_path).unwrap();
|
||||
assert!(
|
||||
patch_user(
|
||||
"alice",
|
||||
patch(serde_json::json!({"data_quota_bytes":null})),
|
||||
None,
|
||||
&fixture.shared
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(std::fs::read(&fixture.shared.config_path).unwrap(), bytes);
|
||||
patch_user(
|
||||
"alice",
|
||||
patch(serde_json::json!({"data_quota_bytes":null,"quota_period":null})),
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
fixture.publish().await;
|
||||
assert_eq!(fixture.store.used("alice"), 300);
|
||||
assert_eq!(
|
||||
fixture
|
||||
.shared
|
||||
.quota_state
|
||||
.schedule_status("alice", None)
|
||||
.quota_schedule_state,
|
||||
"disabled"
|
||||
);
|
||||
patch_user(
|
||||
"alice",
|
||||
patch(serde_json::json!({"data_quota_bytes":1024,"quota_period":"weekly"})),
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
fixture.publish().await;
|
||||
delete_user("alice", None, &fixture.shared).await.unwrap();
|
||||
let cfg = ProxyConfig::load(&fixture.shared.config_path).unwrap();
|
||||
assert!(!cfg.access.users.contains_key("alice"));
|
||||
assert!(!cfg.access.user_quota_periods.contains_key("alice"));
|
||||
assert!(fixture.store.current_handle("alice").is_none());
|
||||
assert!(
|
||||
std::fs::read_to_string(&fixture.shared.config_path)
|
||||
.unwrap()
|
||||
.contains("# preserve this operator comment")
|
||||
);
|
||||
fixture.finish().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_period_api_early_rotation_and_revision_conflicts_are_safe() {
|
||||
let fixture = Fixture::new();
|
||||
let (_, old_revision) = create_user(
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"username":"alice", "data_quota_bytes":1024, "quota_period":"monthly"
|
||||
}))
|
||||
.unwrap(),
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
fixture.store.current_handle("alice").unwrap().charge(200);
|
||||
rotate_secret(
|
||||
"alice",
|
||||
RotateSecretRequest {
|
||||
secret: Some("11111111111111111111111111111111".into()),
|
||||
},
|
||||
None,
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
fixture.publish().await;
|
||||
assert_eq!(fixture.store.used("alice"), 200);
|
||||
assert_eq!(
|
||||
fixture
|
||||
.shared
|
||||
.quota_state
|
||||
.schedule_status("alice", Some(UserQuotaPeriod::Monthly))
|
||||
.quota_schedule_state,
|
||||
"active"
|
||||
);
|
||||
let bytes = std::fs::read(&fixture.shared.config_path).unwrap();
|
||||
let error = patch_user(
|
||||
"alice",
|
||||
patch(serde_json::json!({"quota_period":"daily"})),
|
||||
Some(old_revision),
|
||||
&fixture.shared,
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.unwrap();
|
||||
assert_eq!(error.code, "revision_conflict");
|
||||
assert_eq!(std::fs::read(&fixture.shared.config_path).unwrap(), bytes);
|
||||
fixture.finish().await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quota_period_patch_preserves_omitted_null_and_set_semantics() {
|
||||
assert!(matches!(
|
||||
patch(serde_json::json!({})).quota_period,
|
||||
Patch::Unchanged
|
||||
));
|
||||
assert!(matches!(
|
||||
patch(serde_json::json!({"quota_period":null})).quota_period,
|
||||
Patch::Remove
|
||||
));
|
||||
assert!(matches!(
|
||||
patch(serde_json::json!({"quota_period":"weekly"})).quota_period,
|
||||
Patch::Set(UserQuotaPeriod::Weekly)
|
||||
));
|
||||
assert!(
|
||||
serde_json::from_value::<PatchUserRequest>(serde_json::json!({"quota_period":"Weekly"}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
@@ -48,6 +48,7 @@ async fn patch_user_to_completion(
|
||||
}
|
||||
let expiration = parse_patch_expiration(&body.expiration_rfc3339)?;
|
||||
let _guard = shared.mutation_lock.lock().await;
|
||||
let quota_guard = shared.quota_state.lock().await;
|
||||
let (mut cfg, base_revision) =
|
||||
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
|
||||
|
||||
@@ -93,6 +94,18 @@ async fn patch_user_to_completion(
|
||||
.insert(user.to_string(), expiration);
|
||||
}
|
||||
}
|
||||
let touches_user_quota_periods = !matches!(body.quota_period, Patch::Unchanged);
|
||||
match body.quota_period {
|
||||
Patch::Unchanged => {}
|
||||
Patch::Remove => {
|
||||
cfg.access.user_quota_periods.remove(user);
|
||||
}
|
||||
Patch::Set(period) => {
|
||||
cfg.access
|
||||
.user_quota_periods
|
||||
.insert(user.to_string(), period);
|
||||
}
|
||||
}
|
||||
match body.data_quota_bytes {
|
||||
Patch::Unchanged => {}
|
||||
Patch::Remove => {
|
||||
@@ -169,6 +182,9 @@ async fn patch_user_to_completion(
|
||||
};
|
||||
|
||||
let mut touched_sections = Vec::new();
|
||||
if touches_user_quota_periods {
|
||||
touched_sections.push(AccessSection::UserQuotaPeriods);
|
||||
}
|
||||
if touches_users {
|
||||
touched_sections.push(AccessSection::Users);
|
||||
}
|
||||
@@ -212,6 +228,8 @@ async fn patch_user_to_completion(
|
||||
cfg.access.is_user_enabled(user),
|
||||
);
|
||||
}
|
||||
shared.quota_state.refresh_fresh_schedule(user, &cfg);
|
||||
drop(quota_guard);
|
||||
match max_unique_ips_change {
|
||||
Some(Some(limit)) => shared.ip_tracker.set_user_limit(user, limit).await,
|
||||
Some(None) => shared.ip_tracker.remove_user_limit(user).await,
|
||||
@@ -259,6 +277,7 @@ async fn set_user_enabled_to_completion(
|
||||
shared: &ApiShared,
|
||||
) -> Result<(UserInfo, String), ApiFailure> {
|
||||
let _guard = shared.mutation_lock.lock().await;
|
||||
let quota_guard = shared.quota_state.lock().await;
|
||||
let (mut cfg, base_revision) =
|
||||
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
|
||||
|
||||
@@ -294,6 +313,8 @@ async fn set_user_enabled_to_completion(
|
||||
shared
|
||||
.proxy_shared
|
||||
.stage_user_credential(user, credential_id, enabled);
|
||||
shared.quota_state.refresh_fresh_schedule(user, &cfg);
|
||||
drop(quota_guard);
|
||||
drop(_guard);
|
||||
|
||||
let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips();
|
||||
|
||||
@@ -51,6 +51,7 @@ pub(in crate::api) async fn users_from_config(
|
||||
.get(&username)
|
||||
.map(chrono::DateTime::<chrono::Utc>::to_rfc3339),
|
||||
data_quota_bytes: cfg.access.user_data_quota.get(&username).copied(),
|
||||
quota_period: cfg.access.user_quota_periods.get(&username).copied(),
|
||||
rate_limit_up_bps: cfg
|
||||
.access
|
||||
.user_rate_limits
|
||||
@@ -102,7 +103,12 @@ pub(in crate::api) fn build_user_quota_list(cfg: &ProxyConfig, stats: &Stats) ->
|
||||
.map(|entry| (entry.used_bytes, entry.last_reset_epoch_secs))
|
||||
.unwrap_or((0, 0));
|
||||
users.push(UserQuotaEntry {
|
||||
username,
|
||||
username: username.clone(),
|
||||
quota_period: cfg.access.user_quota_periods.get(&username).copied(),
|
||||
schedule: crate::stats::QuotaScheduleStatus {
|
||||
quota_schedule_state: "disabled",
|
||||
..Default::default()
|
||||
},
|
||||
data_quota_bytes,
|
||||
used_bytes,
|
||||
last_reset_epoch_secs,
|
||||
|
||||
Reference in New Issue
Block a user