visitor.go 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. package server
  2. import (
  3. "golang.org/x/time/rate"
  4. "heckel.io/ntfy/util"
  5. "sync"
  6. "time"
  7. )
  8. const (
  9. visitorExpungeAfter = 30 * time.Minute
  10. )
  11. // visitor represents an API user, and its associated rate.Limiter used for rate limiting
  12. type visitor struct {
  13. config *Config
  14. limiter *rate.Limiter
  15. subscriptions *util.Limiter
  16. seen time.Time
  17. mu sync.Mutex
  18. }
  19. func newVisitor(conf *Config) *visitor {
  20. return &visitor{
  21. config: conf,
  22. limiter: rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst),
  23. subscriptions: util.NewLimiter(int64(conf.VisitorSubscriptionLimit)),
  24. seen: time.Now(),
  25. }
  26. }
  27. func (v *visitor) RequestAllowed() error {
  28. if !v.limiter.Allow() {
  29. return errHTTPTooManyRequests
  30. }
  31. return nil
  32. }
  33. func (v *visitor) AddSubscription() error {
  34. v.mu.Lock()
  35. defer v.mu.Unlock()
  36. if err := v.subscriptions.Add(1); err != nil {
  37. return errHTTPTooManyRequests
  38. }
  39. return nil
  40. }
  41. func (v *visitor) RemoveSubscription() {
  42. v.mu.Lock()
  43. defer v.mu.Unlock()
  44. v.subscriptions.Sub(1)
  45. }
  46. func (v *visitor) Keepalive() {
  47. v.mu.Lock()
  48. defer v.mu.Unlock()
  49. v.seen = time.Now()
  50. }
  51. func (v *visitor) Stale() bool {
  52. v.mu.Lock()
  53. defer v.mu.Unlock()
  54. return time.Since(v.seen) > visitorExpungeAfter
  55. }