github.com/ulule/limiter/v3@v3.11.3-0.20230613131926-4cb9c1da4633/drivers/middleware/gin/middleware.go (about)

     1  package gin
     2  
     3  import (
     4  	"strconv"
     5  
     6  	"github.com/gin-gonic/gin"
     7  
     8  	"github.com/ulule/limiter/v3"
     9  )
    10  
    11  // Middleware is the middleware for gin.
    12  type Middleware struct {
    13  	Limiter        *limiter.Limiter
    14  	OnError        ErrorHandler
    15  	OnLimitReached LimitReachedHandler
    16  	KeyGetter      KeyGetter
    17  	ExcludedKey    func(string) bool
    18  }
    19  
    20  // NewMiddleware return a new instance of a gin middleware.
    21  func NewMiddleware(limiter *limiter.Limiter, options ...Option) gin.HandlerFunc {
    22  	middleware := &Middleware{
    23  		Limiter:        limiter,
    24  		OnError:        DefaultErrorHandler,
    25  		OnLimitReached: DefaultLimitReachedHandler,
    26  		KeyGetter:      DefaultKeyGetter,
    27  		ExcludedKey:    nil,
    28  	}
    29  
    30  	for _, option := range options {
    31  		option.apply(middleware)
    32  	}
    33  
    34  	return func(ctx *gin.Context) {
    35  		middleware.Handle(ctx)
    36  	}
    37  }
    38  
    39  // Handle gin request.
    40  func (middleware *Middleware) Handle(c *gin.Context) {
    41  	key := middleware.KeyGetter(c)
    42  	if middleware.ExcludedKey != nil && middleware.ExcludedKey(key) {
    43  		c.Next()
    44  		return
    45  	}
    46  
    47  	context, err := middleware.Limiter.Get(c, key)
    48  	if err != nil {
    49  		middleware.OnError(c, err)
    50  		c.Abort()
    51  		return
    52  	}
    53  
    54  	c.Header("X-RateLimit-Limit", strconv.FormatInt(context.Limit, 10))
    55  	c.Header("X-RateLimit-Remaining", strconv.FormatInt(context.Remaining, 10))
    56  	c.Header("X-RateLimit-Reset", strconv.FormatInt(context.Reset, 10))
    57  
    58  	if context.Reached {
    59  		middleware.OnLimitReached(c)
    60  		c.Abort()
    61  		return
    62  	}
    63  
    64  	c.Next()
    65  }