Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

🤗 [Question]: limiter supports dynamic setting of “Max” parameter #3065

Closed
3 tasks done
zhangyongding opened this issue Jul 8, 2024 · 8 comments · Fixed by #3070
Closed
3 tasks done

🤗 [Question]: limiter supports dynamic setting of “Max” parameter #3065

zhangyongding opened this issue Jul 8, 2024 · 8 comments · Fixed by #3070

Comments

@zhangyongding
Copy link

Question Description

limiter only supports setting a fixed max value. Can a function be used to obtain a max value
example:

func(c *fiber.Ctx) int 
or
func(key string) int

Code Snippet (optional)

No response

Checklist:

  • I agree to follow Fiber's Code of Conduct.
  • I have checked for existing issues that describe my questions prior to opening this one.
  • I understand that improperly formatted questions may be closed without explanation.
Copy link

welcome bot commented Jul 8, 2024

Thanks for opening your first issue here! 🎉 Be sure to follow the issue template! If you need help or want to chat with us, join us on Discord https://gofiber.io/discord

@ReneWerner87
Copy link
Member

@zhangyongding can you please describe in more detail what you want to achieve and, if necessary, provide an example

@zhangyongding
Copy link
Author

add in middleware\limiter\config.go

type Config struct {
	...

	// MaxGenerator allows you to generate custom max, by default nil is used
	//
	// Default: nil
	MaxGenerator func(fiber.Ctx) int

	...
}

or

type Config struct {
	...

	// MaxGenerator allows you to generate custom max, by default nil is used
	//
	// Default: nil
	MaxGenerator func(key string) int
	
	...
}

example: middleware\limiter\limiter_fixed.go

// New creates a new fixed window middleware handler
func (FixedWindow) New(cfg Config) fiber.Handler {
	var (
		// Limiter variables
		mux        = &sync.RWMutex{}
		expiration = uint64(cfg.Expiration.Seconds())
	)

	// Create manager to simplify storage operations ( see manager.go )
	manager := newManager(cfg.Storage)

	// Update timestamp every second
	utils.StartTimeStampUpdater()

	// Return new handler
	return func(c fiber.Ctx) error {
		// Don't execute middleware if Next returns true
		if cfg.Next != nil && cfg.Next(c) {
			return c.Next()
		}

		// Get key from request
		key := cfg.KeyGenerator(c)

		// Get max from request or config
		max := cfg.Max
		if cfg.MaxGenerator != nil {
			max = cfg.MaxGenerator(c) // or max = cfg.MaxGenerator(key)
		}

		if max == 0 {
			return c.Next()
		}

		// Lock entry
		mux.Lock()

		// Get entry from pool and release when finished
		e := manager.get(key)

		// Get timestamp
		ts := uint64(utils.Timestamp())

		// Set expiration if entry does not exist
		if e.exp == 0 {
			e.exp = ts + expiration
		} else if ts >= e.exp {
			// Check if entry is expired
			e.currHits = 0
			e.exp = ts + expiration
		}

		// Increment hits
		e.currHits++

		// Calculate when it resets in seconds
		resetInSec := e.exp - ts

		// Set how many hits we have left
		remaining := max - e.currHits

		// Update storage
		manager.set(key, e, cfg.Expiration)

		// Unlock entry
		mux.Unlock()

		// Check if hits exceed the max
		if remaining < 0 {
			// Return response with Retry-After header
			// https://tools.ietf.org/html/rfc6584
			c.Set(fiber.HeaderRetryAfter, strconv.FormatUint(resetInSec, 10))

			// Call LimitReached handler
			return cfg.LimitReached(c)
		}

		// Continue stack for reaching c.Response().StatusCode()
		// Store err for returning
		err := c.Next()

		// Check for SkipFailedRequests and SkipSuccessfulRequests
		if (cfg.SkipSuccessfulRequests && c.Response().StatusCode() < fiber.StatusBadRequest) ||
			(cfg.SkipFailedRequests && c.Response().StatusCode() >= fiber.StatusBadRequest) {
			// Lock entry
			mux.Lock()
			e = manager.get(key)
			e.currHits--
			remaining++
			manager.set(key, e, cfg.Expiration)
			// Unlock entry
			mux.Unlock()
		}

		// We can continue, update RateLimit headers
		c.Set(xRateLimitLimit, strconv.Itoa(max))
		c.Set(xRateLimitRemaining, strconv.Itoa(remaining))
		c.Set(xRateLimitReset, strconv.FormatUint(resetInSec, 10))

		return err
	}
}

@ReneWerner87
Copy link
Member

I'm not sure what use case this is intended for
probably for the usecase where you want to manage different accounts with different rate limits

@zhangyongding
Copy link
Author

Yes, I want to manage different accounts with different rate limits

@zhangyongding
Copy link
Author

Does v2 accept PR?

@ReneWerner87
Copy link
Member

Does v2 accept PR?

no, only v3 accept new features

@luk3skyw4lker
Copy link
Contributor

@zhangyongding are you working on a PR? If not I could do it and send the PR today still

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging a pull request may close this issue.

3 participants