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:
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user