Implement rate limiting (#26)

Author: man90es

README.md

@@ -38,6 +38,10 @@ Available flags:
 	- Type: string, space-separated list of IP addresses or URLs
 	- Default value: none, requests are made directly
 	- Also available as `PROXY` environment variable
+- `-ratelimit`
+	- Sets the maximum number of requests per minute per IP address
+	- Type: unsigned integer
+	- Default value: 512
 - `-verbose`
 	- Allows to put the app into verbose mode and print out additional logs to stdout
 	- Default value: none, no additional output is produced

config/config.go

@@ -10,12 +10,13 @@ import (
 )
 
 type config struct {
-	mu             sync.RWMutex
 	cacheTTL       time.Duration
 	maintenanceTTL time.Duration
+	mu             sync.RWMutex
 	port           int
 	proxyList      []string
 	proxySwitcher  colly.ProxyFunc
+	rateLimit      int64
 	verbosity      bool
 }
 
@@ -29,6 +30,7 @@ func getInstance() *config {
 			maintenanceTTL: 5 * time.Minute,
 			port:           8001,
 			proxyList:      nil,
+			rateLimit:      512,
 			verbosity:      false,
 		}
 	})
@@ -102,11 +104,24 @@ func GetVerbosity() bool {
 	return getInstance().verbosity
 }
 
+func SetRateLimit(rateLimit int64) {
+	getInstance().mu.Lock()
+	defer getInstance().mu.Unlock()
+	getInstance().rateLimit = rateLimit
+}
+
+func GetRateLimit() int64 {
+	getInstance().mu.RLock()
+	defer getInstance().mu.RUnlock()
+	return getInstance().rateLimit
+}
+
 func PrintConfig() {
 	fmt.Printf("Configuration:\n" +
 		fmt.Sprintf("\tPort:\t\t%v\n", GetPort()) +
 		fmt.Sprintf("\tProxies:\t%v\n", GetProxyList()) +
 		fmt.Sprintf("\tVerbosity:\t%v\n", GetVerbosity()) +
-		fmt.Sprintf("\tCache TTL:\t%v\n", GetCacheTTL()),
+		fmt.Sprintf("\tCache TTL:\t%v\n", GetCacheTTL()) +
+		fmt.Sprintf("\tRate limit:\t%v/min\n", GetRateLimit()),
 	)
 }

docs/openapi.json

@@ -103,6 +103,10 @@
 										"proxies": {
 											"type": "number"
 										},
+										"rateLimit": {
+											"type": "number",
+											"example": 512
+										},
 										"uptime": {
 											"type": "string",
 											"example": "2m49s"

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/gocolly/colly/v2 v2.1.0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/sa-/slicefunk v0.1.4
+	github.com/ulule/limiter/v3 v3.11.2
 	golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56
 )
 
@@ -19,6 +20,7 @@ require (
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/kennygrant/sanitize v1.2.4 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect
 	github.com/temoto/robotstxt v1.1.2 // indirect
 	golang.org/x/net v0.19.0 // indirect

go.sum

@@ -63,6 +63,8 @@ github.com/kennygrant/sanitize v1.2.4 h1:gN25/otpP5vAsO2djbMhF/LQX6R7+O1TB4yv8Nz
 github.com/kennygrant/sanitize v1.2.4/go.mod h1:LGsjYYtgxbetdg5owWB2mpgUL6e2nfw2eObZ0u0qvak=
 github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
 github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
@@ -73,11 +75,14 @@ github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7
 github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE=
-github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY=
+github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/temoto/robotstxt v1.1.1/go.mod h1:+1AmkuG3IYkh1kv0d2qEB9Le88ehNO0zwOr3ujewlOo=
 github.com/temoto/robotstxt v1.1.2 h1:W2pOjSJ6SWvldyEuiFXNxz3xZ8aiWX5LbfDiOFd7Fxg=
 github.com/temoto/robotstxt v1.1.2/go.mod h1:+1AmkuG3IYkh1kv0d2qEB9Le88ehNO0zwOr3ujewlOo=
+github.com/ulule/limiter/v3 v3.11.2 h1:P4yOrxoEMJbOTfRJR2OzjL90oflzYPPmWg+dvwN2tHA=
+github.com/ulule/limiter/v3 v3.11.2/go.mod h1:QG5GnFOCV+k7lrL5Y8kgEeeflPH3+Cviqlqa8SVSQxI=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
@@ -179,5 +184,7 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
 google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

handlers/ListenAndServe.go

@@ -22,7 +22,8 @@ func ListenAndServe() {
 	mux.HandleFunc("/", catchall)
 
 	middlewareStack := middleware.CreateStack(
-		middleware.SetHeaders,
+		middleware.GetSetHeadersMiddleware(),
+		middleware.GetRateLimitMiddleware(),
 	)
 
 	log.Println("Listening for requests")

handlers/getStatus.go

@@ -10,7 +10,7 @@ import (
 )
 
 var initTime = time.Now()
-var version = "1.9.3"
+var version = "1.9.4"
 
 func getStatus(w http.ResponseWriter, r *http.Request) {
 	json.NewEncoder(w).Encode(map[string]interface{}{
@@ -27,9 +27,10 @@ func getStatus(w http.ResponseWriter, r *http.Request) {
 				"maintenanceStatus": config.GetMaintenanceStatusTTL().Round(time.Minute).String(),
 			},
 		},
-		"docs":    docsLink,
-		"proxies": len(config.GetProxyList()),
-		"uptime":  time.Since(initTime).Round(time.Second).String(),
-		"version": version,
+		"docs":      docsLink,
+		"proxies":   len(config.GetProxyList()),
+		"rateLimit": config.GetRateLimit(),
+		"uptime":    time.Since(initTime).Round(time.Second).String(),
+		"version":   version,
 	})
 }

main.go

@@ -16,6 +16,7 @@ func main() {
 	flagMaintenanceTTL := flag.Int("maintenancettl", 5, "Allows to limit how frequently scraper can check for maintenance end in minutes")
 	flagPort := flag.Int("port", 8001, "Port to catch requests on")
 	flagProxy := flag.String("proxy", "", "Open proxy address to make requests to BDO servers")
+	flagRateLimit := flag.Int64("ratelimit", 512, "Maximum number of requests per minute per IP")
 	flagVerbose := flag.Bool("verbose", false, "Print out additional logs into stdout")
 	flag.Parse()
 
@@ -42,6 +43,7 @@ func main() {
 	config.SetCacheTTL(time.Duration(*flagCacheTTL) * time.Minute)
 	config.SetMaintenanceStatusTTL(time.Duration(*flagMaintenanceTTL) * time.Minute)
 	config.SetVerbosity(*flagVerbose)
+	config.SetRateLimit(*flagRateLimit)
 
 	config.PrintConfig()
 	handlers.ListenAndServe()

middleware/GetRateLimitMiddleware.go

@@ -0,0 +1,23 @@
+package middleware
+
+import (
+	"time"
+
+	"github.com/ulule/limiter/v3"
+	"github.com/ulule/limiter/v3/drivers/middleware/stdlib"
+	"github.com/ulule/limiter/v3/drivers/store/memory"
+
+	"bdo-rest-api/config"
+)
+
+func GetRateLimitMiddleware() Middleware {
+	var rate = limiter.Rate{
+		Limit:  config.GetRateLimit(),
+		Period: time.Minute,
+	}
+	var store = memory.NewStore()
+	var instance = limiter.New(store, rate, limiter.WithClientIPHeader("CF-Connecting-IP"))
+
+	var middleware = stdlib.NewMiddleware(instance)
+	return middleware.Handler
+}

middleware/GetSetHeadersMiddleware.go

@@ -0,0 +1,20 @@
+package middleware
+
+import (
+	"net/http"
+	"time"
+)
+
+func GetSetHeadersMiddleware() Middleware {
+	var setHeaders = func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Access-Control-Allow-Origin", "*")
+			w.Header().Set("Content-Type", "application/json")
+			w.Header().Set("Date", time.Now().Format(time.RFC1123Z))
+
+			next.ServeHTTP(w, r)
+		})
+	}
+
+	return setHeaders
+}