Added first acme implementation
This commit is contained in:
Pablu23
@@ -1,123 +1,122 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
type Limiter struct {
|
||||
currentBuckets map[string]*atomic.Int64
|
||||
bucketSize int
|
||||
refillTicker *time.Ticker
|
||||
cleanupTicker *time.Ticker
|
||||
bucketRefill int
|
||||
rwLock *sync.RWMutex
|
||||
rateChannel chan string
|
||||
}
|
||||
|
||||
func NewLimiter(maxRequests int, refills int, refillInterval time.Duration, cleanupInterval time.Duration) Limiter {
|
||||
return Limiter{
|
||||
currentBuckets: make(map[string]*atomic.Int64),
|
||||
bucketSize: maxRequests,
|
||||
refillTicker: time.NewTicker(refillInterval),
|
||||
cleanupTicker: time.NewTicker(cleanupInterval),
|
||||
bucketRefill: refills,
|
||||
rwLock: &sync.RWMutex{},
|
||||
rateChannel: make(chan string),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Limiter) Start() {
|
||||
go l.Manage()
|
||||
return
|
||||
}
|
||||
|
||||
func (l *Limiter) UpdateCleanupTime(new time.Duration) {
|
||||
l.cleanupTicker.Reset(new)
|
||||
}
|
||||
|
||||
func (l *Limiter) Manage() {
|
||||
for {
|
||||
select {
|
||||
case ip := <-l.rateChannel:
|
||||
if l.AddIfExists(ip) {
|
||||
break
|
||||
}
|
||||
|
||||
l.rwLock.Lock()
|
||||
counter := &atomic.Int64{}
|
||||
l.currentBuckets[ip] = counter
|
||||
l.rwLock.Unlock()
|
||||
case <-l.refillTicker.C:
|
||||
l.rwLock.RLock()
|
||||
for ip := range l.currentBuckets {
|
||||
n := l.currentBuckets[ip].Add(int64(-l.bucketRefill))
|
||||
if n < 0 {
|
||||
l.currentBuckets[ip].Store(0)
|
||||
n = 0
|
||||
}
|
||||
log.Trace().Int64("bucket", n).Str("remote", ip).Msg("Updated limit")
|
||||
}
|
||||
l.rwLock.RUnlock()
|
||||
log.Trace().Msg("Refreshed Limits")
|
||||
case <-l.cleanupTicker.C:
|
||||
start := time.Now()
|
||||
l.rwLock.Lock()
|
||||
deletedBuckets := 0
|
||||
for ip := range l.currentBuckets {
|
||||
if l.currentBuckets[ip].Load() <= 0 {
|
||||
delete(l.currentBuckets, ip)
|
||||
deletedBuckets += 1
|
||||
}
|
||||
}
|
||||
l.rwLock.Unlock()
|
||||
duration := time.Since(start)
|
||||
log.Debug().Str("duration", duration.String()).Int("deleted_buckets", deletedBuckets).Msg("Cleaned up Buckets")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Adds one if ip already exists and returns true
|
||||
// If ip doesnt yet exist only returns false
|
||||
func (l *Limiter) AddIfExists(ip string) bool {
|
||||
l.rwLock.RLock()
|
||||
defer l.rwLock.RUnlock()
|
||||
if counter, ok := l.currentBuckets[ip]; ok {
|
||||
counter.Add(1)
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (l *Limiter) RateLimiter(next http.Handler) http.Handler {
|
||||
log.Info().Int("bucket_size", l.bucketSize).Int("bucket_refill", l.bucketRefill).Msg("Enabling Ratelimits")
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
addr := strings.Split(r.RemoteAddr, ":")[0]
|
||||
l.rwLock.RLock()
|
||||
count, ok := l.currentBuckets[addr]
|
||||
l.rwLock.RUnlock()
|
||||
if ok && int(count.Load()) >= l.bucketSize {
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
r.Body.Close()
|
||||
log.Warn().Str("host", r.Host).Str("uri", r.RequestURI).Str("method", r.Method).Str("remote", addr).Msg("Rate limited")
|
||||
return
|
||||
}
|
||||
conn, _, err := hj.Hijack()
|
||||
if err != nil {
|
||||
log.Error().Err(err).Str("host", r.Host).Str("uri", r.RequestURI).Str("method", r.Method).Str("remote", addr).Msg("Could not hijack connection")
|
||||
}
|
||||
|
||||
log.Warn().Str("host", r.Host).Str("uri", r.RequestURI).Str("method", r.Method).Str("remote", addr).Msg("Rate limited")
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
l.rateChannel <- addr
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
type Limiter struct {
|
||||
currentBuckets map[string]*atomic.Int64
|
||||
bucketSize int
|
||||
refillTicker *time.Ticker
|
||||
cleanupTicker *time.Ticker
|
||||
bucketRefill int
|
||||
rwLock *sync.RWMutex
|
||||
rateChannel chan string
|
||||
}
|
||||
|
||||
func NewLimiter(maxRequests int, refills int, refillInterval time.Duration, cleanupInterval time.Duration) Limiter {
|
||||
return Limiter{
|
||||
currentBuckets: make(map[string]*atomic.Int64),
|
||||
bucketSize: maxRequests,
|
||||
refillTicker: time.NewTicker(refillInterval),
|
||||
cleanupTicker: time.NewTicker(cleanupInterval),
|
||||
bucketRefill: refills,
|
||||
rwLock: &sync.RWMutex{},
|
||||
rateChannel: make(chan string),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Limiter) Start() {
|
||||
go l.Manage()
|
||||
}
|
||||
|
||||
func (l *Limiter) UpdateCleanupTime(new time.Duration) {
|
||||
l.cleanupTicker.Reset(new)
|
||||
}
|
||||
|
||||
func (l *Limiter) Manage() {
|
||||
for {
|
||||
select {
|
||||
case ip := <-l.rateChannel:
|
||||
if l.AddIfExists(ip) {
|
||||
break
|
||||
}
|
||||
|
||||
l.rwLock.Lock()
|
||||
counter := &atomic.Int64{}
|
||||
l.currentBuckets[ip] = counter
|
||||
l.rwLock.Unlock()
|
||||
case <-l.refillTicker.C:
|
||||
l.rwLock.RLock()
|
||||
for ip := range l.currentBuckets {
|
||||
n := l.currentBuckets[ip].Add(int64(-l.bucketRefill))
|
||||
if n < 0 {
|
||||
l.currentBuckets[ip].Store(0)
|
||||
n = 0
|
||||
}
|
||||
log.Trace().Int64("bucket", n).Str("remote", ip).Msg("Updated limit")
|
||||
}
|
||||
l.rwLock.RUnlock()
|
||||
log.Trace().Msg("Refreshed Limits")
|
||||
case <-l.cleanupTicker.C:
|
||||
start := time.Now()
|
||||
l.rwLock.Lock()
|
||||
deletedBuckets := 0
|
||||
for ip := range l.currentBuckets {
|
||||
if l.currentBuckets[ip].Load() <= 0 {
|
||||
delete(l.currentBuckets, ip)
|
||||
deletedBuckets += 1
|
||||
}
|
||||
}
|
||||
l.rwLock.Unlock()
|
||||
duration := time.Since(start)
|
||||
log.Debug().Str("duration", duration.String()).Int("deleted_buckets", deletedBuckets).Msg("Cleaned up Buckets")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Adds one if ip already exists and returns true
|
||||
// If ip doesnt yet exist only returns false
|
||||
func (l *Limiter) AddIfExists(ip string) bool {
|
||||
l.rwLock.RLock()
|
||||
defer l.rwLock.RUnlock()
|
||||
if counter, ok := l.currentBuckets[ip]; ok {
|
||||
counter.Add(1)
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (l *Limiter) RateLimiter(next http.Handler) http.Handler {
|
||||
log.Info().Int("bucket_size", l.bucketSize).Int("bucket_refill", l.bucketRefill).Msg("Enabling Ratelimits")
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
addr := strings.Split(r.RemoteAddr, ":")[0]
|
||||
l.rwLock.RLock()
|
||||
count, ok := l.currentBuckets[addr]
|
||||
l.rwLock.RUnlock()
|
||||
if ok && int(count.Load()) >= l.bucketSize {
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
r.Body.Close()
|
||||
log.Warn().Str("host", r.Host).Str("uri", r.RequestURI).Str("method", r.Method).Str("remote", addr).Msg("Rate limited")
|
||||
return
|
||||
}
|
||||
conn, _, err := hj.Hijack()
|
||||
if err != nil {
|
||||
log.Error().Err(err).Str("host", r.Host).Str("uri", r.RequestURI).Str("method", r.Method).Str("remote", addr).Msg("Could not hijack connection")
|
||||
}
|
||||
|
||||
log.Warn().Str("host", r.Host).Str("uri", r.RequestURI).Str("method", r.Method).Str("remote", addr).Msg("Rate limited")
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
l.rateChannel <- addr
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user