From 0c51dc2ae6fd08a81bb5963f4be782dea28b2d0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Wed, 25 Mar 2026 15:46:16 +0100 Subject: [PATCH 1/4] ACL rule generator --- tools/defguard_generator/Cargo.toml | 9 +++--- tools/defguard_generator/src/acl_rules.rs | 32 +++++++++++++++++++ tools/defguard_generator/src/lib.rs | 5 +-- tools/defguard_generator/src/main.rs | 13 ++++++-- .../src/vpn_session_stats.rs | 15 +++++---- 5 files changed, 57 insertions(+), 17 deletions(-) create mode 100644 tools/defguard_generator/src/acl_rules.rs diff --git a/tools/defguard_generator/Cargo.toml b/tools/defguard_generator/Cargo.toml index 06255adc2b..1b89fe9bc9 100644 --- a/tools/defguard_generator/Cargo.toml +++ b/tools/defguard_generator/Cargo.toml @@ -8,14 +8,13 @@ repository = "https://github.com/DefGuard/defguard" rust-version = "1.87.0" [dependencies] -defguard_common = { workspace = true } +anyhow = { workspace = true } chrono = { workspace = true } -rand = { workspace = true } clap = { workspace = true, features = ["derive"] } +defguard_common = { workspace = true } +defguard_core = { workspace = true } +rand = { workspace = true } sqlx = { workspace = true, features = ["postgres", "runtime-tokio-native-tls", "chrono"] } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } -anyhow = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true, features = ["env-filter"] } - -[dev-dependencies] diff --git a/tools/defguard_generator/src/acl_rules.rs b/tools/defguard_generator/src/acl_rules.rs new file mode 100644 index 0000000000..9121aac231 --- /dev/null +++ b/tools/defguard_generator/src/acl_rules.rs @@ -0,0 +1,32 @@ +use anyhow::Result; +use defguard_core::enterprise::db::models::acl::{AclRule, RuleState}; +use sqlx::{PgPool, query}; + +pub async fn generate_acl_rules(pool: PgPool, num_rules: u32) -> Result<()> { + truncate_with_restart(&pool).await?; + + for index in 1..num_rules { + let mut acl_rule = AclRule::default(); + acl_rule.name = format!("Generated {index}"); + acl_rule.state = RuleState::Applied; + acl_rule.all_locations = true; + acl_rule.allow_all_users = true; + acl_rule.allow_all_groups = true; + acl_rule.allow_all_network_devices = true; + // acl_rule.any_port = false; + // acl_rule.ports = vec![(index as i32..=index as i32).into()]; + acl_rule.save(&pool).await?; + } + + Ok(()) +} + +/// Remove all records from sessions and stats tables. +/// This also resets the auto-incrementing sequences. +async fn truncate_with_restart(pool: &PgPool) -> Result<()> { + query("TRUNCATE aclrule RESTART IDENTITY CASCADE") + .execute(pool) + .await?; + + Ok(()) +} diff --git a/tools/defguard_generator/src/lib.rs b/tools/defguard_generator/src/lib.rs index 38c1a426e9..4cf3f6288d 100644 --- a/tools/defguard_generator/src/lib.rs +++ b/tools/defguard_generator/src/lib.rs @@ -1,3 +1,4 @@ -pub mod user_devices; -pub mod users; +pub mod acl_rules; +mod user_devices; +mod users; pub mod vpn_session_stats; diff --git a/tools/defguard_generator/src/main.rs b/tools/defguard_generator/src/main.rs index 3c6feea0e9..a355e27056 100644 --- a/tools/defguard_generator/src/main.rs +++ b/tools/defguard_generator/src/main.rs @@ -1,8 +1,9 @@ use anyhow::Result; use clap::{Parser, Subcommand}; use defguard_common::db::{Id, init_db}; -use defguard_generator::vpn_session_stats::{ - VpnSessionGeneratorConfig, generate_vpn_session_stats, +use defguard_generator::{ + acl_rules::generate_acl_rules, + vpn_session_stats::{VpnSessionGeneratorConfig, generate_vpn_session_stats}, }; use tracing_subscriber::EnvFilter; @@ -35,7 +36,7 @@ enum Commands { #[arg(long)] location_id: Id, #[arg(long)] - num_users: u16, + num_users: usize, #[arg(long)] devices_per_user: u8, #[arg(long)] @@ -47,6 +48,11 @@ enum Commands { #[arg(long, default_value_t = 1000)] stats_batch_size: u16, }, + /// Generates ACL rules + AclRules { + #[arg(long)] + num_rules: u32, + }, } #[tokio::main] @@ -92,6 +98,7 @@ async fn main() -> Result<()> { generate_vpn_session_stats(pool, config).await?; } + Commands::AclRules { num_rules } => generate_acl_rules(pool, num_rules).await?, }; Ok(()) diff --git a/tools/defguard_generator/src/vpn_session_stats.rs b/tools/defguard_generator/src/vpn_session_stats.rs index ff9519b84e..21749ed663 100644 --- a/tools/defguard_generator/src/vpn_session_stats.rs +++ b/tools/defguard_generator/src/vpn_session_stats.rs @@ -13,7 +13,7 @@ use defguard_common::db::{ }, }; use rand::{Rng, rngs::ThreadRng}; -use sqlx::{PgConnection, PgPool, QueryBuilder}; +use sqlx::{PgConnection, PgPool, QueryBuilder, query}; use tracing::{debug, info}; use crate::{user_devices::prepare_user_devices, users::prepare_users}; @@ -24,7 +24,7 @@ const HANDSHAKE_INTERVAL: Duration = Duration::minutes(2); #[derive(Debug)] pub struct VpnSessionGeneratorConfig { pub location_id: Id, - pub num_users: u16, + pub num_users: usize, pub devices_per_user: u8, pub sessions_per_device: u8, pub no_truncate: bool, @@ -53,7 +53,7 @@ pub async fn generate_vpn_session_stats( let gateway = prepare_gateway(&pool, location.id).await?; // prepare requested number of users - let user_count = config.num_users as usize; + let user_count = config.num_users; let users = prepare_users(&pool, &mut rng, user_count).await?; // generate sessions for each user @@ -148,10 +148,10 @@ pub async fn generate_vpn_session_stats( Ok(()) } -/// Remove all records from sessions & stats tables. -/// This also resets the auto-incrementing sequences +/// Remove all records from sessions and stats tables. +/// This also resets the auto-incrementing sequences. async fn truncate_with_restart(pool: &PgPool) -> Result<()> { - sqlx::query("TRUNCATE TABLE vpn_client_session RESTART IDENTITY CASCADE") + query("TRUNCATE vpn_client_session RESTART IDENTITY CASCADE") .execute(pool) .await?; @@ -248,7 +248,8 @@ async fn insert_stats_batch( } let mut query_builder = QueryBuilder::new( - "INSERT INTO vpn_session_stats (session_id, gateway_id, collected_at, latest_handshake, endpoint, total_upload, total_download, upload_diff, download_diff) ", + "INSERT INTO vpn_session_stats (session_id, gateway_id, collected_at, latest_handshake, \ + endpoint, total_upload, total_download, upload_diff, download_diff) ", ); query_builder.push_values(stats_batch, |mut b, stats| { From e6641ce2fa8a64b3ab6682641b6f13cebfa1c31f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Wed, 25 Mar 2026 15:47:41 +0100 Subject: [PATCH 2/4] Updates --- Cargo.lock | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ad790f5154..355c4febfa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1505,6 +1505,7 @@ dependencies = [ "chrono", "clap", "defguard_common", + "defguard_core", "rand 0.8.5", "sqlx", "tokio", @@ -6098,18 +6099,18 @@ checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" [[package]] name = "struct-patch" -version = "0.10.4" +version = "0.10.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e986d2cf6e819bd843319120453d837dfdfa31497c3fee4cefa614b2d182d8c" +checksum = "16d4caaaccd69c9b56c5f5b33d4dca462464d3275230e4d2d3739ba6d4bf5bcb" dependencies = [ "struct-patch-derive", ] [[package]] name = "struct-patch-derive" -version = "0.10.4" +version = "0.10.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68c6387c1c7b53060605101b63d93edca618c6cf7ce61839f2ec2a527419fdb5" +checksum = "1671c6f0992b1b4cb4f5f8ea4a58f9a5f7f895a7638ef9690633dcec0aa67944" dependencies = [ "proc-macro2", "quote", From 019419492accd823891328d0d3036edf83cb4623 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Wed, 25 Mar 2026 15:48:30 +0100 Subject: [PATCH 3/4] Remove commented-out section --- tools/defguard_generator/src/acl_rules.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/tools/defguard_generator/src/acl_rules.rs b/tools/defguard_generator/src/acl_rules.rs index 9121aac231..4751e3fb83 100644 --- a/tools/defguard_generator/src/acl_rules.rs +++ b/tools/defguard_generator/src/acl_rules.rs @@ -13,8 +13,6 @@ pub async fn generate_acl_rules(pool: PgPool, num_rules: u32) -> Result<()> { acl_rule.allow_all_users = true; acl_rule.allow_all_groups = true; acl_rule.allow_all_network_devices = true; - // acl_rule.any_port = false; - // acl_rule.ports = vec![(index as i32..=index as i32).into()]; acl_rule.save(&pool).await?; } From b522263fa81cc3e2e810fa48ac533bce8ad50184 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Wed, 25 Mar 2026 15:50:31 +0100 Subject: [PATCH 4/4] Fix count --- tools/defguard_generator/src/acl_rules.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/defguard_generator/src/acl_rules.rs b/tools/defguard_generator/src/acl_rules.rs index 4751e3fb83..9766af7e7f 100644 --- a/tools/defguard_generator/src/acl_rules.rs +++ b/tools/defguard_generator/src/acl_rules.rs @@ -5,7 +5,7 @@ use sqlx::{PgPool, query}; pub async fn generate_acl_rules(pool: PgPool, num_rules: u32) -> Result<()> { truncate_with_restart(&pool).await?; - for index in 1..num_rules { + for index in 0..num_rules { let mut acl_rule = AclRule::default(); acl_rule.name = format!("Generated {index}"); acl_rule.state = RuleState::Applied;