Stateless auth proxy that converts AT Protocol native apps from public to confidential OAuth clients. Deploy once, get 180-day refresh tokens instead of 24-hour ones.
1package main
2
3import (
4 "net/http"
5 "strings"
6 "testing"
7)
8
9func TestRateLimiter_PerIP(t *testing.T) {
10 rl := newRateLimiter(3, 0)
11
12 for i := 0; i < 3; i++ {
13 if !rl.allow("1.2.3.4") {
14 t.Fatalf("request %d should be allowed", i+1)
15 }
16 }
17
18 if rl.allow("1.2.3.4") {
19 t.Error("4th request from same IP should be blocked")
20 }
21
22 if !rl.allow("5.6.7.8") {
23 t.Error("request from different IP should be allowed")
24 }
25}
26
27func TestRateLimiter_Global(t *testing.T) {
28 rl := newRateLimiter(0, 5)
29
30 for i := 0; i < 5; i++ {
31 if !rl.allow("1.2.3.4") {
32 t.Fatalf("request %d should be allowed", i+1)
33 }
34 }
35
36 if rl.allow("9.9.9.9") {
37 t.Error("request should be blocked by global limit")
38 }
39}
40
41func TestRateLimiter_Combined(t *testing.T) {
42 rl := newRateLimiter(2, 5)
43
44 if !rl.allow("1.1.1.1") || !rl.allow("1.1.1.1") {
45 t.Fatal("first IP should be allowed twice")
46 }
47 if rl.allow("1.1.1.1") {
48 t.Error("per-IP limit should block")
49 }
50
51 if !rl.allow("2.2.2.2") || !rl.allow("2.2.2.2") {
52 t.Fatal("second IP should be allowed twice")
53 }
54 if !rl.allow("3.3.3.3") {
55 t.Fatal("third IP should be allowed once")
56 }
57 if rl.allow("4.4.4.4") {
58 t.Error("global limit should block")
59 }
60}
61
62func TestRateLimiter_Disabled(t *testing.T) {
63 rl := newRateLimiter(0, 0)
64
65 for i := 0; i < 1000; i++ {
66 if !rl.allow("1.2.3.4") {
67 t.Fatalf("request %d should be allowed when rate limiting is disabled", i+1)
68 }
69 }
70}
71
72func TestRateLimitMiddleware_Returns429(t *testing.T) {
73 srv, cleanup := setupTestServerWithRateLimit(t, 2, 0)
74 defer cleanup()
75
76 for i := 0; i < 2; i++ {
77 resp, err := http.Post(srv.URL+"/oauth/token", "application/json", strings.NewReader(`not json`))
78 if err != nil {
79 t.Fatalf("request %d failed: %v", i+1, err)
80 }
81 resp.Body.Close()
82 if resp.StatusCode == http.StatusTooManyRequests {
83 t.Fatalf("request %d should not be rate limited", i+1)
84 }
85 }
86
87 resp, err := http.Post(srv.URL+"/oauth/token", "application/json", strings.NewReader(`not json`))
88 if err != nil {
89 t.Fatalf("request failed: %v", err)
90 }
91 defer resp.Body.Close()
92
93 if resp.StatusCode != http.StatusTooManyRequests {
94 t.Errorf("expected 429, got %d", resp.StatusCode)
95 }
96 if resp.Header.Get("Retry-After") != "60" {
97 t.Errorf("expected Retry-After: 60, got %s", resp.Header.Get("Retry-After"))
98 }
99}
100
101func TestClientIP(t *testing.T) {
102 tests := []struct {
103 name string
104 remoteAddr string
105 xff string
106 xRealIP string
107 trustProxyHeaders bool
108 expected string
109 }{
110 {name: "remote addr with port", remoteAddr: "1.2.3.4:12345", expected: "1.2.3.4"},
111 {name: "remote addr without port", remoteAddr: "1.2.3.4", expected: "1.2.3.4"},
112 {name: "xff ignored by default", remoteAddr: "9.9.9.9:1234", xff: "1.2.3.4", expected: "9.9.9.9"},
113 {name: "xff trusted when enabled", remoteAddr: "9.9.9.9:1234", xff: "1.2.3.4", trustProxyHeaders: true, expected: "1.2.3.4"},
114 {name: "xff multiple", remoteAddr: "9.9.9.9:1234", xff: "1.2.3.4, 5.6.7.8", trustProxyHeaders: true, expected: "1.2.3.4"},
115 {name: "x-real-ip fallback", remoteAddr: "9.9.9.9:1234", xRealIP: "1.2.3.4", trustProxyHeaders: true, expected: "1.2.3.4"},
116 {name: "invalid xff ignored", remoteAddr: "9.9.9.9:1234", xff: "not-an-ip", trustProxyHeaders: true, expected: "9.9.9.9"},
117 }
118
119 for _, tt := range tests {
120 t.Run(tt.name, func(t *testing.T) {
121 r := &http.Request{
122 RemoteAddr: tt.remoteAddr,
123 Header: http.Header{},
124 }
125 if tt.xff != "" {
126 r.Header.Set("X-Forwarded-For", tt.xff)
127 }
128 if tt.xRealIP != "" {
129 r.Header.Set("X-Real-IP", tt.xRealIP)
130 }
131
132 if got := clientIP(r, tt.trustProxyHeaders); got != tt.expected {
133 t.Errorf("clientIP() = %q, want %q", got, tt.expected)
134 }
135 })
136 }
137}