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)

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()
+}