From fea0c22b482a3910941234fafc27c9defe5a8b8b Mon Sep 17 00:00:00 2001
From: can <can@canx.me>
Date: Mon, 23 Mar 2020 17:32:04 +0800
Subject: [PATCH] Basic rate limit

---
 api-server.go           |  2 ++
 api/generic-handlers.go | 44 ++++++++++++++++++++++++++++++++++++++++-
 error/api-errors.go     |  6 ++++++
 go.mod                  |  6 ++++--
 4 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/api-server.go b/api-server.go
index 2d6382b71..54d30c836 100644
--- a/api-server.go
+++ b/api-server.go
@@ -71,6 +71,8 @@ func configureServerHandler(c *ServerConfig) http.Handler {
 
 		api.NewAccessLogHandler,
 
+		api.SetQosHandler,
+
 		api.SetGenerateContextHandler,
 
 		api.SetRequestIdHandler,
diff --git a/api/generic-handlers.go b/api/generic-handlers.go
index a9729031d..7d6c3e716 100644
--- a/api/generic-handlers.go
+++ b/api/generic-handlers.go
@@ -19,6 +19,8 @@ package api
 import (
 	"context"
 	"fmt"
+	"github.com/go-redis/redis/v7"
+	"github.com/go-redis/redis_rate/v8"
 	"net/http"
 	"strings"
 
@@ -273,6 +275,46 @@ func InReservedOrigins(origin string) bool {
 	return false
 }
 
+type QosHandler struct {
+	handler     http.Handler
+	meta        *meta.Meta
+	rateLimiter *redis_rate.Limiter
+}
+
+func (h QosHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	ctx := getRequestContext(r)
+	if len(ctx.BucketName) == 0 {
+		h.handler.ServeHTTP(w, r)
+		return
+	}
+	k := fmt.Sprintf("bucket_qps_%s", ctx.BucketName)
+	result, err := h.rateLimiter.Allow(k, redis_rate.PerSecond(10000))
+	if err == nil && !result.Allowed {
+		WriteErrorResponse(w, r, ErrRequestLimitExceeded)
+		return
+	}
+	if err != nil {
+		ctx.Logger.Error("rateLimiter:", err)
+		WriteErrorResponse(w, r, ErrInternalError)
+		return
+	}
+	h.handler.ServeHTTP(w, r)
+}
+
+func SetQosHandler(h http.Handler, meta *meta.Meta) http.Handler {
+	rdb := redis.NewClient(&redis.Options{
+		Addr:     helper.CONFIG.RedisAddress,
+		Password: helper.CONFIG.RedisPassword,
+	})
+	limiter := redis_rate.NewLimiter(rdb)
+	qos := QosHandler{
+		handler:     h,
+		meta:        meta,
+		rateLimiter: limiter,
+	}
+	return qos
+}
+
 //// helpers
 
 func GetBucketAndObjectInfoFromRequest(r *http.Request) (bucketName string, objectName string, isBucketDomain bool) {
@@ -300,7 +342,7 @@ func getRequestContext(r *http.Request) RequestContext {
 		return ctx
 	}
 	return RequestContext{
-		Logger: r.Context().Value(ContextLoggerKey).(log.Logger),
+		Logger:    r.Context().Value(ContextLoggerKey).(log.Logger),
 		RequestID: r.Context().Value(RequestIdKey).(string),
 	}
 }
diff --git a/error/api-errors.go b/error/api-errors.go
index 9628ac20b..0b34d3168 100755
--- a/error/api-errors.go
+++ b/error/api-errors.go
@@ -177,6 +177,7 @@ const (
 	ErrInvalidRestoreInfo
 	ErrCreateRestoreObject
 	ErrInvalidGlacierObject
+	ErrRequestLimitExceeded
 )
 
 // error code to APIError structure, these fields carry respective
@@ -828,6 +829,11 @@ var ErrorCodeResponse = map[ApiErrorCode]ApiErrorStruct{
 		Description:    "Create object thaw operation failed",
 		HttpStatusCode: http.StatusInternalServerError,
 	},
+	ErrRequestLimitExceeded: {
+		AwsErrorCode:   "ErrRequestLimitExceeded",
+		Description:    "Request limit exceeded",
+		HttpStatusCode: http.StatusTooManyRequests,
+	},
 }
 
 func (e ApiErrorCode) AwsErrorCode() string {
diff --git a/go.mod b/go.mod
index e1969a98a..8fa8974b3 100644
--- a/go.mod
+++ b/go.mod
@@ -30,9 +30,11 @@ require (
 	github.com/BurntSushi/toml v0.3.1
 	github.com/DATA-DOG/go-sqlmock v1.3.3
 	github.com/cep21/circuit v0.0.0-20181030180945-e893c027dc21
-	github.com/confluentinc/confluent-kafka-go v1.0.0 //indirect
+	github.com/confluentinc/confluent-kafka-go v1.0.0 //ct
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/dustin/go-humanize v1.0.0
+	github.com/go-redis/redis/v7 v7.0.0-beta.4
+	github.com/go-redis/redis_rate/v8 v8.0.0
 	github.com/go-sql-driver/mysql v1.4.1
 	github.com/golang/snappy v0.0.1
 	github.com/gomodule/redigo v1.7.0
@@ -42,7 +44,7 @@ require (
 	github.com/journeymidnight/radoshttpd v0.0.0-20190617133011-609666b51136
 	github.com/minio/highwayhash v1.0.0
 	github.com/prometheus/client_golang v0.9.3-0.20190127221311-3c4408c8b829
-	github.com/stretchr/testify v1.3.0
+	github.com/stretchr/testify v1.4.0
 	github.com/ugorji/go v1.1.4
 	github.com/xxtea/xxtea-go v0.0.0-20170828040851-35c4b17eecf6
 	golang.org/x/crypto v0.0.0-20190325154230-a5d413f7728c // indirect