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

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