github.com/gofiber/fiber/v2@v2.47.0/middleware/limiter/limiter_fixed.go (about)

     1  package limiter
     2  
     3  import (
     4  	"strconv"
     5  	"sync"
     6  	"sync/atomic"
     7  
     8  	"github.com/gofiber/fiber/v2"
     9  	"github.com/gofiber/fiber/v2/utils"
    10  )
    11  
    12  type FixedWindow struct{}
    13  
    14  // New creates a new fixed window middleware handler
    15  func (FixedWindow) New(cfg Config) fiber.Handler {
    16  	var (
    17  		// Limiter variables
    18  		mux        = &sync.RWMutex{}
    19  		max        = strconv.Itoa(cfg.Max)
    20  		expiration = uint64(cfg.Expiration.Seconds())
    21  	)
    22  
    23  	// Create manager to simplify storage operations ( see manager.go )
    24  	manager := newManager(cfg.Storage)
    25  
    26  	// Update timestamp every second
    27  	utils.StartTimeStampUpdater()
    28  
    29  	// Return new handler
    30  	return func(c *fiber.Ctx) error {
    31  		// Don't execute middleware if Next returns true
    32  		if cfg.Next != nil && cfg.Next(c) {
    33  			return c.Next()
    34  		}
    35  
    36  		// Get key from request
    37  		key := cfg.KeyGenerator(c)
    38  
    39  		// Lock entry
    40  		mux.Lock()
    41  
    42  		// Get entry from pool and release when finished
    43  		e := manager.get(key)
    44  
    45  		// Get timestamp
    46  		ts := uint64(atomic.LoadUint32(&utils.Timestamp))
    47  
    48  		// Set expiration if entry does not exist
    49  		if e.exp == 0 {
    50  			e.exp = ts + expiration
    51  		} else if ts >= e.exp {
    52  			// Check if entry is expired
    53  			e.currHits = 0
    54  			e.exp = ts + expiration
    55  		}
    56  
    57  		// Increment hits
    58  		e.currHits++
    59  
    60  		// Calculate when it resets in seconds
    61  		resetInSec := e.exp - ts
    62  
    63  		// Set how many hits we have left
    64  		remaining := cfg.Max - e.currHits
    65  
    66  		// Update storage
    67  		manager.set(key, e, cfg.Expiration)
    68  
    69  		// Unlock entry
    70  		mux.Unlock()
    71  
    72  		// Check if hits exceed the cfg.Max
    73  		if remaining < 0 {
    74  			// Return response with Retry-After header
    75  			// https://tools.ietf.org/html/rfc6584
    76  			c.Set(fiber.HeaderRetryAfter, strconv.FormatUint(resetInSec, 10))
    77  
    78  			// Call LimitReached handler
    79  			return cfg.LimitReached(c)
    80  		}
    81  
    82  		// Continue stack for reaching c.Response().StatusCode()
    83  		// Store err for returning
    84  		err := c.Next()
    85  
    86  		// Check for SkipFailedRequests and SkipSuccessfulRequests
    87  		if (cfg.SkipSuccessfulRequests && c.Response().StatusCode() < fiber.StatusBadRequest) ||
    88  			(cfg.SkipFailedRequests && c.Response().StatusCode() >= fiber.StatusBadRequest) {
    89  			// Lock entry
    90  			mux.Lock()
    91  			e = manager.get(key)
    92  			e.currHits--
    93  			remaining++
    94  			manager.set(key, e, cfg.Expiration)
    95  			// Unlock entry
    96  			mux.Unlock()
    97  		}
    98  
    99  		// We can continue, update RateLimit headers
   100  		c.Set(xRateLimitLimit, max)
   101  		c.Set(xRateLimitRemaining, strconv.Itoa(remaining))
   102  		c.Set(xRateLimitReset, strconv.FormatUint(resetInSec, 10))
   103  
   104  		return err
   105  	}
   106  }