diff options
author | Kamil Trzciński <ayufan@ayufan.eu> | 2020-10-20 13:26:29 +0300 |
---|---|---|
committer | Kamil Trzciński <ayufan@ayufan.eu> | 2020-10-20 13:26:29 +0300 |
commit | 7f8e9bd39def730616a4c7d1d5f00ee6ca9ea76a (patch) | |
tree | 455beeb7ac317059afcad0167bcba4a23ea1aec7 /internal | |
parent | 9cf62c0fc1f31a6e175bb3e8b2321ca19584dee3 (diff) |
Add Host and SNI-based rate limiting
This adds a per-process rate limiting
of the incoming requests and connections.
This assume two:
- Requests generate a pressure on Object Storage
- New TLS connections generate a pressure on CPU
due to TLS handshake (generating and exchanging
asymmetric keys)
Diffstat (limited to 'internal')
-rw-r--r-- | internal/rate_limiting/http_handler.go | 32 | ||||
-rw-r--r-- | internal/rate_limiting/rate_limiting.go | 62 |
2 files changed, 94 insertions, 0 deletions
diff --git a/internal/rate_limiting/http_handler.go b/internal/rate_limiting/http_handler.go new file mode 100644 index 00000000..d6341781 --- /dev/null +++ b/internal/rate_limiting/http_handler.go @@ -0,0 +1,32 @@ +package rate_limiting + +import ( + "crypto/tls" + "errors" + "net/http" + + "gitlab.com/gitlab-org/gitlab-pages/internal/tlsconfig" +) + +func (r *RateLimiting) LimitHostHandler(handler http.Handler) http.Handler { + fn := func(rw http.ResponseWriter, req *http.Request) { + if r.Allow(req.Host) { + handler.ServeHTTP(rw, req) + return + } + + rw.WriteHeader(http.StatusTooManyRequests) + } + + return http.HandlerFunc(fn) +} + +func (r *RateLimiting) LimitServeTLS(handler tlsconfig.GetCertificateFunc) tlsconfig.GetCertificateFunc { + return func(ch *tls.ClientHelloInfo) (*tls.Certificate, error) { + if r.Allow(ch.ServerName) { + return handler(ch) + } + + return nil, errors.New("rate limited") + } +} diff --git a/internal/rate_limiting/rate_limiting.go b/internal/rate_limiting/rate_limiting.go new file mode 100644 index 00000000..d69dd873 --- /dev/null +++ b/internal/rate_limiting/rate_limiting.go @@ -0,0 +1,62 @@ +package rate_limiting + +import ( + "time" + + "github.com/patrickmn/go-cache" + "golang.org/x/time/rate" +) + +type rateLimit struct { + *rate.Limiter +} + +type RateLimiting struct { + cache *cache.Cache + + window time.Duration + limit uint +} + +func NewRateLimiting(window time.Duration, limit uint) *RateLimiting { + return &RateLimiting{ + cache: cache.New(window*2, window), + window: window, + limit: limit, + } +} + +func (r *RateLimiting) newRateLimiter() rateLimit { + // we divide a window by amount of requests + // the bucket is refilled every interval + // allowing to consume up to the defined `limit` + everyNs := r.window.Nanoseconds() / int64(r.limit) + every := time.Duration(everyNs) + + return rateLimit{ + rate.NewLimiter(rate.Every(every), int(r.limit)), + } +} + +func (r *RateLimiting) findOrCreate(key string) rateLimit { + for { + // try to get existing item + if item, expiry, found := r.cache.GetWithExpiration(key); found { + // extend item window + if time.Until(expiry) > r.window { + r.cache.SetDefault(key, item) + } + + return item.(rateLimit) + } + + // add a new item + if rateLimiter := r.newRateLimiter(); r.cache.Add(key, rateLimiter, cache.DefaultExpiration) == nil { + return rateLimiter + } + } +} + +func (r *RateLimiting) Allow(key string) bool { + return r.findOrCreate(key).Allow() +} |