quartzgun/rateLimiter/rateLimiter.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
}