Skip to main content

openauth_fred/
script.rs

1use fred::types::Value;
2use openauth_core::error::OpenAuthError;
3
4#[derive(Debug, Clone, PartialEq, Eq)]
5pub struct RateLimitScriptResult {
6    pub permitted: bool,
7    pub count: u64,
8    pub last_request: i64,
9}
10
11pub(crate) const RATE_LIMIT_SCRIPT: &str = r#"
12local key = KEYS[1]
13local now = tonumber(ARGV[1])
14local window = tonumber(ARGV[2])
15local max = tonumber(ARGV[3])
16
17local data = redis.call("HMGET", key, "count", "last_request")
18local count = tonumber(data[1])
19local last_request = tonumber(data[2])
20
21if count == nil or last_request == nil or (now - last_request) > window then
22  redis.call("HSET", key, "count", 1, "last_request", now)
23  redis.call("PEXPIRE", key, window)
24  return {1, 1, now}
25end
26
27if count >= max then
28  redis.call("PEXPIRE", key, window)
29  return {0, count, last_request}
30end
31
32count = count + 1
33redis.call("HSET", key, "count", count, "last_request", now)
34redis.call("PEXPIRE", key, window)
35return {1, count, now}
36"#;
37
38pub fn parse_rate_limit_script_result(
39    value: Value,
40) -> Result<RateLimitScriptResult, OpenAuthError> {
41    let Value::Array(values) = value else {
42        return Err(OpenAuthError::Adapter(
43            "invalid fred rate limit script result: expected array".to_owned(),
44        ));
45    };
46    let [permitted, count, last_request]: [Value; 3] =
47        values.try_into().map_err(|values: Vec<Value>| {
48            OpenAuthError::Adapter(format!(
49                "invalid fred rate limit script result: expected 3 values, got {}",
50                values.len()
51            ))
52        })?;
53    let permitted = integer_value(permitted, "permitted")? == 1;
54    let count = integer_value(count, "count")?.max(0) as u64;
55    let last_request = integer_value(last_request, "last_request")?;
56    Ok(RateLimitScriptResult {
57        permitted,
58        count,
59        last_request,
60    })
61}
62
63fn integer_value(value: Value, field: &str) -> Result<i64, OpenAuthError> {
64    match value {
65        Value::Integer(value) => Ok(value),
66        _ => Err(OpenAuthError::Adapter(format!(
67            "invalid fred rate limit script result: `{field}` was not an integer"
68        ))),
69    }
70}