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

     1  package fasthttp
     2  
     3  import (
     4  	"github.com/valyala/fasthttp"
     5  )
     6  
     7  // Option is used to define Middleware configuration.
     8  type Option interface {
     9  	apply(middleware *Middleware)
    10  }
    11  
    12  type option func(*Middleware)
    13  
    14  func (o option) apply(middleware *Middleware) {
    15  	o(middleware)
    16  }
    17  
    18  // ErrorHandler is an handler used to inform when an error has occurred.
    19  type ErrorHandler func(ctx *fasthttp.RequestCtx, err error)
    20  
    21  // WithErrorHandler will configure the Middleware to use the given ErrorHandler.
    22  func WithErrorHandler(handler ErrorHandler) Option {
    23  	return option(func(middleware *Middleware) {
    24  		middleware.OnError = handler
    25  	})
    26  }
    27  
    28  // DefaultErrorHandler is the default ErrorHandler used by a new Middleware.
    29  func DefaultErrorHandler(ctx *fasthttp.RequestCtx, err error) {
    30  	panic(err)
    31  }
    32  
    33  // LimitReachedHandler is an handler used to inform when the limit has exceeded.
    34  type LimitReachedHandler func(ctx *fasthttp.RequestCtx)
    35  
    36  // WithLimitReachedHandler will configure the Middleware to use the given LimitReachedHandler.
    37  func WithLimitReachedHandler(handler LimitReachedHandler) Option {
    38  	return option(func(middleware *Middleware) {
    39  		middleware.OnLimitReached = handler
    40  	})
    41  }
    42  
    43  // DefaultLimitReachedHandler is the default LimitReachedHandler used by a new Middleware.
    44  func DefaultLimitReachedHandler(ctx *fasthttp.RequestCtx) {
    45  	ctx.SetStatusCode(fasthttp.StatusTooManyRequests)
    46  	ctx.Response.SetBodyString("Limit exceeded")
    47  }
    48  
    49  // KeyGetter will define the rate limiter key given the fasthttp Context.
    50  type KeyGetter func(ctx *fasthttp.RequestCtx) string
    51  
    52  // WithKeyGetter will configure the Middleware to use the given KeyGetter.
    53  func WithKeyGetter(KeyGetter KeyGetter) Option {
    54  	return option(func(middleware *Middleware) {
    55  		middleware.KeyGetter = KeyGetter
    56  	})
    57  }
    58  
    59  // DefaultKeyGetter is the default KeyGetter used by a new Middleware.
    60  // It returns the Client IP address.
    61  func DefaultKeyGetter(ctx *fasthttp.RequestCtx) string {
    62  	return ctx.RemoteIP().String()
    63  }
    64  
    65  // WithExcludedKey will configure the Middleware to ignore key(s) using the given function.
    66  func WithExcludedKey(handler func(string) bool) Option {
    67  	return option(func(middleware *Middleware) {
    68  		middleware.ExcludedKey = handler
    69  	})
    70  }