feat: add concurrent limiter and router concurrent limiter middleware

go.mod

@@ -4,6 +4,7 @@ go 1.12
 
 require (
 	github.com/stretchr/testify v1.5.1
+	github.com/tidwall/gjson v1.6.0
 	github.com/vicanso/fresh v0.1.1
 	github.com/vicanso/hes v0.2.1
 	github.com/vicanso/intranet-ip v0.0.1

go.sum

@@ -2,12 +2,17 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.5.1 h1:nOGnQDM7FYENwehXlg/kFVnos3rEvtKTjRvOWSzb6H4=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+github.com/tidwall/gjson v1.6.0 h1:9VEQWz6LLMUsUl6PueE49ir4Ka6CzLymOAZDxpFsTDc=
+github.com/tidwall/gjson v1.6.0/go.mod h1:P256ACg0Mn+j1RXIDXoss50DeIABTYK1PULOJHhxOls=
+github.com/tidwall/match v1.0.1 h1:PnKP62LPNxHKTwvHHZZzdOAOCtsJTjo6dZLCwpKm5xc=
+github.com/tidwall/match v1.0.1/go.mod h1:LujAq0jyVjBy028G1WhWfIzbpQfMO8bBZ6Tyb0+pL9E=
+github.com/tidwall/pretty v1.0.0 h1:HsD+QiTn7sK6flMKIvNmpqz1qrpP3Ps6jOKIKMooyg4=
+github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=
 github.com/vicanso/fresh v0.1.1 h1:L3UX/g2uZ4pNFzjc7+PyeYOlIAs2/yUfhcFnK7E+JQU=
 github.com/vicanso/fresh v0.1.1/go.mod h1:gr1RKSFxQ1OnQHzUMBHCigifni7KrXveJjWCTlPjICA=
 github.com/vicanso/hes v0.2.1 h1:jRFEADmiQ30koVY/sKwlkhyXM5B3QbVVizLqrjNJgPw=

middleware/concurrent_limiter.go

@@ -0,0 +1,168 @@
+// MIT License
+
+// Copyright (c) 2020 Tree Xie
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package middleware
+
+import (
+	"errors"
+	"net/http"
+	"strings"
+
+	"github.com/tidwall/gjson"
+	"github.com/vicanso/elton"
+	"github.com/vicanso/hes"
+)
+
+var (
+	// ErrSubmitTooFrequently submit too frequently
+	ErrSubmitTooFrequently = &hes.Error{
+		StatusCode: http.StatusBadRequest,
+		Message:    "submit too frequently",
+		Category:   ErrConcurrentLimiterCategory,
+	}
+	ErrRequireLockFunction = errors.New("require lock function")
+)
+
+const (
+	ipKey     = ":ip"
+	headerKey = "h:"
+	queryKey  = "q:"
+	paramKey  = "p:"
+	// ErrConcurrentLimiterCategory concurrent limiter error category
+	ErrConcurrentLimiterCategory = "elton-concurrent-limiter"
+)
+
+type (
+	// ConcurrentLimiterLock lock the key
+	ConcurrentLimiterLock func(string, *elton.Context) (bool, func(), error)
+	// Config concurrent limiter config
+	ConcurrentLimiterConfig struct {
+		// Keys keys for generate lock id
+		Keys []string
+		// Lock lock function
+		Lock    ConcurrentLimiterLock
+		Skipper elton.Skipper
+	}
+	// concurrentLimiterKeyInfo the concurrent key's info
+	concurrentLimiterKeyInfo struct {
+		Name   string
+		Params bool
+		Query  bool
+		Header bool
+		Body   bool
+		IP     bool
+	}
+)
+
+// New create a concurrent limiter middleware
+func NewConcurrentLimiter(config ConcurrentLimiterConfig) elton.Handler {
+
+	if config.Lock == nil {
+		panic(ErrRequireLockFunction)
+	}
+	keys := make([]*concurrentLimiterKeyInfo, 0)
+	// 根据配置生成key的处理
+	for _, key := range config.Keys {
+		if key == ipKey {
+			keys = append(keys, &concurrentLimiterKeyInfo{
+				IP: true,
+			})
+			continue
+		}
+		if strings.HasPrefix(key, headerKey) {
+			keys = append(keys, &concurrentLimiterKeyInfo{
+				Name:   key[2:],
+				Header: true,
+			})
+			continue
+		}
+		if strings.HasPrefix(key, queryKey) {
+			keys = append(keys, &concurrentLimiterKeyInfo{
+				Name:  key[2:],
+				Query: true,
+			})
+			continue
+		}
+		if strings.HasPrefix(key, paramKey) {
+			keys = append(keys, &concurrentLimiterKeyInfo{
+				Name:   key[2:],
+				Params: true,
+			})
+			continue
+		}
+		keys = append(keys, &concurrentLimiterKeyInfo{
+			Name: key,
+			Body: true,
+		})
+	}
+	skipper := config.Skipper
+	if skipper == nil {
+		skipper = elton.DefaultSkipper
+	}
+	keyLength := len(keys)
+	return func(c *elton.Context) (err error) {
+		if skipper(c) {
+			return c.Next()
+		}
+		sb := new(strings.Builder)
+		// 先申请假定每个value的长度
+		sb.Grow(8 * keyLength)
+		// 获取 lock 的key
+		for i, key := range keys {
+			v := ""
+			name := key.Name
+			if key.IP {
+				v = c.RealIP()
+			} else if key.Header {
+				v = c.GetRequestHeader(name)
+			} else if key.Query {
+				query := c.Query()
+				v = query[name]
+			} else if key.Params {
+				v = c.Param(name)
+			} else {
+				v = gjson.GetBytes(c.RequestBody, name).String()
+			}
+			sb.WriteString(v)
+			if i < keyLength-1 {
+				sb.WriteRune(',')
+			}
+		}
+		lockKey := sb.String()
+
+		success, unlock, err := config.Lock(lockKey, c)
+		if err != nil {
+			err = hes.Wrap(err)
+			return
+		}
+		if !success {
+			err = ErrSubmitTooFrequently
+			return
+		}
+
+		if unlock != nil {
+			defer unlock()
+		}
+
+		return c.Next()
+	}
+}

middleware/concurrent_limiter_test.go

@@ -0,0 +1,122 @@
+// MIT License
+
+// Copyright (c) 2020 Tree Xie
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package middleware
+
+import (
+	"errors"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/vicanso/elton"
+)
+
+func TestNoLockFunction(t *testing.T) {
+	assert := assert.New(t)
+	defer func() {
+		r := recover()
+		assert.Equal(r.(error), ErrRequireLockFunction)
+	}()
+
+	NewConcurrentLimiter(ConcurrentLimiterConfig{})
+}
+
+func TestConcurrentLimiter(t *testing.T) {
+	m := new(sync.Map)
+	fn := NewConcurrentLimiter(ConcurrentLimiterConfig{
+		Keys: []string{
+			":ip",
+			"h:X-Token",
+			"q:type",
+			"p:id",
+			"account",
+		},
+		Lock: func(key string, c *elton.Context) (success bool, unlock func(), err error) {
+			if key != "192.0.2.1,xyz,1,123,tree.xie" {
+				err = errors.New("key is invalid")
+				return
+			}
+			_, loaded := m.LoadOrStore(key, 1)
+			// 如果已存在，则获取销失败
+			if loaded {
+				return
+			}
+			success = true
+			// 删除锁
+			unlock = func() {
+				m.Delete(key)
+			}
+			return
+		},
+	})
+
+	req := httptest.NewRequest("POST", "/users/login?type=1", nil)
+	resp := httptest.NewRecorder()
+	c := elton.NewContext(resp, req)
+	req.Header.Set("X-Token", "xyz")
+	c.RequestBody = []byte(`{
+		"account": "tree.xie"
+	}`)
+	c.Params = new(elton.RouteParams)
+	c.Params.Add("id", "123")
+
+	t.Run("first", func(t *testing.T) {
+		assert := assert.New(t)
+		done := false
+		c.Next = func() error {
+			done = true
+			return nil
+		}
+		err := fn(c)
+		assert.Nil(err)
+		assert.True(done)
+	})
+
+	t.Run("too frequently", func(t *testing.T) {
+		assert := assert.New(t)
+		done := false
+		c.Next = func() error {
+			time.Sleep(100 * time.Millisecond)
+			done = true
+			return nil
+		}
+		go func() {
+			time.Sleep(10 * time.Millisecond)
+			e := fn(c)
+			assert.Equal(e.Error(), "category=elton-concurrent-limiter, message=submit too frequently")
+		}()
+		err := fn(c)
+		// 登录限制,192.0.2.1,xyz,1,123,tree.xie
+		assert.Nil(err)
+		assert.True(done)
+	})
+
+	t.Run("lock function return error", func(t *testing.T) {
+		assert := assert.New(t)
+		c.Params = new(elton.RouteParams)
+		err := fn(c)
+		assert.Equal(err.Error(), "message=key is invalid")
+	})
+}