RateLimit returns a handler that limits requests per-minute based on IP, endpoint, and user ID (if available).
(count int, window time.Duration)
| 21 | // RateLimit returns a handler that limits requests per-minute based |
| 22 | // on IP, endpoint, and user ID (if available). |
| 23 | func RateLimit(count int, window time.Duration) func(http.Handler) http.Handler { |
| 24 | // -1 is no rate limit |
| 25 | if count <= 0 { |
| 26 | return func(handler http.Handler) http.Handler { |
| 27 | return handler |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | return httprate.Limit( |
| 32 | count, |
| 33 | window, |
| 34 | httprate.WithKeyFuncs(func(r *http.Request) (string, error) { |
| 35 | // Identify the caller. We check two sources: |
| 36 | // |
| 37 | // 1. apiKeyPrecheckedContextKey — set by PrecheckAPIKey |
| 38 | // at the root of the router. Only fully validated |
| 39 | // keys are used. |
| 40 | // 2. apiKeyContextKey — set by ExtractAPIKeyMW if it |
| 41 | // has already run (e.g. unit tests, workspace-app |
| 42 | // routes that don't go through PrecheckAPIKey). |
| 43 | // |
| 44 | // If neither is present, fall back to IP. |
| 45 | var userID string |
| 46 | var subject *rbac.Subject |
| 47 | |
| 48 | if pc, ok := r.Context().Value(apiKeyPrecheckedContextKey{}).(APIKeyPrechecked); ok && pc.Result != nil { |
| 49 | userID = pc.Result.Key.UserID.String() |
| 50 | subject = &pc.Result.Subject |
| 51 | } else if ak, ok := r.Context().Value(apiKeyContextKey{}).(database.APIKey); ok { |
| 52 | userID = ak.UserID.String() |
| 53 | if auth, ok := UserAuthorizationOptional(r.Context()); ok { |
| 54 | subject = &auth |
| 55 | } |
| 56 | } else { |
| 57 | return httprate.KeyByIP(r) |
| 58 | } |
| 59 | |
| 60 | if ok, _ := strconv.ParseBool(r.Header.Get(codersdk.BypassRatelimitHeader)); !ok { |
| 61 | // No bypass attempt, just rate limit by user. |
| 62 | return userID, nil |
| 63 | } |
| 64 | |
| 65 | // Allow Owner to bypass rate limiting for load tests |
| 66 | // and automation. We avoid using rbac.Authorizer since |
| 67 | // rego is CPU-intensive and undermines the |
| 68 | // DoS-prevention goal of the rate limiter. |
| 69 | if subject == nil { |
| 70 | // Can't verify roles — rate limit normally. |
| 71 | return userID, nil |
| 72 | } |
| 73 | for _, role := range subject.SafeRoleNames() { |
| 74 | if role == rbac.RoleOwner() { |
| 75 | // HACK: use a random key each time to |
| 76 | // de facto disable rate limiting. The |
| 77 | // httprate package has no support for |
| 78 | // selectively changing the limit for |
| 79 | // particular keys. |
| 80 | return cryptorand.String(16) |