User Management in API

src/api/config_store.rs

@@ -1,13 +1,39 @@
+use std::collections::BTreeMap;
 use std::io::Write;
 use std::path::{Path, PathBuf};
 
+use chrono::{DateTime, Utc};
 use hyper::header::IF_MATCH;
+use serde::Serialize;
 use sha2::{Digest, Sha256};
 
 use crate::config::ProxyConfig;
 
 use super::model::ApiFailure;
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(super) enum AccessSection {
+    Users,
+    UserAdTags,
+    UserMaxTcpConns,
+    UserExpirations,
+    UserDataQuota,
+    UserMaxUniqueIps,
+}
+
+impl AccessSection {
+    fn table_name(self) -> &'static str {
+        match self {
+            Self::Users => "access.users",
+            Self::UserAdTags => "access.user_ad_tags",
+            Self::UserMaxTcpConns => "access.user_max_tcp_conns",
+            Self::UserExpirations => "access.user_expirations",
+            Self::UserDataQuota => "access.user_data_quota",
+            Self::UserMaxUniqueIps => "access.user_max_unique_ips",
+        }
+    }
+}
+
 pub(super) fn parse_if_match(headers: &hyper::HeaderMap) -> Option<String> {
     headers
         .get(IF_MATCH)
@@ -66,6 +92,142 @@ pub(super) async fn save_config_to_disk(
     Ok(compute_revision(&serialized))
 }
 
+pub(super) async fn save_access_sections_to_disk(
+    config_path: &Path,
+    cfg: &ProxyConfig,
+    sections: &[AccessSection],
+) -> Result<String, ApiFailure> {
+    let mut content = tokio::fs::read_to_string(config_path)
+        .await
+        .map_err(|e| ApiFailure::internal(format!("failed to read config: {}", e)))?;
+
+    let mut applied = Vec::new();
+    for section in sections {
+        if applied.contains(section) {
+            continue;
+        }
+        let rendered = render_access_section(cfg, *section)?;
+        content = upsert_toml_table(&content, section.table_name(), &rendered);
+        applied.push(*section);
+    }
+
+    write_atomic(config_path.to_path_buf(), content.clone()).await?;
+    Ok(compute_revision(&content))
+}
+
+fn render_access_section(cfg: &ProxyConfig, section: AccessSection) -> Result<String, ApiFailure> {
+    let body = match section {
+        AccessSection::Users => {
+            let rows: BTreeMap<String, String> = cfg
+                .access
+                .users
+                .iter()
+                .map(|(key, value)| (key.clone(), value.clone()))
+                .collect();
+            serialize_table_body(&rows)?
+        }
+        AccessSection::UserAdTags => {
+            let rows: BTreeMap<String, String> = cfg
+                .access
+                .user_ad_tags
+                .iter()
+                .map(|(key, value)| (key.clone(), value.clone()))
+                .collect();
+            serialize_table_body(&rows)?
+        }
+        AccessSection::UserMaxTcpConns => {
+            let rows: BTreeMap<String, usize> = cfg
+                .access
+                .user_max_tcp_conns
+                .iter()
+                .map(|(key, value)| (key.clone(), *value))
+                .collect();
+            serialize_table_body(&rows)?
+        }
+        AccessSection::UserExpirations => {
+            let rows: BTreeMap<String, DateTime<Utc>> = cfg
+                .access
+                .user_expirations
+                .iter()
+                .map(|(key, value)| (key.clone(), *value))
+                .collect();
+            serialize_table_body(&rows)?
+        }
+        AccessSection::UserDataQuota => {
+            let rows: BTreeMap<String, u64> = cfg
+                .access
+                .user_data_quota
+                .iter()
+                .map(|(key, value)| (key.clone(), *value))
+                .collect();
+            serialize_table_body(&rows)?
+        }
+        AccessSection::UserMaxUniqueIps => {
+            let rows: BTreeMap<String, usize> = cfg
+                .access
+                .user_max_unique_ips
+                .iter()
+                .map(|(key, value)| (key.clone(), *value))
+                .collect();
+            serialize_table_body(&rows)?
+        }
+    };
+
+    let mut out = format!("[{}]\n", section.table_name());
+    if !body.is_empty() {
+        out.push_str(&body);
+    }
+    if !out.ends_with('\n') {
+        out.push('\n');
+    }
+    Ok(out)
+}
+
+fn serialize_table_body<T: Serialize>(value: &T) -> Result<String, ApiFailure> {
+    toml::to_string(value)
+        .map_err(|e| ApiFailure::internal(format!("failed to serialize access section: {}", e)))
+}
+
+fn upsert_toml_table(source: &str, table_name: &str, replacement: &str) -> String {
+    if let Some((start, end)) = find_toml_table_bounds(source, table_name) {
+        let mut out = String::with_capacity(source.len() + replacement.len());
+        out.push_str(&source[..start]);
+        out.push_str(replacement);
+        out.push_str(&source[end..]);
+        return out;
+    }
+
+    let mut out = source.to_string();
+    if !out.is_empty() && !out.ends_with('\n') {
+        out.push('\n');
+    }
+    if !out.is_empty() {
+        out.push('\n');
+    }
+    out.push_str(replacement);
+    out
+}
+
+fn find_toml_table_bounds(source: &str, table_name: &str) -> Option<(usize, usize)> {
+    let target = format!("[{}]", table_name);
+    let mut offset = 0usize;
+    let mut start = None;
+
+    for line in source.split_inclusive('\n') {
+        let trimmed = line.trim();
+        if let Some(start_offset) = start {
+            if trimmed.starts_with('[') {
+                return Some((start_offset, offset));
+            }
+        } else if trimmed == target {
+            start = Some(offset);
+        }
+        offset = offset.saturating_add(line.len());
+    }
+
+    start.map(|start_offset| (start_offset, source.len()))
+}
+
 async fn write_atomic(path: PathBuf, contents: String) -> Result<(), ApiFailure> {
     tokio::task::spawn_blocking(move || write_atomic_sync(&path, &contents))
         .await

src/api/users.rs

@@ -8,7 +8,8 @@ use crate::stats::Stats;
 
 use super::ApiShared;
 use super::config_store::{
-    ensure_expected_revision, load_config_from_disk, save_config_to_disk,
+    AccessSection, ensure_expected_revision, load_config_from_disk, save_access_sections_to_disk,
+    save_config_to_disk,
 };
 use super::model::{
     ApiFailure, CreateUserRequest, CreateUserResponse, PatchUserRequest, RotateSecretRequest,
@@ -21,6 +22,12 @@ pub(super) async fn create_user(
     expected_revision: Option<String>,
     shared: &ApiShared,
 ) -> Result<(CreateUserResponse, String), ApiFailure> {
+    let touches_user_ad_tags = body.user_ad_tag.is_some();
+    let touches_user_max_tcp_conns = body.max_tcp_conns.is_some();
+    let touches_user_expirations = body.expiration_rfc3339.is_some();
+    let touches_user_data_quota = body.data_quota_bytes.is_some();
+    let touches_user_max_unique_ips = body.max_unique_ips.is_some();
+
     if !is_valid_username(&body.username) {
         return Err(ApiFailure::bad_request(
             "username must match [A-Za-z0-9_.-] and be 1..64 chars",
@@ -84,7 +91,24 @@ pub(super) async fn create_user(
     cfg.validate()
         .map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?;
 
-    let revision = save_config_to_disk(&shared.config_path, &cfg).await?;
+    let mut touched_sections = vec![AccessSection::Users];
+    if touches_user_ad_tags {
+        touched_sections.push(AccessSection::UserAdTags);
+    }
+    if touches_user_max_tcp_conns {
+        touched_sections.push(AccessSection::UserMaxTcpConns);
+    }
+    if touches_user_expirations {
+        touched_sections.push(AccessSection::UserExpirations);
+    }
+    if touches_user_data_quota {
+        touched_sections.push(AccessSection::UserDataQuota);
+    }
+    if touches_user_max_unique_ips {
+        touched_sections.push(AccessSection::UserMaxUniqueIps);
+    }
+
+    let revision = save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?;
     drop(_guard);
 
     if let Some(limit) = updated_limit {
@@ -231,7 +255,15 @@ pub(super) async fn rotate_secret(
     cfg.access.users.insert(user.to_string(), secret.clone());
     cfg.validate()
         .map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?;
-    let revision = save_config_to_disk(&shared.config_path, &cfg).await?;
+    let touched_sections = [
+        AccessSection::Users,
+        AccessSection::UserAdTags,
+        AccessSection::UserMaxTcpConns,
+        AccessSection::UserExpirations,
+        AccessSection::UserDataQuota,
+        AccessSection::UserMaxUniqueIps,
+    ];
+    let revision = save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?;
     drop(_guard);
 
     let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips();