@@ -4,11 +4,13 @@ import (
44 "fmt"
55 "net/http"
66 "strconv"
7+ "sync/atomic"
78 "time"
89
910 "github.com/go-chi/httprate"
1011 "golang.org/x/xerrors"
1112
13+ "github.com/coder/coder/v2/coderd/aibridge"
1214 "github.com/coder/coder/v2/coderd/database"
1315 "github.com/coder/coder/v2/coderd/httpapi"
1416 "github.com/coder/coder/v2/coderd/rbac"
@@ -70,3 +72,72 @@ func RateLimit(count int, window time.Duration) func(http.Handler) http.Handler
7072 }),
7173 )
7274}
75+
76+ // RateLimitByAuthToken returns a handler that limits requests based on the
77+ // authentication token in the request.
78+ //
79+ // This differs from [RateLimit] in several ways:
80+ // - It extracts the token directly from request headers (Authorization Bearer
81+ // or X-Api-Key) rather than from the request context, making it suitable for
82+ // endpoints that handle authentication internally (like AI Bridge) rather than
83+ // via [ExtractAPIKeyMW] middleware.
84+ // - It does not support the bypass header for Owners.
85+ // - It does not key by endpoint, so the limit applies across all endpoints using
86+ // this middleware.
87+ // - It includes a Retry-After header in 429 responses for backpressure signaling.
88+ //
89+ // If no token is found in the headers, it falls back to rate limiting by IP address.
90+ func RateLimitByAuthToken (count int , window time.Duration ) func (http.Handler ) http.Handler {
91+ if count <= 0 {
92+ return func (handler http.Handler ) http.Handler {
93+ return handler
94+ }
95+ }
96+
97+ return httprate .Limit (
98+ count ,
99+ window ,
100+ httprate .WithKeyFuncs (func (r * http.Request ) (string , error ) {
101+ // Try to extract auth token for per-user rate limiting using
102+ // AI provider authentication headers (Authorization Bearer or X-Api-Key).
103+ if token := aibridge .ExtractAuthToken (r .Header ); token != "" {
104+ return token , nil
105+ }
106+ // Fall back to IP-based rate limiting if no token present.
107+ return httprate .KeyByIP (r )
108+ }),
109+ httprate .WithLimitHandler (func (w http.ResponseWriter , r * http.Request ) {
110+ // Add Retry-After header for backpressure signaling.
111+ w .Header ().Set ("Retry-After" , fmt .Sprintf ("%d" , int (window .Seconds ())))
112+ httpapi .Write (r .Context (), w , http .StatusTooManyRequests , codersdk.Response {
113+ Message : "You've been rate limited. Please try again later." ,
114+ })
115+ }),
116+ )
117+ }
118+
119+ // ConcurrencyLimit returns a handler that limits the number of concurrent
120+ // requests. When the limit is exceeded, it returns HTTP 503 Service Unavailable.
121+ func ConcurrencyLimit (maxConcurrent int64 , resourceName string ) func (http.Handler ) http.Handler {
122+ if maxConcurrent <= 0 {
123+ return func (handler http.Handler ) http.Handler {
124+ return handler
125+ }
126+ }
127+
128+ var current atomic.Int64
129+ return func (next http.Handler ) http.Handler {
130+ return http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
131+ c := current .Add (1 )
132+ defer current .Add (- 1 )
133+
134+ if c > maxConcurrent {
135+ httpapi .Write (r .Context (), w , http .StatusServiceUnavailable , codersdk.Response {
136+ Message : fmt .Sprintf ("%s is currently at capacity. Please try again later." , resourceName ),
137+ })
138+ return
139+ }
140+ next .ServeHTTP (w , r )
141+ })
142+ }
143+ }
0 commit comments