tower-rate-tier 0.2.0

Tier-based rate limiting middleware for Tower
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

use http::header::HeaderValue;
use http::{Response, StatusCode};

use crate::gcra::{RateLimitInfo, RateLimited};

/// Inject `X-RateLimit-*` headers into a successful response.
///
/// `unix_offset_nanos` is added to `reset_at` to produce a Unix timestamp
/// for the `X-RateLimit-Reset` header.
pub fn inject_headers<B>(response: &mut Response<B>, info: &RateLimitInfo, unix_offset_nanos: u64) {
    let headers = response.headers_mut();
    headers.insert("X-RateLimit-Limit", HeaderValue::from(info.limit));
    headers.insert("X-RateLimit-Remaining", HeaderValue::from(info.remaining));
    headers.insert(
        "X-RateLimit-Reset",
        reset_header_value(info.reset_at, unix_offset_nanos),
    );
}

/// Build a 429 Too Many Requests response with JSON body and rate limit headers.
///
/// `unix_offset_nanos` is added to `reset_at` to produce a Unix timestamp
/// for the `X-RateLimit-Reset` header.
pub fn rate_limited_response(
    limited: &RateLimited,
    tier: &str,
    unix_offset_nanos: u64,
) -> Response<String> {
    let retry_after_secs = limited.retry_after.as_secs();

    let escaped_tier = escape_json_string(tier);
    let body = format!(
        r#"{{"error":"rate limit exceeded","tier":"{}","retry_after":{}}}"#,
        escaped_tier, retry_after_secs
    );

    let mut response = Response::builder()
        .status(StatusCode::TOO_MANY_REQUESTS)
        .header("Content-Type", "application/json")
        .header("Retry-After", retry_after_secs)
        .header("X-RateLimit-Limit", limited.limit)
        .header("X-RateLimit-Remaining", 0u32)
        .header(
            "X-RateLimit-Reset",
            reset_header_value(limited.reset_at, unix_offset_nanos),
        )
        .body(body)
        .unwrap();

    // Ensure Retry-After is at least 1 second when there's a non-zero duration
    if retry_after_secs == 0 && !limited.retry_after.is_zero() {
        response
            .headers_mut()
            .insert("Retry-After", HeaderValue::from(1u64));
    }

    response
}

/// Build a response for when the identifier cannot determine the user/tier
/// and the policy is `OnMissing::Deny(status)`.
pub fn deny_response(status: StatusCode) -> Response<String> {
    Response::builder()
        .status(status)
        .header("Content-Type", "application/json")
        .body(format!(r#"{{"error":"{}"}}"#, canonical_reason(status)))
        .unwrap()
}

/// Build a 503 Service Unavailable response for storage errors.
pub fn storage_error_response() -> Response<String> {
    Response::builder()
        .status(StatusCode::SERVICE_UNAVAILABLE)
        .header("Content-Type", "application/json")
        .body(r#"{"error":"service unavailable"}"#.to_string())
        .unwrap()
}

/// Convert an internal `reset_at` value to a Unix-timestamp `HeaderValue`.
///
/// Adds `unix_offset_nanos` to convert from process-local epoch to Unix epoch,
/// then divides by 1e9 to get seconds.
fn reset_header_value(reset_at_nanos: u64, unix_offset_nanos: u64) -> HeaderValue {
    let unix_nanos = reset_at_nanos.saturating_add(unix_offset_nanos);
    let secs = unix_nanos / 1_000_000_000;
    HeaderValue::from(secs)
}

/// Build a 400 Bad Request response for body read errors.
pub fn bad_request_response() -> Response<String> {
    Response::builder()
        .status(StatusCode::BAD_REQUEST)
        .header("Content-Type", "application/json")
        .body(r#"{"error":"failed to read request body"}"#.to_string())
        .unwrap()
}

fn canonical_reason(status: StatusCode) -> &'static str {
    status.canonical_reason().unwrap_or("request denied")
}

/// Escape a string for safe embedding in a JSON string value.
///
/// Handles `"`, `\`, and control characters per RFC 8259.
fn escape_json_string(s: &str) -> String {
    let mut escaped = String::with_capacity(s.len());
    for ch in s.chars() {
        match ch {
            '"' => escaped.push_str("\\\""),
            '\\' => escaped.push_str("\\\\"),
            c if c.is_control() => {
                escaped.push_str(&format!("\\u{:04x}", c as u32));
            }
            c => escaped.push(c),
        }
    }
    escaped
}