1package http
2
3import (
4 "io"
5 "net/http"
6 "net/netip"
7 "strconv"
8 "strings"
9 "sync/atomic"
10 "time"
11
12 "tug/internal/platform/ratelimit"
13)
14
15type Rate = ratelimit.Rate
16
17type LimitConfig struct {
18 Authorize, LoginForm, Password, Invite, Callback, SSOLogin, LoginClaim Rate
19 Entries int
20 TokenSlots, SigninSlots, MaxWaiters int
21 SlotWait time.Duration
22}
23
24var DefaultLimits = LimitConfig{
25 Authorize: Rate{Burst: 20, Every: 6 * time.Second},
26 LoginForm: Rate{Burst: 20, Every: 6 * time.Second},
27 Password: Rate{Burst: 10, Every: 10 * time.Second},
28 Invite: Rate{Burst: 10, Every: 10 * time.Second},
29 Callback: Rate{Burst: 20, Every: 6 * time.Second},
30 SSOLogin: Rate{Burst: 30, Every: 3 * time.Second},
31 LoginClaim: Rate{Burst: 20, Every: 6 * time.Second},
32 Entries: 20000,
33 TokenSlots: 8,
34 SigninSlots: 8,
35 MaxWaiters: 64,
36 SlotWait: 10 * time.Second,
37}
38
39type rateClass int
40
41const (
42 rateNone rateClass = iota
43 rateAuthorize
44 rateLoginForm
45 ratePassword
46 rateInvite
47 rateCallback
48 rateSSOLogin
49 rateLoginClaim
50)
51
52type limitKey struct {
53 class rateClass
54 prefix netip.Prefix
55}
56
57type gate struct {
58 slots chan struct{}
59 waiting atomic.Int64
60}
61
62type Limits struct {
63 cfg LimitConfig
64 table *ratelimit.Table[limitKey]
65 token *gate
66 signin *gate
67}
68
69func NewLimits(cfg LimitConfig) *Limits { return newLimits(cfg, time.Now) }
70
71func newLimits(cfg LimitConfig, now func() time.Time) *Limits {
72 return &Limits{cfg: cfg, table: ratelimit.New[limitKey](cfg.Entries, now), token: newGate(cfg.TokenSlots), signin: newGate(cfg.SigninSlots)}
73}
74
75type RouterOption func(*router)
76
77func WithLimits(l *Limits) RouterOption {
78 return func(rt *router) { rt.limits = l }
79}
80
81func (l *Limits) rate(c rateClass) Rate {
82 switch c {
83 case rateAuthorize:
84 return l.cfg.Authorize
85 case rateLoginForm:
86 return l.cfg.LoginForm
87 case ratePassword:
88 return l.cfg.Password
89 case rateInvite:
90 return l.cfg.Invite
91 case rateCallback:
92 return l.cfg.Callback
93 case rateSSOLogin:
94 return l.cfg.SSOLogin
95 case rateLoginClaim:
96 return l.cfg.LoginClaim
97 }
98 return Rate{}
99}
100
101func classify(class hostClass, method, p string) rateClass {
102 get, post := method == http.MethodGet, method == http.MethodPost
103 switch class {
104 case classIDP:
105 if get && strings.HasPrefix(p, "/auth/") {
106 return rateAuthorize
107 }
108 case classAccount, classAdmin:
109 switch {
110 case get && p == "/oidc/callback":
111 return rateCallback
112 case post && p == "/sudo":
113 return ratePassword
114 case class == classAdmin:
115 case get && p == "/login/claim":
116 return rateLoginClaim
117 case p == "/login" && get, post && (p == "/login/google" || p == "/google/link"):
118 return rateLoginForm
119 case get && p == "/link/google/callback":
120 return rateCallback
121 case post && (p == "/login" || p == "/password"):
122 return ratePassword
123 case (get || post) && strings.HasPrefix(p, "/invite/"):
124 return rateInvite
125 }
126 case classApp:
127 switch {
128 case get && p == "/.sso/login":
129 return rateSSOLogin
130 case get && p == "/.sso/callback":
131 return rateCallback
132 }
133 }
134 return rateNone
135}
136
137func clientPrefix(r *http.Request) netip.Prefix {
138 return ratelimit.Prefix(ClientOf(r.Context()).Addr)
139}
140
141func (l *Limits) size() int { return l.table.Len() }
142
143func retryAfter(d time.Duration) string {
144 s := int64((d + time.Second - 1) / time.Second)
145 return strconv.FormatInt(max(s, 1), 10)
146}
147
148func (l *Limits) admit(w http.ResponseWriter, r *http.Request, class hostClass, page ErrorPage, o Observer) (func(), bool) {
149 if c := classify(class, r.Method, r.URL.Path); c != rateNone {
150 if wait, ok := l.table.Take(limitKey{class: c, prefix: clientPrefix(r)}, l.rate(c)); !ok {
151 o.Limited(c.name())
152 w.Header().Set("Retry-After", retryAfter(wait))
153 page(w, r, http.StatusTooManyRequests)
154 return nil, false
155 }
156 }
157 g := l.gate(class, r.Method, r.URL.Path)
158 if g == nil {
159 return func() {}, true
160 }
161 if l.acquire(r, g, o) {
162 return func() { <-g.slots }, true
163 }
164 o.Refused(l.gateName(g))
165 w.Header().Set("Retry-After", "1")
166 if g == l.token {
167 h := w.Header()
168 h.Set("Content-Type", "application/json")
169 h.Set("Cache-Control", "no-store")
170 w.WriteHeader(http.StatusServiceUnavailable)
171 io.WriteString(w, tokenBusy)
172 return nil, false
173 }
174 page(w, busyRequest(r), http.StatusServiceUnavailable)
175 return nil, false
176}
177
178const tokenBusy = `{"error":"temporarily_unavailable","error_description":"The server is busy. Try again in a moment."}`
179
180func newGate(n int) *gate {
181 if n <= 0 {
182 return nil
183 }
184 return &gate{slots: make(chan struct{}, n)}
185}
186
187func (l *Limits) gate(class hostClass, method, p string) *gate {
188 get, post := method == http.MethodGet, method == http.MethodPost
189 switch {
190 case class == classIDP && post && p == "/token":
191 return l.token
192 case class == classAccount && post && (p == "/login" || p == "/login/google" || p == "/google/link"),
193 class == classAccount && get && p == "/link/google/callback",
194 (class == classAccount || class == classAdmin) && get && p == "/oidc/callback",
195 class == classApp && get && p == "/.sso/callback":
196 return l.signin
197 }
198 return nil
199}
200
201func (l *Limits) acquire(r *http.Request, g *gate, o Observer) bool {
202 select {
203 case g.slots <- struct{}{}:
204 return true
205 default:
206 }
207 if g.waiting.Add(1) > int64(l.cfg.MaxWaiters) {
208 g.waiting.Add(-1)
209 return false
210 }
211 o.Waited(l.gateName(g))
212 defer g.waiting.Add(-1)
213 t := time.NewTimer(l.cfg.SlotWait)
214 defer t.Stop()
215 select {
216 case g.slots <- struct{}{}:
217 return true
218 case <-t.C:
219 case <-r.Context().Done():
220 }
221 return false
222}