Atomically updates with Includes

Signed-off-by: Alexey <247128645+axkurcom@users.noreply.github.com>

src/api/config_store.rs

@@ -82,6 +82,7 @@ pub(super) async fn load_config_from_disk(config_path: &Path) -> Result<ProxyCon
         .map_err(|e| ApiFailure::internal(format!("failed to load config: {}", e)))
 }
 
+#[allow(dead_code)]
 pub(super) async fn save_config_to_disk(
     config_path: &Path,
     cfg: &ProxyConfig,
@@ -106,6 +107,12 @@ pub(super) async fn save_access_sections_to_disk(
         if applied.contains(section) {
             continue;
         }
+        if find_toml_table_bounds(&content, section.table_name()).is_none()
+            && access_section_is_empty(cfg, *section)
+        {
+            applied.push(*section);
+            continue;
+        }
         let rendered = render_access_section(cfg, *section)?;
         content = upsert_toml_table(&content, section.table_name(), &rendered);
         applied.push(*section);
@@ -183,6 +190,17 @@ fn render_access_section(cfg: &ProxyConfig, section: AccessSection) -> Result<St
     Ok(out)
 }
 
+fn access_section_is_empty(cfg: &ProxyConfig, section: AccessSection) -> bool {
+    match section {
+        AccessSection::Users => cfg.access.users.is_empty(),
+        AccessSection::UserAdTags => cfg.access.user_ad_tags.is_empty(),
+        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::UserMaxUniqueIps => cfg.access.user_max_unique_ips.is_empty(),
+    }
+}
+
 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)))

src/api/users.rs

@@ -8,8 +8,8 @@ use crate::stats::Stats;
 
 use super::ApiShared;
 use super::config_store::{
-    AccessSection, ensure_expected_revision, load_config_from_disk, save_access_sections_to_disk,
-    save_config_to_disk,
+    AccessSection, current_revision, ensure_expected_revision, load_config_from_disk,
+    save_access_sections_to_disk,
 };
 use super::model::{
     ApiFailure, CreateUserRequest, CreateUserResponse, PatchUserRequest, RotateSecretRequest,
@@ -176,6 +176,13 @@ pub(super) async fn patch_user(
     expected_revision: Option<String>,
     shared: &ApiShared,
 ) -> Result<(UserInfo, String), ApiFailure> {
+    let touches_users = body.secret.is_some();
+    let touches_user_ad_tags = !matches!(&body.user_ad_tag, Patch::Unchanged);
+    let touches_user_max_tcp_conns = !matches!(&body.max_tcp_conns, Patch::Unchanged);
+    let touches_user_expirations = !matches!(&body.expiration_rfc3339, Patch::Unchanged);
+    let touches_user_data_quota = !matches!(&body.data_quota_bytes, Patch::Unchanged);
+    let touches_user_max_unique_ips = !matches!(&body.max_unique_ips, Patch::Unchanged);
+
     if let Some(secret) = body.secret.as_ref()
         && !is_valid_user_secret(secret)
     {
@@ -265,7 +272,31 @@ pub(super) async fn patch_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::new();
+    if touches_users {
+        touched_sections.push(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 = if touched_sections.is_empty() {
+        current_revision(&shared.config_path).await?
+    } else {
+        save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?
+    };
     drop(_guard);
     match max_unique_ips_change {
         Some(Some(limit)) => shared.ip_tracker.set_user_limit(user, limit).await,

src/config/defaults.rs

@@ -102,7 +102,7 @@ pub(crate) fn default_fake_cert_len() -> usize {
 }
 
 pub(crate) fn default_tls_front_dir() -> String {
-    "/etc/telemt/tlsfront".to_string()
+    "tlsfront".to_string()
 }
 
 pub(crate) fn default_replay_check_len() -> usize {
@@ -568,7 +568,7 @@ pub(crate) fn default_beobachten_flush_secs() -> u64 {
 }
 
 pub(crate) fn default_beobachten_file() -> String {
-    "/etc/telemt/beobachten.txt".to_string()
+    "beobachten.txt".to_string()
 }
 
 pub(crate) fn default_tls_new_session_tickets() -> u8 {

src/ip_tracker.rs

@@ -278,9 +278,11 @@ impl UserIpTracker {
             return Ok(());
         }
 
+        let is_new_ip = !user_recent.contains_key(&ip);
+
         if let Some(limit) = limit {
             let active_limit_reached = user_active.len() >= limit;
-            let recent_limit_reached = user_recent.len() >= limit;
+            let recent_limit_reached = user_recent.len() >= limit && is_new_ip;
             let deny = match mode {
                 UserMaxUniqueIpsMode::ActiveWindow => active_limit_reached,
                 UserMaxUniqueIpsMode::TimeWindow => recent_limit_reached,
@@ -860,4 +862,19 @@ mod tests {
             .unwrap_or(false);
         assert!(!stale_exists);
     }
+
+    #[tokio::test]
+    async fn test_time_window_allows_same_ip_reconnect() {
+        let tracker = UserIpTracker::new();
+        tracker.set_user_limit("test_user", 1).await;
+        tracker
+            .set_limit_policy(UserMaxUniqueIpsMode::TimeWindow, 1)
+            .await;
+
+        let ip1 = test_ipv4(10, 4, 0, 1);
+
+        assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
+        tracker.remove_ip("test_user", ip1).await;
+        assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
+    }
 }

src/maestro/helpers.rs

@@ -1,6 +1,6 @@
 #![allow(clippy::items_after_test_module)]
 
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::time::Duration;
 
 use tokio::sync::watch;
@@ -17,7 +17,7 @@ use crate::transport::middle_proxy::{
 
 pub(crate) fn resolve_runtime_config_path(
     config_path_cli: &str,
-    startup_cwd: &std::path::Path,
+    startup_cwd: &Path,
     config_path_explicit: bool,
 ) -> PathBuf {
     if config_path_explicit {
@@ -46,6 +46,39 @@ pub(crate) fn resolve_runtime_config_path(
     startup_cwd.join("config.toml")
 }
 
+pub(crate) fn resolve_runtime_base_dir(
+    config_path: &Path,
+    startup_cwd: &Path,
+    config_path_explicit: bool,
+    data_path: Option<&Path>,
+) -> PathBuf {
+    if let Some(path) = data_path {
+        return normalize_runtime_dir(path, startup_cwd);
+    }
+
+    if startup_cwd != Path::new("/") {
+        return normalize_runtime_dir(startup_cwd, startup_cwd);
+    }
+
+    if config_path_explicit
+        && let Some(parent) = config_path.parent()
+        && !parent.as_os_str().is_empty()
+    {
+        return normalize_runtime_dir(parent, startup_cwd);
+    }
+
+    PathBuf::from("/etc/telemt")
+}
+
+fn normalize_runtime_dir(path: &Path, startup_cwd: &Path) -> PathBuf {
+    let absolute = if path.is_absolute() {
+        path.to_path_buf()
+    } else {
+        startup_cwd.join(path)
+    };
+    absolute.canonicalize().unwrap_or(absolute)
+}
+
 /// Parsed CLI arguments.
 pub(crate) struct CliArgs {
     pub config_path: String,
@@ -231,9 +264,11 @@ fn print_help() {
 
 #[cfg(test)]
 mod tests {
+    use std::path::{Path, PathBuf};
+
     use super::{
         expected_handshake_close_description, is_expected_handshake_eof, peer_close_description,
-        resolve_runtime_config_path,
+        resolve_runtime_base_dir, resolve_runtime_config_path,
     };
     use crate::error::{ProxyError, StreamError};
 
@@ -304,6 +339,92 @@ mod tests {
         let _ = std::fs::remove_dir(&startup_cwd);
     }
 
+    #[test]
+    fn resolve_runtime_base_dir_prefers_cli_data_path() {
+        let nonce = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_cwd_{nonce}"));
+        let data_path = std::env::temp_dir().join(format!("telemt_runtime_base_data_{nonce}"));
+        std::fs::create_dir_all(&startup_cwd).unwrap();
+        std::fs::create_dir_all(&data_path).unwrap();
+
+        let resolved = resolve_runtime_base_dir(
+            &startup_cwd.join("config.toml"),
+            &startup_cwd,
+            true,
+            Some(&data_path),
+        );
+        assert_eq!(resolved, data_path.canonicalize().unwrap());
+
+        let _ = std::fs::remove_dir(&data_path);
+        let _ = std::fs::remove_dir(&startup_cwd);
+    }
+
+    #[test]
+    fn resolve_runtime_base_dir_uses_working_directory_before_explicit_config_parent() {
+        let nonce = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_start_{nonce}"));
+        let config_dir = std::env::temp_dir().join(format!("telemt_runtime_base_cfg_{nonce}"));
+        std::fs::create_dir_all(&startup_cwd).unwrap();
+        std::fs::create_dir_all(&config_dir).unwrap();
+
+        let resolved =
+            resolve_runtime_base_dir(&config_dir.join("telemt.toml"), &startup_cwd, true, None);
+        assert_eq!(resolved, startup_cwd.canonicalize().unwrap());
+
+        let _ = std::fs::remove_dir(&config_dir);
+        let _ = std::fs::remove_dir(&startup_cwd);
+    }
+
+    #[test]
+    fn resolve_runtime_base_dir_uses_explicit_config_parent_from_root() {
+        let nonce = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let config_dir = std::env::temp_dir().join(format!("telemt_runtime_base_root_cfg_{nonce}"));
+        std::fs::create_dir_all(&config_dir).unwrap();
+
+        let resolved =
+            resolve_runtime_base_dir(&config_dir.join("telemt.toml"), Path::new("/"), true, None);
+        assert_eq!(resolved, config_dir.canonicalize().unwrap());
+
+        let _ = std::fs::remove_dir(&config_dir);
+    }
+
+    #[test]
+    fn resolve_runtime_base_dir_uses_systemd_working_directory_before_etc() {
+        let nonce = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let startup_cwd =
+            std::env::temp_dir().join(format!("telemt_runtime_base_systemd_{nonce}"));
+        std::fs::create_dir_all(&startup_cwd).unwrap();
+
+        let resolved =
+            resolve_runtime_base_dir(&startup_cwd.join("config.toml"), &startup_cwd, false, None);
+        assert_eq!(resolved, startup_cwd.canonicalize().unwrap());
+
+        let _ = std::fs::remove_dir(&startup_cwd);
+    }
+
+    #[test]
+    fn resolve_runtime_base_dir_falls_back_to_etc_from_root() {
+        let resolved = resolve_runtime_base_dir(
+            Path::new("/etc/telemt/config.toml"),
+            Path::new("/"),
+            false,
+            None,
+        );
+        assert_eq!(resolved, PathBuf::from("/etc/telemt"));
+    }
+
     #[test]
     fn expected_handshake_eof_matches_connection_reset() {
         let err = ProxyError::Io(std::io::Error::from(std::io::ErrorKind::ConnectionReset));

src/maestro/mod.rs

@@ -47,7 +47,7 @@ use crate::stats::{ReplayChecker, Stats};
 use crate::stream::BufferPool;
 use crate::transport::UpstreamManager;
 use crate::transport::middle_proxy::MePool;
-use helpers::{parse_cli, resolve_runtime_config_path};
+use helpers::{parse_cli, resolve_runtime_base_dir, resolve_runtime_config_path};
 
 #[cfg(unix)]
 use crate::daemon::{DaemonOptions, PidFile, drop_privileges};
@@ -112,8 +112,51 @@ async fn run_telemt_core(
             std::process::exit(1);
         }
     };
+    if let Some(ref data_path) = data_path
+        && !data_path.is_absolute()
+    {
+        eprintln!(
+            "[telemt] data_path must be absolute: {}",
+            data_path.display()
+        );
+        std::process::exit(1);
+    }
     let mut config_path =
         resolve_runtime_config_path(&config_path_cli, &startup_cwd, config_path_explicit);
+    let runtime_base_dir = resolve_runtime_base_dir(
+        &config_path,
+        &startup_cwd,
+        config_path_explicit,
+        data_path.as_deref(),
+    );
+
+    if !runtime_base_dir.exists()
+        && let Err(e) = std::fs::create_dir_all(&runtime_base_dir)
+    {
+        eprintln!(
+            "[telemt] Can't create runtime directory {}: {}",
+            runtime_base_dir.display(),
+            e
+        );
+        std::process::exit(1);
+    }
+
+    if !runtime_base_dir.is_dir() {
+        eprintln!(
+            "[telemt] Runtime path exists but is not a directory: {}",
+            runtime_base_dir.display()
+        );
+        std::process::exit(1);
+    }
+
+    if let Err(e) = std::env::set_current_dir(&runtime_base_dir) {
+        eprintln!(
+            "[telemt] Can't use runtime directory {}: {}",
+            runtime_base_dir.display(),
+            e
+        );
+        std::process::exit(1);
+    }
 
     let mut config = match ProxyConfig::load(&config_path) {
         Ok(c) => c,
@@ -156,16 +199,15 @@ async fn run_telemt_core(
                         );
                     }
                 } else {
-                    let system_dir = std::path::Path::new("/etc/telemt");
-                    let system_config_path = system_dir.join("telemt.toml");
-                    let startup_config_path = startup_cwd.join("config.toml");
+                    let runtime_config_path = runtime_base_dir.join("telemt.toml");
+                    let fallback_config_path = runtime_base_dir.join("config.toml");
                     let mut persisted = false;
 
                     if let Some(serialized) = serialized.as_ref() {
-                        match std::fs::create_dir_all(system_dir) {
-                            Ok(()) => match std::fs::write(&system_config_path, serialized) {
+                        match std::fs::create_dir_all(&runtime_base_dir) {
+                            Ok(()) => match std::fs::write(&runtime_config_path, serialized) {
                                 Ok(()) => {
-                                    config_path = system_config_path;
+                                    config_path = runtime_config_path;
                                     eprintln!(
                                         "[telemt] Created default config at {}",
                                         config_path.display()
@@ -175,7 +217,7 @@ async fn run_telemt_core(
                                 Err(write_error) => {
                                     eprintln!(
                                         "[telemt] Warning: failed to write default config at {}: {}",
-                                        system_config_path.display(),
+                                        runtime_config_path.display(),
                                         write_error
                                     );
                                 }
@@ -183,16 +225,16 @@ async fn run_telemt_core(
                             Err(create_error) => {
                                 eprintln!(
                                     "[telemt] Warning: failed to create {}: {}",
-                                    system_dir.display(),
+                                    runtime_base_dir.display(),
                                     create_error
                                 );
                             }
                         }
 
                         if !persisted {
-                            match std::fs::write(&startup_config_path, serialized) {
+                            match std::fs::write(&fallback_config_path, serialized) {
                                 Ok(()) => {
-                                    config_path = startup_config_path;
+                                    config_path = fallback_config_path;
                                     eprintln!(
                                         "[telemt] Created default config at {}",
                                         config_path.display()
@@ -202,7 +244,7 @@ async fn run_telemt_core(
                                 Err(write_error) => {
                                     eprintln!(
                                         "[telemt] Warning: failed to write default config at {}: {}",
-                                        startup_config_path.display(),
+                                        fallback_config_path.display(),
                                         write_error
                                     );
                                 }