Skip to main content

tower_rate_limiter/limiter/
response.rs

1//! Middleware response construction and Rate Limit Fields.
2//!
3//! Owns the public response customization seam and the private finalization
4//! path shared by every response produced before the inner service is called.
5
6use std::time::Duration;
7
8use http::{HeaderValue, Request, Response, StatusCode, header::HeaderName};
9
10use super::{
11    error::RateLimitError,
12    policy::{Policy, ResponseMetadata},
13};
14
15/// The `RateLimit` header name.
16const RATE_LIMIT: HeaderName = HeaderName::from_static("ratelimit");
17/// The `RateLimit-Policy` header name.
18const RATE_LIMIT_POLICY: HeaderName = HeaderName::from_static("ratelimit-policy");
19/// The `Retry-After` header name.
20const RETRY_AFTER: HeaderName = HeaderName::from_static("retry-after");
21
22/// The Rate Limit Fields revision emitted in responses.
23///
24/// Draft 7 represents `RateLimit` as a dictionary containing `limit`, `remaining`, and `reset`.
25/// Its `RateLimit-Policy` field contains the quota and window without a policy identifier.
26///
27/// Draft 11 represents both fields as lists of named items. This crate emits its fixed-window
28/// `q`/`w` and `r`/`t` parameters; optional quota-unit (`qu`) and partition-key (`pk`) parameters
29/// are not emitted.
30///
31/// See the field definitions for [draft 7] and [draft 11].
32///
33/// [draft 7]: https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers-07#name-ratelimit-header-field-def
34/// [draft 11]: https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers-11#name-ratelimit-policy-field
35#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
36#[non_exhaustive]
37pub enum RateLimitFields {
38    /// Emit fields compatible with draft 7.
39    Draft7,
40    /// Emit fields compatible with draft 11.
41    #[default]
42    Draft11,
43    /// Do not emit `RateLimit` or `RateLimit-Policy` fields.
44    /// Rate-limited responses still include `Retry-After`.
45    Disabled,
46}
47
48/// The structured reason passed to a [`ResponseFactory`].
49#[derive(Debug)]
50pub enum ResponseReason {
51    /// The complete policy state after the request exceeded its resolved quota.
52    RateLimited(Policy),
53    /// The middleware could not resolve or charge the request's policy.
54    Error(RateLimitError),
55}
56
57impl ResponseReason {
58    /// Return the default HTTP status for this reason.
59    pub const fn status_code(&self) -> StatusCode {
60        match self {
61            Self::RateLimited(_) => StatusCode::TOO_MANY_REQUESTS,
62            Self::Error(RateLimitError::Key(_, _)) | Self::Error(RateLimitError::Quota(_, _)) => {
63                StatusCode::INTERNAL_SERVER_ERROR
64            },
65            Self::Error(RateLimitError::Store(_, _)) => StatusCode::SERVICE_UNAVAILABLE,
66        }
67    }
68}
69
70/// A factory for responses to rate-limit middleware results.
71pub trait ResponseFactory<ReqBody, ResBody>: Clone {
72    /// Build a response using the original request and structured reason.
73    fn build(&self, request: Request<ReqBody>, reason: ResponseReason) -> Response<ResBody>;
74}
75
76/// The default empty-body response factory.
77#[derive(Clone, Copy, Debug, Default)]
78#[non_exhaustive]
79pub struct DefaultResponseFactory;
80
81impl<ReqBody, ResBody> ResponseFactory<ReqBody, ResBody> for DefaultResponseFactory
82where
83    ResBody: Default,
84{
85    fn build(&self, _request: Request<ReqBody>, reason: ResponseReason) -> Response<ResBody> {
86        let mut response = Response::new(ResBody::default());
87        *response.status_mut() = reason.status_code();
88        response
89    }
90}
91
92/// A response produced by the middleware before the inner service is called.
93///
94/// Response construction and rate-limit field decoration stay deferred until
95/// the future's common ready state.
96pub(super) enum MiddlewareResponse<ReqBody> {
97    RateLimited(Request<ReqBody>, ResponseMetadata),
98    Error(Request<ReqBody>, RateLimitError),
99}
100
101impl<ReqBody> MiddlewareResponse<ReqBody> {
102    /// Build and decorate the final HTTP response through one middleware path.
103    pub(super) fn finalize<ResBody, Factory>(self, factory: &Factory) -> Response<ResBody>
104    where
105        Factory: ResponseFactory<ReqBody, ResBody>,
106    {
107        match self {
108            Self::RateLimited(request, metadata) => {
109                let reason = ResponseReason::RateLimited(metadata.policy.clone());
110                let response = factory.build(request, reason);
111
112                append_rate_limited_response_headers(response, metadata)
113            },
114            Self::Error(request, error) => factory.build(request, ResponseReason::Error(error)),
115        }
116    }
117}
118
119/// Decorate an allowed (or fail-open) inner response with Rate Limit Fields.
120///
121/// `None` means no quota metadata is available after bypass or fail-open, so no fields are
122/// written. When present, fields follow [`RateLimitFields`]; `Retry-After` is never added on this
123/// path.
124pub(super) fn append_inner_response_headers<B>(
125    response: Response<B>,
126    metadata: Option<ResponseMetadata>,
127) -> Response<B> {
128    match metadata {
129        Some(metadata) => append_rate_limit_fields(response, &metadata),
130        None => response,
131    }
132}
133
134/// Decorate a rate-limited middleware response with Rate Limit Fields and Retry-After.
135///
136/// Rate Limit Fields still respect [`RateLimitFields`]. `Retry-After` is always added.
137fn append_rate_limited_response_headers<B>(response: Response<B>, metadata: ResponseMetadata) -> Response<B> {
138    let mut response = append_rate_limit_fields(response, &metadata);
139    append_header(
140        &mut response,
141        RETRY_AFTER,
142        &ceil_seconds(metadata.policy.reset_after).to_string(),
143    );
144    response
145}
146
147/// Shared Rate Limit / RateLimit-Policy field writer.
148fn append_rate_limit_fields<B>(response: Response<B>, metadata: &ResponseMetadata) -> Response<B> {
149    let Some((policy, rate_limit)) = format_rate_limit_fields(metadata) else {
150        return response;
151    };
152
153    let mut response = response;
154    append_header(&mut response, RATE_LIMIT_POLICY, &policy);
155    append_header(&mut response, RATE_LIMIT, &rate_limit);
156    response
157}
158
159/// Format Rate Limit Fields for the configured revision.
160fn format_rate_limit_fields(metadata: &ResponseMetadata) -> Option<(String, String)> {
161    if metadata.fields == RateLimitFields::Disabled {
162        return None;
163    }
164
165    let limit = metadata.policy.limit;
166    let remaining = metadata.policy.remaining();
167    let reset_after = ceil_seconds(metadata.policy.reset_after);
168    let window = ceil_seconds(metadata.policy.window);
169
170    Some(match metadata.fields {
171        RateLimitFields::Draft7 => (
172            format!("{limit};w={window}"),
173            format!("limit={limit}, remaining={remaining}, reset={reset_after}"),
174        ),
175        RateLimitFields::Draft11 => {
176            let policy_name = &metadata.policy.name;
177
178            (
179                format!(r#""{policy_name}";q={limit};w={window}"#),
180                format!(r#""{policy_name}";r={remaining};t={reset_after}"#),
181            )
182        },
183        RateLimitFields::Disabled => return None,
184    })
185}
186
187/// Append a header to the response.
188fn append_header<B>(response: &mut Response<B>, name: HeaderName, value: &str) {
189    if let Ok(value) = HeaderValue::from_str(value) {
190        response.headers_mut().append(name, value);
191    }
192}
193
194/// Ceil the duration to the nearest second.
195fn ceil_seconds(duration: Duration) -> u64 {
196    duration
197        .as_secs()
198        .saturating_add(u64::from(duration.subsec_nanos() != 0))
199        .max(1)
200}