45 lines
724 B
Go
45 lines
724 B
Go
|
package middlewares
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func RateLimiter(handler http.Handler, clearInterval time.Duration, maxRequests int) http.HandlerFunc {
|
||
|
|
requests := make(map[string]int)
|
||
|
|
lock := sync.Mutex{}
|
||
|
|
|
||
|
|
ticker := time.NewTicker(clearInterval)
|
||
|
|
go func() {
|
||
|
|
for {
|
||
|
|
<-ticker.C
|
||
|
|
|
||
|
|
lock.Lock()
|
||
|
|
clear(requests)
|
||
|
|
lock.Unlock()
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
|
||
|
|
return func(res http.ResponseWriter, req *http.Request) {
|
||
|
|
addr := req.RemoteAddr
|
||
|
|
|
||
|
|
lock.Lock()
|
||
|
|
count, found := requests[addr]
|
||
|
|
if !found {
|
||
|
|
count = 0
|
||
|
|
}
|
||
|
|
|
||
|
|
count += 1
|
||
|
|
requests[addr] = count
|
||
|
|
lock.Unlock()
|
||
|
|
|
||
|
|
if count > maxRequests {
|
||
|
|
http.Error(res, "Too many requests", http.StatusTooManyRequests)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
handler.ServeHTTP(res, req)
|
||
|
|
}
|
||
|
|
}
|