Add monthly quotas, per-IP limits, user management CLI

Quota system now supports:
- Per-user: daily, weekly, monthly limits
- Per-IP: daily, weekly, monthly limits (abuse prevention)
- Per-IP connection limit
- Max test duration

New CLI flags:
  --monthly-quota, --ip-daily, --ip-weekly, --ip-monthly

User management subcommands:
  btest-server-pro useradd <user> <pass>
  btest-server-pro userdel <user>
  btest-server-pro userlist
  btest-server-pro userset <user> --enabled true/false --daily N --weekly N

New DB tables: ip_usage (per-IP daily tracking)
New methods: get_monthly_usage, get_ip_*_usage, start/end_session,
  delete_user, set_user_enabled, set_user_quota

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Siavash Sameni
2026-04-01 14:58:19 +04:00
parent d2fdc9c6ae
commit d61fdb1b94
3 changed files with 350 additions and 29 deletions

View File

@@ -62,6 +62,22 @@ struct Cli {
#[arg(long = "weekly-quota", default_value_t = 0)]
weekly_quota: u64,
/// Default monthly quota per user in bytes (0 = unlimited)
#[arg(long = "monthly-quota", default_value_t = 0)]
monthly_quota: u64,
/// Daily bandwidth limit per IP in bytes (0 = unlimited)
#[arg(long = "ip-daily", default_value_t = 0)]
ip_daily: u64,
/// Weekly bandwidth limit per IP in bytes (0 = unlimited)
#[arg(long = "ip-weekly", default_value_t = 0)]
ip_weekly: u64,
/// Monthly bandwidth limit per IP in bytes (0 = unlimited)
#[arg(long = "ip-monthly", default_value_t = 0)]
ip_monthly: u64,
/// Maximum concurrent connections per IP (0 = unlimited)
#[arg(long = "max-conn-per-ip", default_value_t = 5)]
max_conn_per_ip: u32,
@@ -85,6 +101,46 @@ struct Cli {
/// Verbose logging
#[arg(short = 'v', long = "verbose", action = clap::ArgAction::Count)]
verbose: u8,
/// User management subcommand
#[command(subcommand)]
command: Option<UserCommand>,
}
#[derive(clap::Subcommand, Debug)]
enum UserCommand {
/// Add a user
#[command(name = "useradd")]
UserAdd {
/// Username
username: String,
/// Password
password: String,
},
/// Delete a user
#[command(name = "userdel")]
UserDel {
/// Username
username: String,
},
/// List all users
#[command(name = "userlist")]
UserList,
/// Enable/disable a user
#[command(name = "userset")]
UserSet {
/// Username
username: String,
/// Enable (true/false)
#[arg(long)]
enabled: Option<bool>,
/// Daily quota in bytes
#[arg(long)]
daily: Option<i64>,
/// Weekly quota in bytes
#[arg(long)]
weekly: Option<i64>,
},
}
#[tokio::main]
@@ -119,10 +175,60 @@ async fn main() -> anyhow::Result<()> {
}
// Initialize user database
tracing::info!("Opening user database: {}", cli.users_db);
let db = user_db::UserDb::open(&cli.users_db)?;
db.ensure_tables()?;
tracing::info!("User database ready ({} users)", db.user_count()?);
// Handle user management subcommands (exit after)
if let Some(cmd) = &cli.command {
match cmd {
UserCommand::UserAdd { username, password } => {
db.add_user(username, password)?;
println!("User '{}' added.", username);
return Ok(());
}
UserCommand::UserDel { username } => {
if db.delete_user(username)? {
println!("User '{}' deleted.", username);
} else {
println!("User '{}' not found.", username);
}
return Ok(());
}
UserCommand::UserList => {
let users = db.list_users()?;
if users.is_empty() {
println!("No users.");
} else {
println!("{:<20} {:<10} {:<15} {:<15}", "USERNAME", "ENABLED", "DAILY_QUOTA", "WEEKLY_QUOTA");
println!("{}", "-".repeat(60));
for u in &users {
println!("{:<20} {:<10} {:<15} {:<15}",
u.username,
if u.enabled { "yes" } else { "no" },
if u.daily_quota == 0 { "default".to_string() } else { format!("{}B", u.daily_quota) },
if u.weekly_quota == 0 { "default".to_string() } else { format!("{}B", u.weekly_quota) },
);
}
}
return Ok(());
}
UserCommand::UserSet { username, enabled, daily, weekly } => {
if let Some(e) = enabled {
db.set_user_enabled(username, *e)?;
println!("User '{}' enabled={}", username, e);
}
if daily.is_some() || weekly.is_some() {
let d = daily.unwrap_or(0);
let w = weekly.unwrap_or(0);
db.set_user_quota(username, d, w, 0)?;
println!("User '{}' quota: daily={}, weekly={}", username, d, w);
}
return Ok(());
}
}
}
tracing::info!("User database: {} ({} users)", cli.users_db, db.user_count()?);
// Initialize LDAP if configured
if let Some(ref url) = cli.ldap_url {
@@ -130,19 +236,30 @@ async fn main() -> anyhow::Result<()> {
}
// Initialize quota manager
let quota_mgr = quota::QuotaManager::new(
let _quota_mgr = quota::QuotaManager::new(
db.clone(),
cli.daily_quota,
cli.weekly_quota,
cli.monthly_quota,
cli.ip_daily,
cli.ip_weekly,
cli.ip_monthly,
cli.max_conn_per_ip,
cli.max_duration,
);
let fmt_q = |v: u64| if v == 0 { "unlimited".to_string() } else { format!("{}B", v) };
tracing::info!(
"Quotas: daily={}, weekly={}, max_conn_per_ip={}, max_duration={}s",
if cli.daily_quota == 0 { "unlimited".to_string() } else { format!("{}", cli.daily_quota) },
if cli.weekly_quota == 0 { "unlimited".to_string() } else { format!("{}", cli.weekly_quota) },
cli.max_conn_per_ip,
cli.max_duration,
"User quotas: daily={}, weekly={}, monthly={}",
fmt_q(cli.daily_quota), fmt_q(cli.weekly_quota), fmt_q(cli.monthly_quota),
);
tracing::info!(
"IP quotas: daily={}, weekly={}, monthly={}",
fmt_q(cli.ip_daily), fmt_q(cli.ip_weekly), fmt_q(cli.ip_monthly),
);
tracing::info!(
"Limits: max_conn_per_ip={}, max_duration={}s",
cli.max_conn_per_ip, cli.max_duration,
);
tracing::info!("btest-server-pro starting on port {}", cli.port);