ip.go 902 B

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. package limiter
  2. import (
  3. "golang.org/x/time/rate"
  4. "sync"
  5. )
  6. var _ IPRateLimiter = (*ipRateLimiter)(nil)
  7. type IPRateLimiter interface {
  8. addIP(ip string) *rate.Limiter
  9. getLimiter(ip string) *rate.Limiter
  10. Allow(ip string) bool
  11. }
  12. type ipRateLimiter struct {
  13. ips *sync.Map
  14. limit rate.Limit
  15. burst int
  16. }
  17. func NewIPRateLimiter(limit int, burst int) IPRateLimiter {
  18. return &ipRateLimiter{
  19. ips: new(sync.Map),
  20. limit: rate.Limit(limit),
  21. burst: burst,
  22. }
  23. }
  24. func (i *ipRateLimiter) addIP(ip string) *rate.Limiter {
  25. store, _ := i.ips.LoadOrStore(ip, rate.NewLimiter(i.limit, i.burst))
  26. return store.(*rate.Limiter)
  27. }
  28. func (i *ipRateLimiter) getLimiter(ip string) *rate.Limiter {
  29. limiter, exists := i.ips.Load(ip)
  30. if !exists {
  31. return i.addIP(ip)
  32. }
  33. return limiter.(*rate.Limiter)
  34. }
  35. func (i *ipRateLimiter) Allow(ip string) bool {
  36. limiter := i.getLimiter(ip)
  37. return limiter.Allow()
  38. }