78 lines
1.4 KiB
Go
78 lines
1.4 KiB
Go
|
package rateLimiter
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type RateLimitData struct {
|
||
|
|
Attempts int
|
||
|
|
LastAccess time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
// limit by IP
|
||
|
|
|
||
|
|
type IpRateLimiter struct {
|
||
|
|
Data map[string]*RateLimitData
|
||
|
|
Seconds int
|
||
|
|
AttemptsAllowed int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (self *IpRateLimiter) RateLimit(req *http.Request) bool {
|
||
|
|
curtime := time.Now()
|
||
|
|
|
||
|
|
addrWithPort := req.RemoteAddr
|
||
|
|
addrParts := strings.Split(addrWithPort, ":")
|
||
|
|
addr := strings.Join(addrParts[:len(addrParts)-1], ":")
|
||
|
|
|
||
|
|
data, exists := self.Data[addr]
|
||
|
|
if !exists {
|
||
|
|
self.Data[addr] = &RateLimitData{1, curtime}
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
if curtime.Sub(data.LastAccess).Seconds() > float64(self.Seconds) {
|
||
|
|
data.Attempts = 1
|
||
|
|
data.LastAccess = curtime
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
if self.AttemptsAllowed > data.Attempts {
|
||
|
|
data.Attempts++
|
||
|
|
data.LastAccess = curtime
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
data.Attempts++
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
// limit regardless of IP
|
||
|
|
|
||
|
|
type IndiscriminateRateLimiter struct {
|
||
|
|
Attempts int
|
||
|
|
AttemptsAllowed int
|
||
|
|
Seconds int
|
||
|
|
LastAccess time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
func (self *IndiscriminateRateLimiter) RateLimit(req *http.Request) bool {
|
||
|
|
curtime := time.Now()
|
||
|
|
|
||
|
|
if curtime.Sub(self.LastAccess).Seconds() > float64(self.Seconds) {
|
||
|
|
self.Attempts = 1
|
||
|
|
self.LastAccess = curtime
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
if self.AttemptsAllowed > self.Attempts {
|
||
|
|
self.Attempts++
|
||
|
|
self.LastAccess = curtime
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
self.Attempts++
|
||
|
|
return false
|
||
|
|
}
|