/
limit.go
63 lines (52 loc) · 1.55 KB
/
limit.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
package limit
import (
"net/http"
"time"
"github.com/AlexanderChen1989/xrest"
"github.com/AlexanderChen1989/xrest/utils"
"github.com/didip/tollbooth"
"github.com/didip/tollbooth/config"
"github.com/didip/tollbooth/errors"
"golang.org/x/net/context"
)
// limiter cancel context when timeout
type limiter struct {
limiter *config.Limiter
ErrHandleFn func(http.ResponseWriter, *errors.HTTPError)
next xrest.Handler
}
func errHandleFn(w http.ResponseWriter, err *errors.HTTPError) {
utils.JSON.Code(err.StatusCode).Dump(w, map[string]interface{}{
"status": "error",
"msg": err.Message,
})
}
// New create a new request rate limiter plug, max requests in ttl time duration
func New(max int64, ttl time.Duration) xrest.Plugger {
return &limiter{
limiter: tollbooth.NewLimiter(max, ttl),
ErrHandleFn: errHandleFn,
}
}
// NewLimiter create a new request rate limiter with conf and error handle function
func NewLimiter(conf *config.Limiter, handlefn func(http.ResponseWriter, *errors.HTTPError)) xrest.Plugger {
return &limiter{
limiter: conf,
ErrHandleFn: handlefn,
}
}
// ServeHTTP implement Handler.ServeHTTP
func (limiter *limiter) ServeHTTP(ctx context.Context, w http.ResponseWriter, r *http.Request) {
err := tollbooth.LimitByRequest(limiter.limiter, r)
if err != nil {
limiter.ErrHandleFn(w, err)
return
}
// There's no rate-limit error, serve the next handler.
limiter.next.ServeHTTP(ctx, w, r)
}
// Plug implement Plugger.Plug
func (limiter *limiter) Plug(h xrest.Handler) xrest.Handler {
limiter.next = h
return limiter
}