robinpath_modules/modules/

ratelimit_mod.rs

use robinpath::{RobinPath, Value};
use std::sync::{Arc, Mutex};
use std::collections::HashMap;

enum Limiter {
    TokenBucket {
        tokens: f64,
        max_tokens: f64,
        refill_rate: f64,
        last_refill: u64,
    },
    SlidingWindow {
        requests: Vec<u64>,
        window_ms: u64,
        max_requests: u64,
    },
    FixedWindow {
        count: u64,
        window_ms: u64,
        max_requests: u64,
        window_start: u64,
    },
}

fn now_ms() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

pub fn register(rp: &mut RobinPath) {
    let state = Arc::new(Mutex::new(HashMap::<String, Limiter>::new()));

    // ratelimit.create name type options → config
    let s = state.clone();
    rp.register_builtin("ratelimit.create", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
        let limiter_type = args.get(1).map(|v| v.to_display_string()).unwrap_or_else(|| "token-bucket".to_string());
        let opts = args.get(2).cloned().unwrap_or(Value::Null);
        let now = now_ms();
        let limiter = match limiter_type.as_str() {
            "sliding-window" => {
                let window_ms = get_opt_num(&opts, "windowMs", 60000.0) as u64;
                let max_requests = get_opt_num(&opts, "maxRequests", 100.0) as u64;
                Limiter::SlidingWindow { requests: Vec::new(), window_ms, max_requests }
            }
            "fixed-window" => {
                let window_ms = get_opt_num(&opts, "windowMs", 60000.0) as u64;
                let max_requests = get_opt_num(&opts, "maxRequests", 100.0) as u64;
                Limiter::FixedWindow { count: 0, window_ms, max_requests, window_start: now }
            }
            _ => {
                let max_tokens = get_opt_num(&opts, "maxTokens", 10.0);
                let refill_rate = get_opt_num(&opts, "refillRate", 1.0);
                Limiter::TokenBucket { tokens: max_tokens, max_tokens, refill_rate, last_refill: now }
            }
        };
        s.lock().unwrap().insert(name.clone(), limiter);
        let mut obj = indexmap::IndexMap::new();
        obj.insert("name".to_string(), Value::String(name));
        obj.insert("type".to_string(), Value::String(limiter_type));
        Ok(Value::Object(obj))
    });

    // ratelimit.acquire name count? → bool
    let s = state.clone();
    rp.register_builtin("ratelimit.acquire", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
        let count = args.get(1).map(|v| v.to_number()).unwrap_or(1.0);
        let mut limiters = s.lock().unwrap();
        let limiter = limiters.get_mut(&name)
            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
        let now = now_ms();
        let allowed = match limiter {
            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
                let elapsed = (now - *last_refill) as f64 / 1000.0;
                *tokens = (*tokens + elapsed * *refill_rate).min(*max_tokens);
                *last_refill = now;
                if *tokens >= count {
                    *tokens -= count;
                    true
                } else {
                    false
                }
            }
            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
                requests.retain(|&t| now - t < *window_ms);
                if (requests.len() as u64) < *max_requests {
                    for _ in 0..count as u64 {
                        requests.push(now);
                    }
                    true
                } else {
                    false
                }
            }
            Limiter::FixedWindow { count: c, window_ms, max_requests, window_start } => {
                if now - *window_start >= *window_ms {
                    *c = 0;
                    *window_start = now;
                }
                if *c + count as u64 <= *max_requests {
                    *c += count as u64;
                    true
                } else {
                    false
                }
            }
        };
        Ok(Value::Bool(allowed))
    });

    // ratelimit.check name → bool
    let s = state.clone();
    rp.register_builtin("ratelimit.check", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
        let mut limiters = s.lock().unwrap();
        let limiter = limiters.get_mut(&name)
            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
        let now = now_ms();
        let allowed = match limiter {
            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
                let elapsed = (now - *last_refill) as f64 / 1000.0;
                let current = (*tokens + elapsed * *refill_rate).min(*max_tokens);
                current >= 1.0
            }
            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
                let active = requests.iter().filter(|&&t| now - t < *window_ms).count() as u64;
                active < *max_requests
            }
            Limiter::FixedWindow { count, window_ms, max_requests, window_start } => {
                (if now - *window_start >= *window_ms { 0 } else { *count }) < *max_requests
            }
        };
        Ok(Value::Bool(allowed))
    });

    // ratelimit.remaining name → number
    let s = state.clone();
    rp.register_builtin("ratelimit.remaining", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
        let mut limiters = s.lock().unwrap();
        let limiter = limiters.get_mut(&name)
            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
        let now = now_ms();
        let remaining = match limiter {
            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
                let elapsed = (now - *last_refill) as f64 / 1000.0;
                (*tokens + elapsed * *refill_rate).min(*max_tokens).floor()
            }
            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
                let active = requests.iter().filter(|&&t| now - t < *window_ms).count() as f64;
                (*max_requests as f64 - active).max(0.0)
            }
            Limiter::FixedWindow { count, window_ms, max_requests, window_start } => {
                let c = if now - *window_start >= *window_ms { 0 } else { *count };
                (*max_requests - c) as f64
            }
        };
        Ok(Value::Number(remaining))
    });

    // ratelimit.reset name → bool
    let s = state.clone();
    rp.register_builtin("ratelimit.reset", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
        let mut limiters = s.lock().unwrap();
        let limiter = limiters.get_mut(&name)
            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
        let now = now_ms();
        match limiter {
            Limiter::TokenBucket { tokens, max_tokens, last_refill, .. } => {
                *tokens = *max_tokens;
                *last_refill = now;
            }
            Limiter::SlidingWindow { requests, .. } => { requests.clear(); }
            Limiter::FixedWindow { count, window_start, .. } => {
                *count = 0;
                *window_start = now;
            }
        }
        Ok(Value::Bool(true))
    });

    // ratelimit.status name → {name, remaining, ...}
    let s = state.clone();
    rp.register_builtin("ratelimit.status", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
        let mut limiters = s.lock().unwrap();
        let limiter = limiters.get_mut(&name)
            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
        let now = now_ms();
        let mut obj = indexmap::IndexMap::new();
        obj.insert("name".to_string(), Value::String(name));
        match limiter {
            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
                let elapsed = (now - *last_refill) as f64 / 1000.0;
                let current = (*tokens + elapsed * *refill_rate).min(*max_tokens);
                obj.insert("type".to_string(), Value::String("token-bucket".to_string()));
                obj.insert("tokens".to_string(), Value::Number(current.floor()));
                obj.insert("maxTokens".to_string(), Value::Number(*max_tokens));
                obj.insert("refillRate".to_string(), Value::Number(*refill_rate));
            }
            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
                let active = requests.iter().filter(|&&t| now - t < *window_ms).count() as f64;
                obj.insert("type".to_string(), Value::String("sliding-window".to_string()));
                obj.insert("used".to_string(), Value::Number(active));
                obj.insert("remaining".to_string(), Value::Number((*max_requests as f64 - active).max(0.0)));
                obj.insert("maxRequests".to_string(), Value::Number(*max_requests as f64));
            }
            Limiter::FixedWindow { count, window_ms, max_requests, window_start } => {
                let c = if now - *window_start >= *window_ms { 0 } else { *count };
                obj.insert("type".to_string(), Value::String("fixed-window".to_string()));
                obj.insert("used".to_string(), Value::Number(c as f64));
                obj.insert("remaining".to_string(), Value::Number((*max_requests - c) as f64));
                obj.insert("maxRequests".to_string(), Value::Number(*max_requests as f64));
            }
        }
        Ok(Value::Object(obj))
    });

    // ratelimit.destroy name → bool
    let s = state.clone();
    rp.register_builtin("ratelimit.destroy", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
        Ok(Value::Bool(s.lock().unwrap().remove(&name).is_some()))
    });
}

fn get_opt_num(opts: &Value, key: &str, default: f64) -> f64 {
    if let Value::Object(obj) = opts {
        obj.get(key).map(|v| v.to_number()).unwrap_or(default)
    } else {
        default
    }
}