Merge pull request #85 from micanzhang/feature/exclude_key

drivers/middleware/gin: add ExcludedKey Option
ulule · Apr 7, 2020 · 4572462 · 4572462
2 parents 1366201 + 53840c1
commit 4572462
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/drivers/middleware/gin/middleware.go b/drivers/middleware/gin/middleware.go
@@ -14,6 +14,7 @@ type Middleware struct {
 	OnError        ErrorHandler
 	OnLimitReached LimitReachedHandler
 	KeyGetter      KeyGetter
+	ExcludedKey    func(string) bool
 }
 
 // NewMiddleware return a new instance of a gin middleware.
@@ -37,6 +38,11 @@ func NewMiddleware(limiter *limiter.Limiter, options ...Option) gin.HandlerFunc
 // Handle gin request.
 func (middleware *Middleware) Handle(c *gin.Context) {
 	key := middleware.KeyGetter(c)
+	if middleware.ExcludedKey != nil && middleware.ExcludedKey(key) {
+		c.Next()
+		return
+	}
+
 	context, err := middleware.Limiter.Get(c, key)
 	if err != nil {
 		middleware.OnError(c, err)

diff --git a/drivers/middleware/gin/middleware_test.go b/drivers/middleware/gin/middleware_test.go
@@ -125,4 +125,37 @@ func TestHTTPMiddleware(t *testing.T) {
 		is.Equal(http.StatusOK, resp.Code, strconv.FormatInt(i, 10))
 	}
 
+	//
+	// Test ExcludedKey
+	//
+	store = memory.NewStore()
+	is.NotZero(store)
+	counter = int64(0)
+	excludedKeyFn := func(key string) bool {
+		return key == "1"
+	}
+	middleware = gin.NewMiddleware(limiter.New(store, rate),
+		gin.WithKeyGetter(func(c *libgin.Context) string {
+			v := atomic.AddInt64(&counter, 1)
+			return strconv.FormatInt(v%2, 10)
+		}),
+		gin.WithExcludedKey(excludedKeyFn),
+	)
+	is.NotZero(middleware)
+
+	router = libgin.New()
+	router.Use(middleware)
+	router.GET("/", func(c *libgin.Context) {
+		c.String(http.StatusOK, "hello")
+	})
+	success = 20
+	for i := int64(1); i < clients; i++ {
+		resp := httptest.NewRecorder()
+		router.ServeHTTP(resp, request)
+		if i <= success || i%2 == 1 {
+			is.Equal(http.StatusOK, resp.Code, strconv.FormatInt(i, 10))
+		} else {
+			is.Equal(resp.Code, http.StatusTooManyRequests)
+		}
+	}
 }
diff --git a/drivers/middleware/gin/options.go b/drivers/middleware/gin/options.go
@@ -62,3 +62,10 @@ func WithKeyGetter(KeyGetter KeyGetter) Option {
 func DefaultKeyGetter(c *gin.Context) string {
 	return c.ClientIP()
 }
+
+// WithExcludedKey will configure the Middleware to use the given function.
+func WithExcludedKey(fn func(string) bool) Option {
+	return option(func(middleware *Middleware) {
+		middleware.ExcludedKey = fn
+	})
+}