github.com/status-im/status-go@v1.1.0/mailserver/limiter.go (about)

     1  package mailserver
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  type rateLimiter struct {
     9  	sync.RWMutex
    10  
    11  	lifespan time.Duration // duration of the limit
    12  	db       map[string]time.Time
    13  
    14  	period time.Duration
    15  	cancel chan struct{}
    16  }
    17  
    18  func newRateLimiter(duration time.Duration) *rateLimiter {
    19  	return &rateLimiter{
    20  		lifespan: duration,
    21  		db:       make(map[string]time.Time),
    22  		period:   time.Second,
    23  	}
    24  }
    25  
    26  func (l *rateLimiter) Start() {
    27  	cancel := make(chan struct{})
    28  
    29  	l.Lock()
    30  	l.cancel = cancel
    31  	l.Unlock()
    32  
    33  	go l.cleanUp(l.period, cancel)
    34  }
    35  
    36  func (l *rateLimiter) Stop() {
    37  	l.Lock()
    38  	defer l.Unlock()
    39  
    40  	if l.cancel == nil {
    41  		return
    42  	}
    43  	close(l.cancel)
    44  	l.cancel = nil
    45  }
    46  
    47  func (l *rateLimiter) Add(id string) {
    48  	l.Lock()
    49  	l.db[id] = time.Now()
    50  	l.Unlock()
    51  }
    52  
    53  func (l *rateLimiter) IsAllowed(id string) bool {
    54  	l.RLock()
    55  	defer l.RUnlock()
    56  
    57  	if lastRequestTime, ok := l.db[id]; ok {
    58  		return lastRequestTime.Add(l.lifespan).Before(time.Now())
    59  	}
    60  
    61  	return true
    62  }
    63  
    64  func (l *rateLimiter) cleanUp(period time.Duration, cancel <-chan struct{}) {
    65  	t := time.NewTicker(period)
    66  	defer t.Stop()
    67  
    68  	for {
    69  		select {
    70  		case <-t.C:
    71  			l.deleteExpired()
    72  		case <-cancel:
    73  			return
    74  		}
    75  	}
    76  }
    77  
    78  func (l *rateLimiter) deleteExpired() {
    79  	l.Lock()
    80  	defer l.Unlock()
    81  
    82  	now := time.Now()
    83  	for id, lastRequestTime := range l.db {
    84  		if lastRequestTime.Add(l.lifespan).Before(now) {
    85  			delete(l.db, id)
    86  		}
    87  	}
    88  }