use derusted::{JwtClaims, JwtValidator, RateLimiter, RateLimiterConfig};
use std::sync::Arc;
use jsonwebtoken::{encode, EncodingKey, Header};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct TieredClaims {
#[serde(default)]
tier: Option<String>,
#[serde(default)]
rate_limit_per_minute: Option<usize>,
}
fn get_tier_rate_limit(tier: Option<&str>) -> usize {
match tier {
Some("enterprise") => 100_000,
Some("pro") => 10_000,
Some("free") | None => 100,
Some(other) => {
eprintln!("Unknown tier '{}', using free tier limit", other);
100
}
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = RateLimiterConfig {
requests_per_minute: 100, burst_size: 20, bucket_ttl_seconds: 3600, max_buckets: 100_000, };
let rate_limiter = Arc::new(RateLimiter::new(config));
let secret = "your-256-bit-secret-key-here-min-32-chars";
let validator: JwtValidator<TieredClaims> = JwtValidator::new(
secret.to_string(),
"HS256".to_string(),
"us-east".to_string(),
None,
None,
)?;
let free_user = create_token(secret, "free_user_1", "free", None)?;
let pro_user = create_token(secret, "pro_user_1", "pro", None)?;
let enterprise_user = create_token(secret, "enterprise_user_1", "enterprise", None)?;
let vip_user = create_token(secret, "vip_user_1", "enterprise", Some(500_000))?;
println!("=== Tiered Rate Limiting Demo ===\n");
for (name, token, expected_limit) in [
("Free User", &free_user, 100),
("Pro User", &pro_user, 10_000),
("Enterprise User", &enterprise_user, 100_000),
("VIP User (custom)", &vip_user, 500_000),
] {
println!("--- {} ---", name);
let claims = validator.validate(&format!("Bearer {}", token))?;
let effective_limit = claims
.extra
.rate_limit_per_minute
.unwrap_or_else(|| get_tier_rate_limit(claims.extra.tier.as_deref()));
println!("Token ID: {}", claims.token_id);
println!("Tier: {:?}", claims.extra.tier);
println!("Custom limit: {:?}", claims.extra.rate_limit_per_minute);
println!("Effective limit: {} req/min", effective_limit);
assert_eq!(effective_limit, expected_limit);
let mut success_count = 0;
let mut fail_count = 0;
for i in 0..21 {
let result = rate_limiter
.check_limit_with_override(&claims.token_id, Some(effective_limit))
.await;
match result {
Ok(()) => success_count += 1,
Err(_) => {
fail_count += 1;
if i == 20 {
println!("Request {} rate limited (expected after burst)", i + 1);
}
}
}
}
println!("Successful requests: {}", success_count);
println!("Rate limited: {}", fail_count);
println!();
}
println!("=== Concurrent Users Demo ===\n");
let users: Vec<(&str, String)> = vec![
("free_a", create_token(secret, "free_a", "free", None)?),
("free_b", create_token(secret, "free_b", "free", None)?),
("pro_a", create_token(secret, "pro_a", "pro", None)?),
("enterprise_a", create_token(secret, "enterprise_a", "enterprise", None)?),
];
let mut handles = vec![];
for (name, token) in users {
let rate_limiter = Arc::clone(&rate_limiter);
let validator_secret = secret.to_string();
let name = name.to_string();
let token = token.clone();
let handle = tokio::spawn(async move {
let validator: JwtValidator<TieredClaims> = JwtValidator::new(
validator_secret,
"HS256".to_string(),
"us-east".to_string(),
None,
None,
)
.unwrap();
let claims = validator.validate(&format!("Bearer {}", token)).unwrap();
let limit = get_tier_rate_limit(claims.extra.tier.as_deref());
let mut success = 0;
for _ in 0..10 {
if rate_limiter
.check_limit_with_override(&claims.token_id, Some(limit))
.await
.is_ok()
{
success += 1;
}
}
(name, success)
});
handles.push(handle);
}
for handle in handles {
let (name, success) = handle.await?;
println!("{}: {} successful requests", name, success);
}
println!("\n=== Rate Limiter Stats ===");
let stats = rate_limiter.get_stats().await;
println!("Active token buckets: {}", stats.active_tokens);
println!("Max buckets: {}", stats.max_tokens);
println!("Default requests/min: {}", stats.requests_per_minute);
println!("Burst size: {}", stats.burst_size);
Ok(())
}
fn create_token(
secret: &str,
token_id: &str,
tier: &str,
custom_limit: Option<usize>,
) -> Result<String, Box<dyn std::error::Error>> {
let claims: JwtClaims<TieredClaims> = JwtClaims {
token_id: token_id.to_string(),
user_id: 1,
allowed_regions: vec!["us-east".to_string()],
exp: (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp(),
iat: chrono::Utc::now().timestamp(),
iss: None,
aud: None,
extra: TieredClaims {
tier: Some(tier.to_string()),
rate_limit_per_minute: custom_limit,
},
};
let token = encode(
&Header::default(),
&claims,
&EncodingKey::from_secret(secret.as_bytes()),
)?;
Ok(token)
}