openauth_fred/
store.rs

1use fred::clients::Client;
2use fred::interfaces::ClientLike;
3use fred::prelude::{Builder, Config};
4use fred::types::scripts::Script;
5use openauth_core::error::OpenAuthError;
6use openauth_core::options::{
7    RateLimitConsumeInput, RateLimitDecision, RateLimitFuture, RateLimitStore,
8};
9
10use crate::script::{parse_rate_limit_script_result, RATE_LIMIT_SCRIPT};
11use crate::{normalize_fred_url, FredRateLimitOptions};
12
13#[derive(Clone)]
14pub struct FredRateLimitStore {
15    client: Client,
16    options: FredRateLimitOptions,
17    script: Script,
18}
19
20impl FredRateLimitStore {
21    pub async fn connect(url: &str) -> Result<Self, OpenAuthError> {
22        Self::connect_with_options(url, FredRateLimitOptions::default()).await
23    }
24
25    pub async fn connect_redis(url: &str) -> Result<Self, OpenAuthError> {
26        Self::connect(url).await
27    }
28
29    pub async fn connect_valkey(url: &str) -> Result<Self, OpenAuthError> {
30        Self::connect(url).await
31    }
32
33    pub async fn connect_with_options(
34        url: &str,
35        options: FredRateLimitOptions,
36    ) -> Result<Self, OpenAuthError> {
37        let url = normalize_fred_url(url);
38        let config = Config::from_url(url.as_ref())
39            .map_err(|error| OpenAuthError::Adapter(error.to_string()))?;
40        let client = Builder::from_config(config)
41            .build()
42            .map_err(|error| OpenAuthError::Adapter(error.to_string()))?;
43        client
44            .init()
45            .await
46            .map_err(|error| OpenAuthError::Adapter(error.to_string()))?;
47        Ok(Self::new(client, options))
48    }
49
50    pub fn new(client: Client, options: FredRateLimitOptions) -> Self {
51        Self {
52            client,
53            options,
54            script: Script::from_lua(RATE_LIMIT_SCRIPT),
55        }
56    }
57
58    fn key(&self, key: &str) -> String {
59        format!("{}rate-limit:{key}", self.options.key_prefix)
60    }
61}
62
63impl RateLimitStore for FredRateLimitStore {
64    fn consume<'a>(&'a self, input: RateLimitConsumeInput) -> RateLimitFuture<'a> {
65        Box::pin(async move {
66            let redis_key = self.key(&input.key);
67            let window_ms = input.rule.window.saturating_mul(1000);
68            let result = self
69                .script
70                .evalsha_with_reload(
71                    &self.client,
72                    vec![redis_key],
73                    vec![
74                        input.now_ms.to_string(),
75                        window_ms.to_string(),
76                        input.rule.max.to_string(),
77                    ],
78                )
79                .await
80                .map_err(|error| OpenAuthError::Adapter(error.to_string()))?;
81            let result = parse_rate_limit_script_result(result)?;
82            let retry_ms = result
83                .last_request
84                .saturating_add(window_ms as i64)
85                .saturating_sub(input.now_ms)
86                .max(0);
87            Ok(RateLimitDecision {
88                permitted: result.permitted,
89                retry_after: if result.permitted {
90                    0
91                } else {
92                    ceil_millis_to_seconds(retry_ms)
93                },
94                limit: input.rule.max,
95                remaining: input.rule.max.saturating_sub(result.count),
96                reset_after: ceil_millis_to_seconds(retry_ms),
97            })
98        })
99    }
100}
101
102fn ceil_millis_to_seconds(milliseconds: i64) -> u64 {
103    if milliseconds <= 0 {
104        return 0;
105    }
106    ((milliseconds as u64).saturating_add(999)) / 1000
107}
openauth_fred/store.rs

openauth_fred/
store.rs
