76 lines
1.8 KiB
Go
76 lines
1.8 KiB
Go
package api
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func resetAgentWSRateLim(t *testing.T) {
|
|
t.Helper()
|
|
agentWSRateLim.mu.Lock()
|
|
agentWSRateLim.attempts = make(map[string][]time.Time)
|
|
agentWSRateLim.mu.Unlock()
|
|
}
|
|
|
|
func TestAllowAgentWSUpgradeRateLimit(t *testing.T) {
|
|
t.Run("rejects 31st attempt within window", func(t *testing.T) {
|
|
resetAgentWSRateLim(t)
|
|
ip := "203.0.113.42"
|
|
|
|
for i := 1; i <= agentWSRateLimitMax; i++ {
|
|
if !allowAgentWSUpgrade(ip) {
|
|
t.Fatalf("attempt %d: expected allow, got reject", i)
|
|
}
|
|
}
|
|
if allowAgentWSUpgrade(ip) {
|
|
t.Fatal("31st attempt: expected reject, got allow")
|
|
}
|
|
if allowAgentWSUpgrade(ip) {
|
|
t.Fatal("32nd attempt: expected reject, got allow")
|
|
}
|
|
})
|
|
|
|
t.Run("empty IP bypasses limit", func(t *testing.T) {
|
|
resetAgentWSRateLim(t)
|
|
|
|
for i := 1; i <= agentWSRateLimitMax+5; i++ {
|
|
if !allowAgentWSUpgrade("") {
|
|
t.Fatalf("empty IP attempt %d: expected allow, got reject", i)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("stale attempts outside window are pruned", func(t *testing.T) {
|
|
resetAgentWSRateLim(t)
|
|
ip := "198.51.100.7"
|
|
stale := time.Now().Add(-agentWSRateLimitWindow - time.Second)
|
|
|
|
agentWSRateLim.mu.Lock()
|
|
staleAttempts := make([]time.Time, agentWSRateLimitMax)
|
|
for i := range staleAttempts {
|
|
staleAttempts[i] = stale
|
|
}
|
|
agentWSRateLim.attempts[ip] = staleAttempts
|
|
agentWSRateLim.mu.Unlock()
|
|
|
|
if !allowAgentWSUpgrade(ip) {
|
|
t.Fatal("expected allow after stale attempts pruned")
|
|
}
|
|
})
|
|
|
|
t.Run("different IPs have independent limits", func(t *testing.T) {
|
|
resetAgentWSRateLim(t)
|
|
ipA := "192.0.2.1"
|
|
ipB := "192.0.2.2"
|
|
|
|
for i := 1; i <= agentWSRateLimitMax; i++ {
|
|
if !allowAgentWSUpgrade(ipA) {
|
|
t.Fatalf("ipA attempt %d: expected allow, got reject", i)
|
|
}
|
|
}
|
|
if !allowAgentWSUpgrade(ipB) {
|
|
t.Fatal("ipB first attempt: expected allow after ipA exhausted")
|
|
}
|
|
})
|
|
}
|