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

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