mirror of
https://github.com/rcourtman/Pulse.git
synced 2026-02-18 00:17:39 +01:00
a55bdb7a3a
- Require admin + settings:write scope for setup-script-url endpoint - Add license enforcement for long-term metrics (30d/90d require Pro) - Add downsampling step calculation for metrics history queries - Add isContainerSSHRestricted helper for SSH restriction checks - Clean up temperature proxy references from config handlers - Minor OIDC and rate limit improvements
437 lines
9.8 KiB
Go
437 lines
9.8 KiB
Go
package api
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestNewRateLimiter(t *testing.T) {
|
|
rl := NewRateLimiter(5, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
if rl == nil {
|
|
t.Fatal("NewRateLimiter returned nil")
|
|
}
|
|
if rl.limit != 5 {
|
|
t.Errorf("limit = %d, want 5", rl.limit)
|
|
}
|
|
if rl.window != time.Minute {
|
|
t.Errorf("window = %v, want %v", rl.window, time.Minute)
|
|
}
|
|
if rl.attempts == nil {
|
|
t.Error("attempts map is nil")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Allow_Basic(t *testing.T) {
|
|
rl := NewRateLimiter(3, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
ip := "192.168.1.1"
|
|
|
|
// First 3 attempts should be allowed
|
|
for i := 1; i <= 3; i++ {
|
|
if !rl.Allow(ip) {
|
|
t.Errorf("attempt %d should be allowed", i)
|
|
}
|
|
}
|
|
|
|
// 4th attempt should be denied
|
|
if rl.Allow(ip) {
|
|
t.Error("4th attempt should be denied")
|
|
}
|
|
|
|
// 5th attempt should also be denied
|
|
if rl.Allow(ip) {
|
|
t.Error("5th attempt should be denied")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Allow_DifferentIPs(t *testing.T) {
|
|
rl := NewRateLimiter(2, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
ip1 := "192.168.1.1"
|
|
ip2 := "192.168.1.2"
|
|
|
|
// Both IPs should have independent limits
|
|
if !rl.Allow(ip1) {
|
|
t.Error("ip1 attempt 1 should be allowed")
|
|
}
|
|
if !rl.Allow(ip2) {
|
|
t.Error("ip2 attempt 1 should be allowed")
|
|
}
|
|
if !rl.Allow(ip1) {
|
|
t.Error("ip1 attempt 2 should be allowed")
|
|
}
|
|
if !rl.Allow(ip2) {
|
|
t.Error("ip2 attempt 2 should be allowed")
|
|
}
|
|
|
|
// Both should now be at their limit
|
|
if rl.Allow(ip1) {
|
|
t.Error("ip1 attempt 3 should be denied")
|
|
}
|
|
if rl.Allow(ip2) {
|
|
t.Error("ip2 attempt 3 should be denied")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Allow_WindowExpiry(t *testing.T) {
|
|
// Use a very short window for testing
|
|
rl := NewRateLimiter(2, 50*time.Millisecond)
|
|
defer rl.Stop()
|
|
|
|
ip := "192.168.1.1"
|
|
|
|
// Use up the limit
|
|
if !rl.Allow(ip) {
|
|
t.Error("attempt 1 should be allowed")
|
|
}
|
|
if !rl.Allow(ip) {
|
|
t.Error("attempt 2 should be allowed")
|
|
}
|
|
if rl.Allow(ip) {
|
|
t.Error("attempt 3 should be denied")
|
|
}
|
|
|
|
// Wait for window to expire
|
|
time.Sleep(60 * time.Millisecond)
|
|
|
|
// Should be allowed again
|
|
if !rl.Allow(ip) {
|
|
t.Error("attempt after window expiry should be allowed")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Allow_SlidingWindow(t *testing.T) {
|
|
// Test that the window is truly sliding (not fixed intervals)
|
|
rl := NewRateLimiter(2, 100*time.Millisecond)
|
|
defer rl.Stop()
|
|
|
|
ip := "192.168.1.1"
|
|
|
|
// First attempt
|
|
if !rl.Allow(ip) {
|
|
t.Error("attempt 1 should be allowed")
|
|
}
|
|
|
|
// Wait half the window
|
|
time.Sleep(60 * time.Millisecond)
|
|
|
|
// Second attempt
|
|
if !rl.Allow(ip) {
|
|
t.Error("attempt 2 should be allowed")
|
|
}
|
|
|
|
// Third attempt should be denied (both still in window)
|
|
if rl.Allow(ip) {
|
|
t.Error("attempt 3 should be denied")
|
|
}
|
|
|
|
// Wait for first attempt to expire (another 50ms should be enough)
|
|
time.Sleep(50 * time.Millisecond)
|
|
|
|
// Now should be allowed (first attempt expired, second still valid)
|
|
if !rl.Allow(ip) {
|
|
t.Error("attempt 4 should be allowed after first expires")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Allow_ZeroLimit(t *testing.T) {
|
|
rl := NewRateLimiter(0, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
// With limit 0, nothing should be allowed
|
|
if rl.Allow("192.168.1.1") {
|
|
t.Error("should deny with limit 0")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Allow_LargeLimit(t *testing.T) {
|
|
rl := NewRateLimiter(1000, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
ip := "192.168.1.1"
|
|
|
|
// All 1000 attempts should be allowed
|
|
for i := 1; i <= 1000; i++ {
|
|
if !rl.Allow(ip) {
|
|
t.Errorf("attempt %d should be allowed", i)
|
|
}
|
|
}
|
|
|
|
// 1001st should be denied
|
|
if rl.Allow(ip) {
|
|
t.Error("attempt 1001 should be denied")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Allow_EmptyIP(t *testing.T) {
|
|
rl := NewRateLimiter(2, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
// Empty string IP should still work as a valid key
|
|
if !rl.Allow("") {
|
|
t.Error("empty IP attempt 1 should be allowed")
|
|
}
|
|
if !rl.Allow("") {
|
|
t.Error("empty IP attempt 2 should be allowed")
|
|
}
|
|
if rl.Allow("") {
|
|
t.Error("empty IP attempt 3 should be denied")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Cleanup(t *testing.T) {
|
|
rl := NewRateLimiter(5, 50*time.Millisecond)
|
|
defer rl.Stop()
|
|
|
|
// Add some attempts
|
|
rl.Allow("192.168.1.1")
|
|
rl.Allow("192.168.1.2")
|
|
rl.Allow("192.168.1.3")
|
|
|
|
// Verify attempts are tracked
|
|
rl.mu.RLock()
|
|
beforeCount := len(rl.attempts)
|
|
rl.mu.RUnlock()
|
|
|
|
if beforeCount != 3 {
|
|
t.Errorf("expected 3 IPs tracked, got %d", beforeCount)
|
|
}
|
|
|
|
// Wait for window to expire
|
|
time.Sleep(60 * time.Millisecond)
|
|
|
|
// Manually trigger cleanup
|
|
rl.cleanup()
|
|
|
|
// All should be cleaned up
|
|
rl.mu.RLock()
|
|
afterCount := len(rl.attempts)
|
|
rl.mu.RUnlock()
|
|
|
|
if afterCount != 0 {
|
|
t.Errorf("expected 0 IPs after cleanup, got %d", afterCount)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Cleanup_PartialExpiry(t *testing.T) {
|
|
rl := NewRateLimiter(5, 50*time.Millisecond)
|
|
defer rl.Stop()
|
|
|
|
// Add attempt for IP1
|
|
rl.Allow("192.168.1.1")
|
|
|
|
// Wait for it to expire
|
|
time.Sleep(60 * time.Millisecond)
|
|
|
|
// Add attempt for IP2 (this one is fresh)
|
|
rl.Allow("192.168.1.2")
|
|
|
|
// Cleanup should remove IP1 but keep IP2
|
|
rl.cleanup()
|
|
|
|
rl.mu.RLock()
|
|
_, hasIP1 := rl.attempts["192.168.1.1"]
|
|
_, hasIP2 := rl.attempts["192.168.1.2"]
|
|
rl.mu.RUnlock()
|
|
|
|
if hasIP1 {
|
|
t.Error("IP1 should have been cleaned up")
|
|
}
|
|
if !hasIP2 {
|
|
t.Error("IP2 should still be present")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Concurrent(t *testing.T) {
|
|
rl := NewRateLimiter(100, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
var wg sync.WaitGroup
|
|
numGoroutines := 10
|
|
attemptsPerGoroutine := 20
|
|
|
|
// Track results per goroutine
|
|
results := make([][]bool, numGoroutines)
|
|
for i := 0; i < numGoroutines; i++ {
|
|
results[i] = make([]bool, attemptsPerGoroutine)
|
|
}
|
|
|
|
// Concurrent access from same IP
|
|
for i := 0; i < numGoroutines; i++ {
|
|
wg.Add(1)
|
|
go func(idx int) {
|
|
defer wg.Done()
|
|
for j := 0; j < attemptsPerGoroutine; j++ {
|
|
results[idx][j] = rl.Allow("shared-ip")
|
|
}
|
|
}(i)
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
// Count allowed attempts
|
|
allowed := 0
|
|
for i := 0; i < numGoroutines; i++ {
|
|
for j := 0; j < attemptsPerGoroutine; j++ {
|
|
if results[i][j] {
|
|
allowed++
|
|
}
|
|
}
|
|
}
|
|
|
|
// Should have exactly 100 allowed (the limit)
|
|
if allowed != 100 {
|
|
t.Errorf("expected exactly 100 allowed attempts, got %d", allowed)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Stop(t *testing.T) {
|
|
rl := NewRateLimiter(5, time.Minute)
|
|
|
|
// Should be able to stop without issues
|
|
rl.Stop()
|
|
|
|
// After stop, Allow should still work (stop only affects cleanup goroutine)
|
|
if !rl.Allow("192.168.1.1") {
|
|
t.Error("Allow should still work after Stop")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Middleware_Allowed(t *testing.T) {
|
|
rl := NewRateLimiter(5, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
handlerCalled := false
|
|
handler := func(w http.ResponseWriter, r *http.Request) {
|
|
handlerCalled = true
|
|
w.WriteHeader(http.StatusOK)
|
|
}
|
|
|
|
wrapped := rl.Middleware(handler)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
|
req.RemoteAddr = "192.168.1.1:12345"
|
|
w := httptest.NewRecorder()
|
|
|
|
wrapped(w, req)
|
|
|
|
if !handlerCalled {
|
|
t.Error("handler should have been called")
|
|
}
|
|
if w.Code != http.StatusOK {
|
|
t.Errorf("status = %d, want %d", w.Code, http.StatusOK)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Middleware_Denied(t *testing.T) {
|
|
rl := NewRateLimiter(1, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
handlerCalls := 0
|
|
handler := func(w http.ResponseWriter, r *http.Request) {
|
|
handlerCalls++
|
|
w.WriteHeader(http.StatusOK)
|
|
}
|
|
|
|
wrapped := rl.Middleware(handler)
|
|
|
|
// First request should succeed
|
|
req1 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
|
req1.RemoteAddr = "192.168.1.1:12345"
|
|
w1 := httptest.NewRecorder()
|
|
wrapped(w1, req1)
|
|
|
|
if w1.Code != http.StatusOK {
|
|
t.Errorf("first request status = %d, want %d", w1.Code, http.StatusOK)
|
|
}
|
|
|
|
// Second request should be rate limited
|
|
req2 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
|
req2.RemoteAddr = "192.168.1.1:12345"
|
|
w2 := httptest.NewRecorder()
|
|
wrapped(w2, req2)
|
|
|
|
if w2.Code != http.StatusTooManyRequests {
|
|
t.Errorf("second request status = %d, want %d", w2.Code, http.StatusTooManyRequests)
|
|
}
|
|
|
|
if handlerCalls != 1 {
|
|
t.Errorf("handler called %d times, want 1", handlerCalls)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Middleware_XForwardedFor(t *testing.T) {
|
|
t.Setenv("PULSE_TRUSTED_PROXY_CIDRS", "127.0.0.1/32")
|
|
resetTrustedProxyConfig()
|
|
|
|
rl := NewRateLimiter(1, time.Minute)
|
|
defer rl.Stop()
|
|
|
|
handlerCalls := 0
|
|
handler := func(w http.ResponseWriter, r *http.Request) {
|
|
handlerCalls++
|
|
}
|
|
|
|
wrapped := rl.Middleware(handler)
|
|
|
|
// Request with X-Forwarded-For header
|
|
req1 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
|
req1.RemoteAddr = "127.0.0.1:12345"
|
|
req1.Header.Set("X-Forwarded-For", "203.0.113.1")
|
|
w1 := httptest.NewRecorder()
|
|
wrapped(w1, req1)
|
|
|
|
// Second request with same X-Forwarded-For should be rate limited
|
|
req2 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
|
req2.RemoteAddr = "127.0.0.1:12345"
|
|
req2.Header.Set("X-Forwarded-For", "203.0.113.1")
|
|
w2 := httptest.NewRecorder()
|
|
wrapped(w2, req2)
|
|
|
|
if w2.Code != http.StatusTooManyRequests {
|
|
t.Errorf("second request with same X-Forwarded-For should be rate limited, got %d", w2.Code)
|
|
}
|
|
|
|
// Request with different X-Forwarded-For should be allowed
|
|
req3 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
|
req3.RemoteAddr = "127.0.0.1:12345"
|
|
req3.Header.Set("X-Forwarded-For", "203.0.113.2")
|
|
w3 := httptest.NewRecorder()
|
|
wrapped(w3, req3)
|
|
|
|
if w3.Code == http.StatusTooManyRequests {
|
|
t.Error("request with different X-Forwarded-For should be allowed")
|
|
}
|
|
|
|
if handlerCalls != 2 {
|
|
t.Errorf("handler called %d times, want 2", handlerCalls)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Middleware_ResponseBody(t *testing.T) {
|
|
rl := NewRateLimiter(0, time.Minute) // Immediately rate limit
|
|
defer rl.Stop()
|
|
|
|
handler := func(w http.ResponseWriter, r *http.Request) {}
|
|
wrapped := rl.Middleware(handler)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
|
req.RemoteAddr = "192.168.1.1:12345"
|
|
w := httptest.NewRecorder()
|
|
|
|
wrapped(w, req)
|
|
|
|
body := w.Body.String()
|
|
expected := "Rate limit exceeded. Please try again later.\n"
|
|
if body != expected {
|
|
t.Errorf("body = %q, want %q", body, expected)
|
|
}
|
|
}
|