use std::time::Duration;
use http::Request;
use super::{RateLimitConfig, error::RateLimitError, response::RateLimitFields, store::Usage};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RateLimitPolicy {
pub policy_name: String,
pub limit: u64,
pub used: u64,
pub remaining: u64,
pub reset_after: Duration,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct RateLimitContext {
policies: Vec<RateLimitPolicy>,
}
impl RateLimitContext {
pub const fn new() -> Self {
Self { policies: Vec::new() }
}
pub fn policies(&self) -> &[RateLimitPolicy] {
&self.policies
}
}
#[derive(Clone, Debug)]
pub(super) struct ChargeMetadata {
pub(super) policy_name: String,
pub(super) limit: u64,
pub(super) usage: Usage,
pub(super) window: Duration,
pub(super) rate_limit_fields: RateLimitFields,
}
impl ChargeMetadata {
fn to_policy(&self) -> RateLimitPolicy {
RateLimitPolicy {
policy_name: self.policy_name.clone(),
limit: self.limit,
used: self.usage.used,
remaining: self.remaining(),
reset_after: self.usage.reset_after,
}
}
pub(super) fn remaining(&self) -> u64 {
self.limit.saturating_sub(self.usage.used)
}
}
pub(super) enum ChargeOutcome {
Allowed(ChargeMetadata),
RateLimited(ChargeMetadata),
}
impl ChargeOutcome {
pub(super) fn evaluate(usage: Usage, limit: u64, config: &RateLimitConfig) -> Result<Self, RateLimitError> {
if usage.used == 0 {
return Err(RateLimitError::Store(
String::from("invalid_usage"),
String::from("rate-limit store returned zero usage"),
));
}
let metadata = ChargeMetadata {
policy_name: config.policy_name.clone(),
limit,
usage,
window: config.window,
rate_limit_fields: config.rate_limit_fields,
};
Ok(if usage.used > limit {
Self::RateLimited(metadata)
} else {
Self::Allowed(metadata)
})
}
}
pub(super) fn make_key(policy_name: &str, client_key: &str) -> String {
format!("{}:{}", escape_key_part(policy_name), escape_key_part(client_key))
}
fn escape_key_part(value: &str) -> String {
value.replace('%', "%25").replace(':', "%3A")
}
pub(super) fn append_context<B>(request: &mut Request<B>, metadata: &ChargeMetadata) {
request
.extensions_mut()
.get_or_insert_default::<RateLimitContext>()
.policies
.push(metadata.to_policy());
}