nora-registry 0.2.35

Cloud-Native Artifact Registry - Fast, lightweight, multi-protocol
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
119
120

// Copyright (c) 2026 Volkov Pavel | DevITWay
// SPDX-License-Identifier: MIT

//! Rate limiting configuration and middleware
//!
//! Provides rate limiting to protect against:
//! - Brute-force authentication attacks
//! - DoS attacks on upload endpoints
//! - General API abuse

use crate::config::RateLimitConfig;
use tower_governor::governor::GovernorConfigBuilder;
use tower_governor::key_extractor::SmartIpKeyExtractor;

/// Create rate limiter layer for auth endpoints (strict protection against brute-force)
pub fn auth_rate_limiter(
    config: &RateLimitConfig,
) -> tower_governor::GovernorLayer<
    tower_governor::key_extractor::PeerIpKeyExtractor,
    governor::middleware::StateInformationMiddleware,
    axum::body::Body,
> {
    let gov_config = GovernorConfigBuilder::default()
        .per_second(config.auth_rps)
        .burst_size(config.auth_burst)
        .use_headers()
        .finish()
        .expect("Failed to build auth rate limiter");

    tower_governor::GovernorLayer::new(gov_config)
}

/// Create rate limiter layer for upload endpoints
///
/// High limits to accommodate Docker client's aggressive parallel layer uploads
pub fn upload_rate_limiter(
    config: &RateLimitConfig,
) -> tower_governor::GovernorLayer<
    SmartIpKeyExtractor,
    governor::middleware::StateInformationMiddleware,
    axum::body::Body,
> {
    let gov_config = GovernorConfigBuilder::default()
        .key_extractor(SmartIpKeyExtractor)
        .per_second(config.upload_rps)
        .burst_size(config.upload_burst)
        .use_headers()
        .finish()
        .expect("Failed to build upload rate limiter");

    tower_governor::GovernorLayer::new(gov_config)
}

/// Create rate limiter layer for general endpoints (lenient)
pub fn general_rate_limiter(
    config: &RateLimitConfig,
) -> tower_governor::GovernorLayer<
    SmartIpKeyExtractor,
    governor::middleware::StateInformationMiddleware,
    axum::body::Body,
> {
    let gov_config = GovernorConfigBuilder::default()
        .key_extractor(SmartIpKeyExtractor)
        .per_second(config.general_rps)
        .burst_size(config.general_burst)
        .use_headers()
        .finish()
        .expect("Failed to build general rate limiter");

    tower_governor::GovernorLayer::new(gov_config)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::RateLimitConfig;

    #[test]
    fn test_default_config() {
        let config = RateLimitConfig::default();
        assert_eq!(config.auth_rps, 1);
        assert_eq!(config.auth_burst, 5);
        assert_eq!(config.upload_rps, 200);
        assert_eq!(config.general_rps, 100);
    }

    #[test]
    fn test_auth_rate_limiter_creation() {
        let config = RateLimitConfig::default();
        let _limiter = auth_rate_limiter(&config);
    }

    #[test]
    fn test_upload_rate_limiter_creation() {
        let config = RateLimitConfig::default();
        let _limiter = upload_rate_limiter(&config);
    }

    #[test]
    fn test_general_rate_limiter_creation() {
        let config = RateLimitConfig::default();
        let _limiter = general_rate_limiter(&config);
    }

    #[test]
    fn test_custom_config() {
        let config = RateLimitConfig {
            enabled: true,
            auth_rps: 10,
            auth_burst: 20,
            upload_rps: 500,
            upload_burst: 1000,
            general_rps: 200,
            general_burst: 400,
        };
        let _auth = auth_rate_limiter(&config);
        let _upload = upload_rate_limiter(&config);
        let _general = general_rate_limiter(&config);
    }
}