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
64
65
66
67
68
69
70
71
72
|
package utils
import (
"context"
"net/http"
"strconv"
"time"
"donetick.com/core/config"
"github.com/gin-gonic/gin"
"github.com/ulule/limiter/v3"
"github.com/ulule/limiter/v3/drivers/store/memory"
)
const (
XRequestIdKey = "X-Request-ID" // request id header key
)
func NewRateLimiter(cfg *config.Config) *limiter.Limiter {
store := memory.NewStore()
// rate, err := limiter.NewRateFromFormatted("10-H")
rate := limiter.Rate{
Period: cfg.Server.RatePeriod,
Limit: int64(cfg.Server.RateLimit),
}
// Then, create the limiter instance which takes the store and the rate as arguments.
// Now, you can give this instance to any supported middleware.
return limiter.New(store, rate)
}
// wrapper ratelimiter and have it as a middkewatr function:
func RateLimitMiddleware(limiter *limiter.Limiter) gin.HandlerFunc {
return func(c *gin.Context) {
// Use the IP as the key, which is the client IP.
// And set the expiration time to 10 seconds.
context, err := limiter.Get(c.Request.Context(), c.ClientIP())
if err != nil {
panic(err) // perhaps handle this nicer
}
// Check if the client is ratelimited.
if context.Reached {
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"message": "Too many requests"})
return
}
// Add a header in response to inform the current quota.
c.Header("X-RateLimit-Limit", strconv.FormatInt(context.Limit, 10))
// Add a header in response to inform the remaining quota.
c.Header("X-RateLimit-Remaining", strconv.FormatInt(context.Remaining, 10))
// Add a header in response to inform the time to wait before retry.
c.Header("X-RateLimit-Reset", strconv.FormatInt(context.Reset, 10))
c.Next()
}
}
func TimeoutMiddleware(timeout time.Duration) gin.HandlerFunc {
return func(c *gin.Context) {
ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
defer func() {
if ctx.Err() == context.DeadlineExceeded {
c.AbortWithStatus(http.StatusGatewayTimeout)
}
cancel()
}()
c.Request = c.Request.WithContext(ctx)
c.Next()
}
}
|
