apikeys_rs/limiters/
redis_limiter.rs

1use async_trait::async_trait;
2use redis::{AsyncCommands, Client, RedisError};
3
4use crate::{
5    errors::ApiKeyLimiterError,
6    traits::ApiKeyLimiter,
7    types::{ApiKey, ApiKeyLimit},
8};
9
10#[derive(Clone)]
11pub struct RedisLimiter {
12    redis_client: Client,
13}
14
15impl RedisLimiter {
16    pub async fn new(uri: &str) -> Result<Self, RedisError> {
17        tracing::debug!("Creating redis client from uri: {}", uri);
18        let redis_client = Client::open(uri)?;
19        Ok(Self { redis_client })
20    }
21}
22
23#[async_trait]
24impl ApiKeyLimiter for RedisLimiter {
25    async fn use_key(&self, api_key: &ApiKey) -> Result<(), ApiKeyLimiterError> {
26        match api_key.limits.max_reads_per_minute {
27            ApiKeyLimit::Limited(max_reads_per_minute) => {
28                let mut connection = self.redis_client.get_async_connection().await?;
29
30                let key = format!("{}_read_count", api_key.key);
31
32                match connection.exists(&key).await? {
33                    true => {}
34                    false => {
35                        let _: Result<String, RedisError> = connection.set(&key, 0).await;
36
37                        let _: Result<String, RedisError> = connection.expire(&key, 60).await;
38                    }
39                }
40
41                let result: i32 = connection.incr(&key, 1).await?;
42
43                if result > max_reads_per_minute as i32 {
44                    return Err(ApiKeyLimiterError::RateLimitExceeded);
45                }
46            }
47            ApiKeyLimit::Unlimited => {}
48        }
49
50        Ok(())
51    }
52}
53
54impl From<RedisError> for ApiKeyLimiterError {
55    fn from(error: RedisError) -> Self {
56        ApiKeyLimiterError::Other(error.to_string())
57    }
58}
apikeys_rs/limiters/redis_limiter.rs

apikeys_rs/limiters/
redis_limiter.rs
