upstash_ratelimit/
lib.rs

1#![allow(unused)]
2
3pub mod builder;
4
5use anyhow::{anyhow, Result};
6use redis::Client as RedisClient;
7use std::time::{Duration, SystemTime, UNIX_EPOCH};
8
9pub enum Limiter {
10    FixedWindow {
11        tokens: u64,
12        window: Duration,
13    },
14    SlidingLogs {
15        tokens: u64,
16        window: Duration,
17    },
18    SlidingWindow {
19        tokens: u64,
20        window: Duration,
21    },
22    TokenBucket {
23        refill_rate: u64,
24        interval: Duration,
25        max_tokens: u64,
26    },
27}
28
29pub enum Response {
30    Success {
31        limit: u64,
32        remaining: u64,
33        reset: Duration,
34    },
35    Failure {
36        limit: u64,
37        reset: Duration,
38    },
39}
40
41pub struct RateLimit {
42    redis: RedisClient,
43    limiter: Limiter,
44    prefix: String,
45    timeout: Option<Duration>,
46    analytics: bool,
47}
48
49impl RateLimit {
50    /// Creates a RateLimit builder instance.
51    pub fn builder() -> builder::RateLimitBuilder {
52        builder::RateLimitBuilder::default()
53    }
54
55    /// Apply limiting based on a given unique identifier.
56    ///
57    /// The identifier could be a user id,
58    pub fn limit<T>(&self, identifier: T) -> Result<Response>
59    where
60        T: Into<String> + std::fmt::Display,
61    {
62        return match self.limiter {
63            Limiter::FixedWindow { tokens, window } => {
64                let script = redis::Script::new(
65                    r#"
66                    local key     = KEYS[1]
67                    local window  = ARGV[1]
68
69                    local r = redis.call("INCR", key)
70                    if r == 1 then
71                    -- The first time this key is set, the value will be 1.
72                    -- So we only need the expire command once
73                    redis.call("PEXPIRE", key, window)
74                    end
75
76                    return r
77                "#,
78                );
79
80                let window_duration = window.as_millis();
81
82                let since_epoch = SystemTime::now()
83                    .duration_since(UNIX_EPOCH)
84                    .expect("Time went backwards")
85                    .as_millis();
86                let bucket = since_epoch / window_duration;
87                let key = format!("{}:{}", identifier, bucket);
88
89                // todo: implement in-memory chache
90
91                let mut redis = self.redis.get_connection()?;
92                let tokens_used: u64 = script
93                    .key(key)
94                    .arg(format!("{}", window_duration))
95                    .invoke(&mut redis)?;
96
97                let reset = Duration::from_millis(((bucket + 1) * window_duration) as u64);
98
99                if tokens_used <= tokens {
100                    Ok(Response::Success {
101                        limit: tokens,
102                        remaining: tokens - tokens_used,
103                        reset,
104                    })
105                } else {
106                    Ok(Response::Failure {
107                        limit: tokens,
108                        reset,
109                    })
110                }
111            }
112            _ => unimplemented!("Limiter method not implemented."),
113        };
114    }
115}
upstash_ratelimit/lib.rs

upstash_ratelimit/
lib.rs
