-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
early proof-of-concept for providing pre-auth rate limiting at the Nomad servers, with the ability to observe which users are creating the load.
- Loading branch information
Showing
6 changed files
with
99 additions
and
8 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,78 @@ | ||
| package nomad | ||
|
|
||
| import ( | ||
| "context" | ||
| "time" | ||
|
|
||
| "github.com/armon/go-metrics" | ||
| "github.com/hashicorp/nomad/acl" | ||
| "github.com/hashicorp/nomad/nomad/structs" | ||
| limiter "github.com/sethvargo/go-limiter" | ||
| "github.com/sethvargo/go-limiter/memorystore" | ||
| ) | ||
|
|
||
| func (srv *Server) CheckRateLimit(limiter *RateLimiter, secretID, op string) error { | ||
| token, err := srv.ResolveSecretToken(secretID) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| return limiter.check(srv.shutdownCtx, token.AccessorID, op) | ||
| } | ||
|
|
||
| type RateLimiter struct { | ||
| endpoint string | ||
| write limiter.Store | ||
| read limiter.Store | ||
| list limiter.Store | ||
| } | ||
|
|
||
| func newRateLimiter(endpoint string) *RateLimiter { | ||
| // TODO: needs to take a configuration object so we can set tokens | ||
| // value | ||
| return &RateLimiter{ | ||
| endpoint: endpoint, | ||
| write: newRateLimiterStore(10), | ||
| read: newRateLimiterStore(10), | ||
| list: newRateLimiterStore(10), | ||
| } | ||
| } | ||
|
|
||
| func (r *RateLimiter) check(ctx context.Context, op, key string) error { | ||
| var tokens, remaining uint64 | ||
| var ok bool | ||
| var err error | ||
| switch op { | ||
| case acl.PolicyWrite: | ||
| tokens, remaining, _, ok, err = r.write.Take(ctx, key) | ||
| case acl.PolicyRead: | ||
| tokens, remaining, _, ok, err = r.read.Take(ctx, key) | ||
| case acl.PolicyList: | ||
| tokens, remaining, _, ok, err = r.list.Take(ctx, key) | ||
| } | ||
| used := tokens - remaining | ||
| metrics.IncrCounterWithLabels( | ||
| []string{"nomad", "rpc", r.endpoint, op}, 1, | ||
| []metrics.Label{{Name: "id", Value: key}}) | ||
| metrics.AddSampleWithLabels( | ||
| []string{"nomad", "rpc", r.endpoint, op, "used"}, float32(used), | ||
| []metrics.Label{{Name: "id", Value: key}}) | ||
|
|
||
| if err != nil { | ||
| return err // TODO: what's the source of these errors? | ||
| } | ||
| if !ok { | ||
| metrics.IncrCounterWithLabels([]string{"nomad", "rpc", r.endpoint, op, "limited"}, 1, | ||
| []metrics.Label{{Name: "id", Value: key}}) | ||
| return structs.ErrTooManyRequests | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| func newRateLimiterStore(tokens uint64) limiter.Store { | ||
| // note: the memorystore implementation never returns an error | ||
| store, _ := memorystore.New(&memorystore.Config{ | ||
| Tokens: tokens, // Number of tokens allowed per interval. | ||
| Interval: time.Minute, // Interval until tokens reset. | ||
| }) | ||
| return store | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters