fix(security): add rate limiting to registry API
- commit
- d713603278fdcbbcac072af17ac86f9ab10b7e6b
- author
- bnema · · 2026-01-17 20:58 UTC
- parent
- 1a4fae4
- signature
- unsigned
2 files · +350 -0
A internal/adapters/in/http/registry/ratelimit.go +153 -0
| @@ -0,0 +1,153 @@ | |||
| 1 | + | // Package registry implements the HTTP adapter for the registry API. | |
| 2 | + | package registry | |
| 3 | + | ||
| 4 | + | import ( | |
| 5 | + | "encoding/json" | |
| 6 | + | "net" | |
| 7 | + | "net/http" | |
| 8 | + | "strings" | |
| 9 | + | "sync" | |
| 10 | + | ||
| 11 | + | "golang.org/x/time/rate" | |
| 12 | + | ||
| 13 | + | "gordon/internal/adapters/dto" | |
| 14 | + | ) | |
| 15 | + | ||
| 16 | + | // RateLimitConfig holds rate limiting configuration for the registry API. | |
| 17 | + | // These settings help prevent DoS attacks and brute force attempts. | |
| 18 | + | type RateLimitConfig struct { | |
| 19 | + | // Enabled controls whether rate limiting is active. | |
| 20 | + | // Default: true | |
| 21 | + | Enabled bool | |
| 22 | + | ||
| 23 | + | // GlobalRPS is the maximum requests per second across all clients. | |
| 24 | + | // Default: 500 | |
| 25 | + | GlobalRPS float64 | |
| 26 | + | ||
| 27 | + | // PerIPRPS is the maximum requests per second per client IP. | |
| 28 | + | // Default: 50 | |
| 29 | + | PerIPRPS float64 | |
| 30 | + | ||
| 31 | + | // Burst is the maximum burst size for rate limiters. | |
| 32 | + | // Default: 100 | |
| 33 | + | Burst int | |
| 34 | + | } | |
| 35 | + | ||
| 36 | + | // DefaultRateLimitConfig returns sensible default rate limiting configuration. | |
| 37 | + | func DefaultRateLimitConfig() RateLimitConfig { | |
| 38 | + | return RateLimitConfig{ | |
| 39 | + | Enabled: true, | |
| 40 | + | GlobalRPS: 500, // 500 requests/second globally | |
| 41 | + | PerIPRPS: 50, // 50 requests/second per IP | |
| 42 | + | Burst: 100, // Allow burst of 100 requests | |
| 43 | + | } | |
| 44 | + | } | |
| 45 | + | ||
| 46 | + | // ipRateLimiter manages per-IP rate limiters. | |
| 47 | + | type ipRateLimiter struct { | |
| 48 | + | limiters map[string]*rate.Limiter | |
| 49 | + | mu sync.RWMutex | |
| 50 | + | rps float64 | |
| 51 | + | burst int | |
| 52 | + | } | |
| 53 | + | ||
| 54 | + | func newIPRateLimiter(rps float64, burst int) *ipRateLimiter { | |
| 55 | + | return &ipRateLimiter{ | |
| 56 | + | limiters: make(map[string]*rate.Limiter), | |
| 57 | + | rps: rps, | |
| 58 | + | burst: burst, | |
| 59 | + | } | |
| 60 | + | } | |
| 61 | + | ||
| 62 | + | func (l *ipRateLimiter) getLimiter(ip string) *rate.Limiter { | |
| 63 | + | l.mu.RLock() | |
| 64 | + | limiter, exists := l.limiters[ip] | |
| 65 | + | l.mu.RUnlock() | |
| 66 | + | ||
| 67 | + | if exists { | |
| 68 | + | return limiter | |
| 69 | + | } | |
| 70 | + | ||
| 71 | + | l.mu.Lock() | |
| 72 | + | defer l.mu.Unlock() | |
| 73 | + | ||
| 74 | + | // Double-check after acquiring write lock | |
| 75 | + | if limiter, exists = l.limiters[ip]; exists { | |
| 76 | + | return limiter | |
| 77 | + | } | |
| 78 | + | ||
| 79 | + | limiter = rate.NewLimiter(rate.Limit(l.rps), l.burst) | |
| 80 | + | l.limiters[ip] = limiter | |
| 81 | + | return limiter | |
| 82 | + | } | |
| 83 | + | ||
| 84 | + | // RateLimitMiddleware creates rate limiting middleware for the registry API. | |
| 85 | + | // Returns nil if rate limiting is disabled. | |
| 86 | + | func RateLimitMiddleware(config RateLimitConfig) func(http.Handler) http.Handler { | |
| 87 | + | if !config.Enabled { | |
| 88 | + | return func(next http.Handler) http.Handler { | |
| 89 | + | return next // Pass through if disabled | |
| 90 | + | } | |
| 91 | + | } | |
| 92 | + | ||
| 93 | + | globalLimiter := rate.NewLimiter(rate.Limit(config.GlobalRPS), config.Burst) | |
| 94 | + | ipLimiter := newIPRateLimiter(config.PerIPRPS, config.Burst) | |
| 95 | + | ||
| 96 | + | return func(next http.Handler) http.Handler { | |
| 97 | + | return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 98 | + | // Check global rate limit | |
| 99 | + | if !globalLimiter.Allow() { | |
| 100 | + | sendRateLimitError(w) | |
| 101 | + | return | |
| 102 | + | } | |
| 103 | + | ||
| 104 | + | // Check per-IP rate limit | |
| 105 | + | ip := getClientIP(r) | |
| 106 | + | if !ipLimiter.getLimiter(ip).Allow() { | |
| 107 | + | sendRateLimitError(w) | |
| 108 | + | return | |
| 109 | + | } | |
| 110 | + | ||
| 111 | + | next.ServeHTTP(w, r) | |
| 112 | + | }) | |
| 113 | + | } | |
| 114 | + | } | |
| 115 | + | ||
| 116 | + | // sendRateLimitError sends an HTTP 429 response in Docker Registry API format. | |
| 117 | + | func sendRateLimitError(w http.ResponseWriter) { | |
| 118 | + | w.Header().Set("Content-Type", "application/json") | |
| 119 | + | w.Header().Set("Docker-Distribution-API-Version", "registry/2.0") | |
| 120 | + | w.Header().Set("Retry-After", "1") | |
| 121 | + | w.WriteHeader(http.StatusTooManyRequests) | |
| 122 | + | _ = json.NewEncoder(w).Encode(dto.RegistryErrorResponse{ | |
| 123 | + | Errors: []dto.RegistryErrorItem{{ | |
| 124 | + | Code: "TOOMANYREQUESTS", | |
| 125 | + | Message: "rate limit exceeded", | |
| 126 | + | }}, | |
| 127 | + | }) | |
| 128 | + | } | |
| 129 | + | ||
| 130 | + | // getClientIP extracts the client IP address from the request. | |
| 131 | + | // It checks X-Forwarded-For and X-Real-IP headers for proxied requests. | |
| 132 | + | func getClientIP(r *http.Request) string { | |
| 133 | + | // Check X-Forwarded-For for proxied requests | |
| 134 | + | if xff := r.Header.Get("X-Forwarded-For"); xff != "" { | |
| 135 | + | // Take first IP in chain (original client) | |
| 136 | + | if idx := strings.Index(xff, ","); idx != -1 { | |
| 137 | + | return strings.TrimSpace(xff[:idx]) | |
| 138 | + | } | |
| 139 | + | return strings.TrimSpace(xff) | |
| 140 | + | } | |
| 141 | + | ||
| 142 | + | // Check X-Real-IP | |
| 143 | + | if xri := r.Header.Get("X-Real-IP"); xri != "" { | |
| 144 | + | return xri | |
| 145 | + | } | |
| 146 | + | ||
| 147 | + | // Fall back to RemoteAddr | |
| 148 | + | ip, _, err := net.SplitHostPort(r.RemoteAddr) | |
| 149 | + | if err != nil { | |
| 150 | + | return r.RemoteAddr | |
| 151 | + | } | |
| 152 | + | return ip | |
| 153 | + | } | |
| @@ -0,0 +1,153 @@ | |||||
| 1 | + | // Package registry implements the HTTP adapter for the registry API. | |||
| 2 | + | package registry | |||
| 3 | + | ||||
| 4 | + | import ( | |||
| 5 | + | "encoding/json" | |||
| 6 | + | "net" | |||
| 7 | + | "net/http" | |||
| 8 | + | "strings" | |||
| 9 | + | "sync" | |||
| 10 | + | ||||
| 11 | + | "golang.org/x/time/rate" | |||
| 12 | + | ||||
| 13 | + | "gordon/internal/adapters/dto" | |||
| 14 | + | ) | |||
| 15 | + | ||||
| 16 | + | // RateLimitConfig holds rate limiting configuration for the registry API. | |||
| 17 | + | // These settings help prevent DoS attacks and brute force attempts. | |||
| 18 | + | type RateLimitConfig struct { | |||
| 19 | + | // Enabled controls whether rate limiting is active. | |||
| 20 | + | // Default: true | |||
| 21 | + | Enabled bool | |||
| 22 | + | ||||
| 23 | + | // GlobalRPS is the maximum requests per second across all clients. | |||
| 24 | + | // Default: 500 | |||
| 25 | + | GlobalRPS float64 | |||
| 26 | + | ||||
| 27 | + | // PerIPRPS is the maximum requests per second per client IP. | |||
| 28 | + | // Default: 50 | |||
| 29 | + | PerIPRPS float64 | |||
| 30 | + | ||||
| 31 | + | // Burst is the maximum burst size for rate limiters. | |||
| 32 | + | // Default: 100 | |||
| 33 | + | Burst int | |||
| 34 | + | } | |||
| 35 | + | ||||
| 36 | + | // DefaultRateLimitConfig returns sensible default rate limiting configuration. | |||
| 37 | + | func DefaultRateLimitConfig() RateLimitConfig { | |||
| 38 | + | return RateLimitConfig{ | |||
| 39 | + | Enabled: true, | |||
| 40 | + | GlobalRPS: 500, // 500 requests/second globally | |||
| 41 | + | PerIPRPS: 50, // 50 requests/second per IP | |||
| 42 | + | Burst: 100, // Allow burst of 100 requests | |||
| 43 | + | } | |||
| 44 | + | } | |||
| 45 | + | ||||
| 46 | + | // ipRateLimiter manages per-IP rate limiters. | |||
| 47 | + | type ipRateLimiter struct { | |||
| 48 | + | limiters map[string]*rate.Limiter | |||
| 49 | + | mu sync.RWMutex | |||
| 50 | + | rps float64 | |||
| 51 | + | burst int | |||
| 52 | + | } | |||
| 53 | + | ||||
| 54 | + | func newIPRateLimiter(rps float64, burst int) *ipRateLimiter { | |||
| 55 | + | return &ipRateLimiter{ | |||
| 56 | + | limiters: make(map[string]*rate.Limiter), | |||
| 57 | + | rps: rps, | |||
| 58 | + | burst: burst, | |||
| 59 | + | } | |||
| 60 | + | } | |||
| 61 | + | ||||
| 62 | + | func (l *ipRateLimiter) getLimiter(ip string) *rate.Limiter { | |||
| 63 | + | l.mu.RLock() | |||
| 64 | + | limiter, exists := l.limiters[ip] | |||
| 65 | + | l.mu.RUnlock() | |||
| 66 | + | ||||
| 67 | + | if exists { | |||
| 68 | + | return limiter | |||
| 69 | + | } | |||
| 70 | + | ||||
| 71 | + | l.mu.Lock() | |||
| 72 | + | defer l.mu.Unlock() | |||
| 73 | + | ||||
| 74 | + | // Double-check after acquiring write lock | |||
| 75 | + | if limiter, exists = l.limiters[ip]; exists { | |||
| 76 | + | return limiter | |||
| 77 | + | } | |||
| 78 | + | ||||
| 79 | + | limiter = rate.NewLimiter(rate.Limit(l.rps), l.burst) | |||
| 80 | + | l.limiters[ip] = limiter | |||
| 81 | + | return limiter | |||
| 82 | + | } | |||
| 83 | + | ||||
| 84 | + | // RateLimitMiddleware creates rate limiting middleware for the registry API. | |||
| 85 | + | // Returns nil if rate limiting is disabled. | |||
| 86 | + | func RateLimitMiddleware(config RateLimitConfig) func(http.Handler) http.Handler { | |||
| 87 | + | if !config.Enabled { | |||
| 88 | + | return func(next http.Handler) http.Handler { | |||
| 89 | + | return next // Pass through if disabled | |||
| 90 | + | } | |||
| 91 | + | } | |||
| 92 | + | ||||
| 93 | + | globalLimiter := rate.NewLimiter(rate.Limit(config.GlobalRPS), config.Burst) | |||
| 94 | + | ipLimiter := newIPRateLimiter(config.PerIPRPS, config.Burst) | |||
| 95 | + | ||||
| 96 | + | return func(next http.Handler) http.Handler { | |||
| 97 | + | return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| 98 | + | // Check global rate limit | |||
| 99 | + | if !globalLimiter.Allow() { | |||
| 100 | + | sendRateLimitError(w) | |||
| 101 | + | return | |||
| 102 | + | } | |||
| 103 | + | ||||
| 104 | + | // Check per-IP rate limit | |||
| 105 | + | ip := getClientIP(r) | |||
| 106 | + | if !ipLimiter.getLimiter(ip).Allow() { | |||
| 107 | + | sendRateLimitError(w) | |||
| 108 | + | return | |||
| 109 | + | } | |||
| 110 | + | ||||
| 111 | + | next.ServeHTTP(w, r) | |||
| 112 | + | }) | |||
| 113 | + | } | |||
| 114 | + | } | |||
| 115 | + | ||||
| 116 | + | // sendRateLimitError sends an HTTP 429 response in Docker Registry API format. | |||
| 117 | + | func sendRateLimitError(w http.ResponseWriter) { | |||
| 118 | + | w.Header().Set("Content-Type", "application/json") | |||
| 119 | + | w.Header().Set("Docker-Distribution-API-Version", "registry/2.0") | |||
| 120 | + | w.Header().Set("Retry-After", "1") | |||
| 121 | + | w.WriteHeader(http.StatusTooManyRequests) | |||
| 122 | + | _ = json.NewEncoder(w).Encode(dto.RegistryErrorResponse{ | |||
| 123 | + | Errors: []dto.RegistryErrorItem{{ | |||
| 124 | + | Code: "TOOMANYREQUESTS", | |||
| 125 | + | Message: "rate limit exceeded", | |||
| 126 | + | }}, | |||
| 127 | + | }) | |||
| 128 | + | } | |||
| 129 | + | ||||
| 130 | + | // getClientIP extracts the client IP address from the request. | |||
| 131 | + | // It checks X-Forwarded-For and X-Real-IP headers for proxied requests. | |||
| 132 | + | func getClientIP(r *http.Request) string { | |||
| 133 | + | // Check X-Forwarded-For for proxied requests | |||
| 134 | + | if xff := r.Header.Get("X-Forwarded-For"); xff != "" { | |||
| 135 | + | // Take first IP in chain (original client) | |||
| 136 | + | if idx := strings.Index(xff, ","); idx != -1 { | |||
| 137 | + | return strings.TrimSpace(xff[:idx]) | |||
| 138 | + | } | |||
| 139 | + | return strings.TrimSpace(xff) | |||
| 140 | + | } | |||
| 141 | + | ||||
| 142 | + | // Check X-Real-IP | |||
| 143 | + | if xri := r.Header.Get("X-Real-IP"); xri != "" { | |||
| 144 | + | return xri | |||
| 145 | + | } | |||
| 146 | + | ||||
| 147 | + | // Fall back to RemoteAddr | |||
| 148 | + | ip, _, err := net.SplitHostPort(r.RemoteAddr) | |||
| 149 | + | if err != nil { | |||
| 150 | + | return r.RemoteAddr | |||
| 151 | + | } | |||
| 152 | + | return ip | |||
| 153 | + | } | |||
A internal/adapters/in/http/registry/ratelimit_test.go +197 -0
| @@ -0,0 +1,197 @@ | |||
| 1 | + | package registry | |
| 2 | + | ||
| 3 | + | import ( | |
| 4 | + | "net/http" | |
| 5 | + | "net/http/httptest" | |
| 6 | + | "testing" | |
| 7 | + | ||
| 8 | + | "github.com/stretchr/testify/assert" | |
| 9 | + | "github.com/stretchr/testify/require" | |
| 10 | + | ) | |
| 11 | + | ||
| 12 | + | func TestRateLimitMiddleware_Disabled(t *testing.T) { | |
| 13 | + | config := RateLimitConfig{ | |
| 14 | + | Enabled: false, | |
| 15 | + | } | |
| 16 | + | ||
| 17 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 18 | + | w.WriteHeader(http.StatusOK) | |
| 19 | + | }) | |
| 20 | + | ||
| 21 | + | middleware := RateLimitMiddleware(config) | |
| 22 | + | wrappedHandler := middleware(handler) | |
| 23 | + | ||
| 24 | + | // Should pass through without rate limiting | |
| 25 | + | for i := 0; i < 1000; i++ { | |
| 26 | + | req := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |
| 27 | + | rec := httptest.NewRecorder() | |
| 28 | + | wrappedHandler.ServeHTTP(rec, req) | |
| 29 | + | assert.Equal(t, http.StatusOK, rec.Code) | |
| 30 | + | } | |
| 31 | + | } | |
| 32 | + | ||
| 33 | + | func TestRateLimitMiddleware_PerIPLimit(t *testing.T) { | |
| 34 | + | config := RateLimitConfig{ | |
| 35 | + | Enabled: true, | |
| 36 | + | GlobalRPS: 1000, // High global limit | |
| 37 | + | PerIPRPS: 1, // Very low per-IP limit for testing | |
| 38 | + | Burst: 100, // High burst to not interfere | |
| 39 | + | } | |
| 40 | + | ||
| 41 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 42 | + | w.WriteHeader(http.StatusOK) | |
| 43 | + | }) | |
| 44 | + | ||
| 45 | + | middleware := RateLimitMiddleware(config) | |
| 46 | + | wrappedHandler := middleware(handler) | |
| 47 | + | ||
| 48 | + | // First request should succeed | |
| 49 | + | req1 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |
| 50 | + | req1.RemoteAddr = "192.168.1.100:12345" | |
| 51 | + | rec1 := httptest.NewRecorder() | |
| 52 | + | wrappedHandler.ServeHTTP(rec1, req1) | |
| 53 | + | assert.Equal(t, http.StatusOK, rec1.Code) | |
| 54 | + | ||
| 55 | + | // Request from different IP should succeed (independent rate limiter) | |
| 56 | + | req2 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |
| 57 | + | req2.RemoteAddr = "192.168.1.101:12345" | |
| 58 | + | rec2 := httptest.NewRecorder() | |
| 59 | + | wrappedHandler.ServeHTTP(rec2, req2) | |
| 60 | + | assert.Equal(t, http.StatusOK, rec2.Code) | |
| 61 | + | } | |
| 62 | + | ||
| 63 | + | func TestRateLimitMiddleware_GlobalLimit(t *testing.T) { | |
| 64 | + | config := RateLimitConfig{ | |
| 65 | + | Enabled: true, | |
| 66 | + | GlobalRPS: 1, // Very low global limit for testing | |
| 67 | + | PerIPRPS: 100, | |
| 68 | + | Burst: 1, | |
| 69 | + | } | |
| 70 | + | ||
| 71 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 72 | + | w.WriteHeader(http.StatusOK) | |
| 73 | + | }) | |
| 74 | + | ||
| 75 | + | middleware := RateLimitMiddleware(config) | |
| 76 | + | wrappedHandler := middleware(handler) | |
| 77 | + | ||
| 78 | + | // First request should succeed | |
| 79 | + | req1 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |
| 80 | + | req1.RemoteAddr = "192.168.1.100:12345" | |
| 81 | + | rec1 := httptest.NewRecorder() | |
| 82 | + | wrappedHandler.ServeHTTP(rec1, req1) | |
| 83 | + | assert.Equal(t, http.StatusOK, rec1.Code) | |
| 84 | + | ||
| 85 | + | // Second request should hit global limit even from different IP | |
| 86 | + | req2 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |
| 87 | + | req2.RemoteAddr = "192.168.1.101:12345" | |
| 88 | + | rec2 := httptest.NewRecorder() | |
| 89 | + | wrappedHandler.ServeHTTP(rec2, req2) | |
| 90 | + | assert.Equal(t, http.StatusTooManyRequests, rec2.Code) | |
| 91 | + | } | |
| 92 | + | ||
| 93 | + | func TestRateLimitMiddleware_ResponseFormat(t *testing.T) { | |
| 94 | + | config := RateLimitConfig{ | |
| 95 | + | Enabled: true, | |
| 96 | + | GlobalRPS: 1, | |
| 97 | + | PerIPRPS: 1, | |
| 98 | + | Burst: 1, | |
| 99 | + | } | |
| 100 | + | ||
| 101 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 102 | + | w.WriteHeader(http.StatusOK) | |
| 103 | + | }) | |
| 104 | + | ||
| 105 | + | middleware := RateLimitMiddleware(config) | |
| 106 | + | wrappedHandler := middleware(handler) | |
| 107 | + | ||
| 108 | + | // First request succeeds | |
| 109 | + | req1 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |
| 110 | + | req1.RemoteAddr = "192.168.1.100:12345" | |
| 111 | + | rec1 := httptest.NewRecorder() | |
| 112 | + | wrappedHandler.ServeHTTP(rec1, req1) | |
| 113 | + | ||
| 114 | + | // Second request gets rate limited | |
| 115 | + | req2 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |
| 116 | + | req2.RemoteAddr = "192.168.1.100:12346" | |
| 117 | + | rec2 := httptest.NewRecorder() | |
| 118 | + | wrappedHandler.ServeHTTP(rec2, req2) | |
| 119 | + | ||
| 120 | + | // Verify response format | |
| 121 | + | assert.Equal(t, http.StatusTooManyRequests, rec2.Code) | |
| 122 | + | assert.Equal(t, "application/json", rec2.Header().Get("Content-Type")) | |
| 123 | + | assert.Equal(t, "registry/2.0", rec2.Header().Get("Docker-Distribution-API-Version")) | |
| 124 | + | assert.Equal(t, "1", rec2.Header().Get("Retry-After")) | |
| 125 | + | assert.Contains(t, rec2.Body.String(), "TOOMANYREQUESTS") | |
| 126 | + | } | |
| 127 | + | ||
| 128 | + | func TestGetClientIP(t *testing.T) { | |
| 129 | + | tests := []struct { | |
| 130 | + | name string | |
| 131 | + | remoteAddr string | |
| 132 | + | xff string | |
| 133 | + | xRealIP string | |
| 134 | + | wantIP string | |
| 135 | + | }{ | |
| 136 | + | { | |
| 137 | + | name: "from RemoteAddr with port", | |
| 138 | + | remoteAddr: "192.168.1.100:12345", | |
| 139 | + | wantIP: "192.168.1.100", | |
| 140 | + | }, | |
| 141 | + | { | |
| 142 | + | name: "from RemoteAddr without port", | |
| 143 | + | remoteAddr: "192.168.1.100", | |
| 144 | + | wantIP: "192.168.1.100", | |
| 145 | + | }, | |
| 146 | + | { | |
| 147 | + | name: "from X-Forwarded-For single IP", | |
| 148 | + | remoteAddr: "127.0.0.1:12345", | |
| 149 | + | xff: "203.0.113.50", | |
| 150 | + | wantIP: "203.0.113.50", | |
| 151 | + | }, | |
| 152 | + | { | |
| 153 | + | name: "from X-Forwarded-For multiple IPs", | |
| 154 | + | remoteAddr: "127.0.0.1:12345", | |
| 155 | + | xff: "203.0.113.50, 10.0.0.1, 172.16.0.1", | |
| 156 | + | wantIP: "203.0.113.50", | |
| 157 | + | }, | |
| 158 | + | { | |
| 159 | + | name: "from X-Real-IP", | |
| 160 | + | remoteAddr: "127.0.0.1:12345", | |
| 161 | + | xRealIP: "203.0.113.50", | |
| 162 | + | wantIP: "203.0.113.50", | |
| 163 | + | }, | |
| 164 | + | { | |
| 165 | + | name: "X-Forwarded-For takes precedence", | |
| 166 | + | remoteAddr: "127.0.0.1:12345", | |
| 167 | + | xff: "203.0.113.50", | |
| 168 | + | xRealIP: "203.0.113.60", | |
| 169 | + | wantIP: "203.0.113.50", | |
| 170 | + | }, | |
| 171 | + | } | |
| 172 | + | ||
| 173 | + | for _, tt := range tests { | |
| 174 | + | t.Run(tt.name, func(t *testing.T) { | |
| 175 | + | req := httptest.NewRequest(http.MethodGet, "/", nil) | |
| 176 | + | req.RemoteAddr = tt.remoteAddr | |
| 177 | + | if tt.xff != "" { | |
| 178 | + | req.Header.Set("X-Forwarded-For", tt.xff) | |
| 179 | + | } | |
| 180 | + | if tt.xRealIP != "" { | |
| 181 | + | req.Header.Set("X-Real-IP", tt.xRealIP) | |
| 182 | + | } | |
| 183 | + | ||
| 184 | + | ip := getClientIP(req) | |
| 185 | + | assert.Equal(t, tt.wantIP, ip) | |
| 186 | + | }) | |
| 187 | + | } | |
| 188 | + | } | |
| 189 | + | ||
| 190 | + | func TestDefaultRateLimitConfig(t *testing.T) { | |
| 191 | + | config := DefaultRateLimitConfig() | |
| 192 | + | ||
| 193 | + | require.True(t, config.Enabled) | |
| 194 | + | assert.Equal(t, float64(500), config.GlobalRPS) | |
| 195 | + | assert.Equal(t, float64(50), config.PerIPRPS) | |
| 196 | + | assert.Equal(t, 100, config.Burst) | |
| 197 | + | } | |
| @@ -0,0 +1,197 @@ | |||||
| 1 | + | package registry | |||
| 2 | + | ||||
| 3 | + | import ( | |||
| 4 | + | "net/http" | |||
| 5 | + | "net/http/httptest" | |||
| 6 | + | "testing" | |||
| 7 | + | ||||
| 8 | + | "github.com/stretchr/testify/assert" | |||
| 9 | + | "github.com/stretchr/testify/require" | |||
| 10 | + | ) | |||
| 11 | + | ||||
| 12 | + | func TestRateLimitMiddleware_Disabled(t *testing.T) { | |||
| 13 | + | config := RateLimitConfig{ | |||
| 14 | + | Enabled: false, | |||
| 15 | + | } | |||
| 16 | + | ||||
| 17 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| 18 | + | w.WriteHeader(http.StatusOK) | |||
| 19 | + | }) | |||
| 20 | + | ||||
| 21 | + | middleware := RateLimitMiddleware(config) | |||
| 22 | + | wrappedHandler := middleware(handler) | |||
| 23 | + | ||||
| 24 | + | // Should pass through without rate limiting | |||
| 25 | + | for i := 0; i < 1000; i++ { | |||
| 26 | + | req := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |||
| 27 | + | rec := httptest.NewRecorder() | |||
| 28 | + | wrappedHandler.ServeHTTP(rec, req) | |||
| 29 | + | assert.Equal(t, http.StatusOK, rec.Code) | |||
| 30 | + | } | |||
| 31 | + | } | |||
| 32 | + | ||||
| 33 | + | func TestRateLimitMiddleware_PerIPLimit(t *testing.T) { | |||
| 34 | + | config := RateLimitConfig{ | |||
| 35 | + | Enabled: true, | |||
| 36 | + | GlobalRPS: 1000, // High global limit | |||
| 37 | + | PerIPRPS: 1, // Very low per-IP limit for testing | |||
| 38 | + | Burst: 100, // High burst to not interfere | |||
| 39 | + | } | |||
| 40 | + | ||||
| 41 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| 42 | + | w.WriteHeader(http.StatusOK) | |||
| 43 | + | }) | |||
| 44 | + | ||||
| 45 | + | middleware := RateLimitMiddleware(config) | |||
| 46 | + | wrappedHandler := middleware(handler) | |||
| 47 | + | ||||
| 48 | + | // First request should succeed | |||
| 49 | + | req1 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |||
| 50 | + | req1.RemoteAddr = "192.168.1.100:12345" | |||
| 51 | + | rec1 := httptest.NewRecorder() | |||
| 52 | + | wrappedHandler.ServeHTTP(rec1, req1) | |||
| 53 | + | assert.Equal(t, http.StatusOK, rec1.Code) | |||
| 54 | + | ||||
| 55 | + | // Request from different IP should succeed (independent rate limiter) | |||
| 56 | + | req2 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |||
| 57 | + | req2.RemoteAddr = "192.168.1.101:12345" | |||
| 58 | + | rec2 := httptest.NewRecorder() | |||
| 59 | + | wrappedHandler.ServeHTTP(rec2, req2) | |||
| 60 | + | assert.Equal(t, http.StatusOK, rec2.Code) | |||
| 61 | + | } | |||
| 62 | + | ||||
| 63 | + | func TestRateLimitMiddleware_GlobalLimit(t *testing.T) { | |||
| 64 | + | config := RateLimitConfig{ | |||
| 65 | + | Enabled: true, | |||
| 66 | + | GlobalRPS: 1, // Very low global limit for testing | |||
| 67 | + | PerIPRPS: 100, | |||
| 68 | + | Burst: 1, | |||
| 69 | + | } | |||
| 70 | + | ||||
| 71 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| 72 | + | w.WriteHeader(http.StatusOK) | |||
| 73 | + | }) | |||
| 74 | + | ||||
| 75 | + | middleware := RateLimitMiddleware(config) | |||
| 76 | + | wrappedHandler := middleware(handler) | |||
| 77 | + | ||||
| 78 | + | // First request should succeed | |||
| 79 | + | req1 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |||
| 80 | + | req1.RemoteAddr = "192.168.1.100:12345" | |||
| 81 | + | rec1 := httptest.NewRecorder() | |||
| 82 | + | wrappedHandler.ServeHTTP(rec1, req1) | |||
| 83 | + | assert.Equal(t, http.StatusOK, rec1.Code) | |||
| 84 | + | ||||
| 85 | + | // Second request should hit global limit even from different IP | |||
| 86 | + | req2 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |||
| 87 | + | req2.RemoteAddr = "192.168.1.101:12345" | |||
| 88 | + | rec2 := httptest.NewRecorder() | |||
| 89 | + | wrappedHandler.ServeHTTP(rec2, req2) | |||
| 90 | + | assert.Equal(t, http.StatusTooManyRequests, rec2.Code) | |||
| 91 | + | } | |||
| 92 | + | ||||
| 93 | + | func TestRateLimitMiddleware_ResponseFormat(t *testing.T) { | |||
| 94 | + | config := RateLimitConfig{ | |||
| 95 | + | Enabled: true, | |||
| 96 | + | GlobalRPS: 1, | |||
| 97 | + | PerIPRPS: 1, | |||
| 98 | + | Burst: 1, | |||
| 99 | + | } | |||
| 100 | + | ||||
| 101 | + | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| 102 | + | w.WriteHeader(http.StatusOK) | |||
| 103 | + | }) | |||
| 104 | + | ||||
| 105 | + | middleware := RateLimitMiddleware(config) | |||
| 106 | + | wrappedHandler := middleware(handler) | |||
| 107 | + | ||||
| 108 | + | // First request succeeds | |||
| 109 | + | req1 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |||
| 110 | + | req1.RemoteAddr = "192.168.1.100:12345" | |||
| 111 | + | rec1 := httptest.NewRecorder() | |||
| 112 | + | wrappedHandler.ServeHTTP(rec1, req1) | |||
| 113 | + | ||||
| 114 | + | // Second request gets rate limited | |||
| 115 | + | req2 := httptest.NewRequest(http.MethodGet, "/v2/", nil) | |||
| 116 | + | req2.RemoteAddr = "192.168.1.100:12346" | |||
| 117 | + | rec2 := httptest.NewRecorder() | |||
| 118 | + | wrappedHandler.ServeHTTP(rec2, req2) | |||
| 119 | + | ||||
| 120 | + | // Verify response format | |||
| 121 | + | assert.Equal(t, http.StatusTooManyRequests, rec2.Code) | |||
| 122 | + | assert.Equal(t, "application/json", rec2.Header().Get("Content-Type")) | |||
| 123 | + | assert.Equal(t, "registry/2.0", rec2.Header().Get("Docker-Distribution-API-Version")) | |||
| 124 | + | assert.Equal(t, "1", rec2.Header().Get("Retry-After")) | |||
| 125 | + | assert.Contains(t, rec2.Body.String(), "TOOMANYREQUESTS") | |||
| 126 | + | } | |||
| 127 | + | ||||
| 128 | + | func TestGetClientIP(t *testing.T) { | |||
| 129 | + | tests := []struct { | |||
| 130 | + | name string | |||
| 131 | + | remoteAddr string | |||
| 132 | + | xff string | |||
| 133 | + | xRealIP string | |||
| 134 | + | wantIP string | |||
| 135 | + | }{ | |||
| 136 | + | { | |||
| 137 | + | name: "from RemoteAddr with port", | |||
| 138 | + | remoteAddr: "192.168.1.100:12345", | |||
| 139 | + | wantIP: "192.168.1.100", | |||
| 140 | + | }, | |||
| 141 | + | { | |||
| 142 | + | name: "from RemoteAddr without port", | |||
| 143 | + | remoteAddr: "192.168.1.100", | |||
| 144 | + | wantIP: "192.168.1.100", | |||
| 145 | + | }, | |||
| 146 | + | { | |||
| 147 | + | name: "from X-Forwarded-For single IP", | |||
| 148 | + | remoteAddr: "127.0.0.1:12345", | |||
| 149 | + | xff: "203.0.113.50", | |||
| 150 | + | wantIP: "203.0.113.50", | |||
| 151 | + | }, | |||
| 152 | + | { | |||
| 153 | + | name: "from X-Forwarded-For multiple IPs", | |||
| 154 | + | remoteAddr: "127.0.0.1:12345", | |||
| 155 | + | xff: "203.0.113.50, 10.0.0.1, 172.16.0.1", | |||
| 156 | + | wantIP: "203.0.113.50", | |||
| 157 | + | }, | |||
| 158 | + | { | |||
| 159 | + | name: "from X-Real-IP", | |||
| 160 | + | remoteAddr: "127.0.0.1:12345", | |||
| 161 | + | xRealIP: "203.0.113.50", | |||
| 162 | + | wantIP: "203.0.113.50", | |||
| 163 | + | }, | |||
| 164 | + | { | |||
| 165 | + | name: "X-Forwarded-For takes precedence", | |||
| 166 | + | remoteAddr: "127.0.0.1:12345", | |||
| 167 | + | xff: "203.0.113.50", | |||
| 168 | + | xRealIP: "203.0.113.60", | |||
| 169 | + | wantIP: "203.0.113.50", | |||
| 170 | + | }, | |||
| 171 | + | } | |||
| 172 | + | ||||
| 173 | + | for _, tt := range tests { | |||
| 174 | + | t.Run(tt.name, func(t *testing.T) { | |||
| 175 | + | req := httptest.NewRequest(http.MethodGet, "/", nil) | |||
| 176 | + | req.RemoteAddr = tt.remoteAddr | |||
| 177 | + | if tt.xff != "" { | |||
| 178 | + | req.Header.Set("X-Forwarded-For", tt.xff) | |||
| 179 | + | } | |||
| 180 | + | if tt.xRealIP != "" { | |||
| 181 | + | req.Header.Set("X-Real-IP", tt.xRealIP) | |||
| 182 | + | } | |||
| 183 | + | ||||
| 184 | + | ip := getClientIP(req) | |||
| 185 | + | assert.Equal(t, tt.wantIP, ip) | |||
| 186 | + | }) | |||
| 187 | + | } | |||
| 188 | + | } | |||
| 189 | + | ||||
| 190 | + | func TestDefaultRateLimitConfig(t *testing.T) { | |||
| 191 | + | config := DefaultRateLimitConfig() | |||
| 192 | + | ||||
| 193 | + | require.True(t, config.Enabled) | |||
| 194 | + | assert.Equal(t, float64(500), config.GlobalRPS) | |||
| 195 | + | assert.Equal(t, float64(50), config.PerIPRPS) | |||
| 196 | + | assert.Equal(t, 100, config.Burst) | |||
| 197 | + | } | |||