idanoo/GoScrobble
1
0
Fork 0
mirror of https://github.com/idanoo/GoScrobble.git synced 2024-11-22 08:25:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Tidy up rate limiter

Browse Source
This commit is contained in:
Daniel Mason 2021-03-25 23:30:35 +13:00
parent 74108820f5
commit 12f4fb6d89
1 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
internal/goscrobble/server.go
View File

@ -24,7 +24,10 @@ type jsonResponse struct {
} }
// Limits to 1 req / 10 sec // Limits to 1 req / 10 sec
var limiter = NewIPRateLimiter(0.1, 1) var heavyLimiter = NewIPRateLimiter(0.1, 1)
// Limits to 5 req / sec
var standardLimiter = NewIPRateLimiter(5, 5)
// HandleRequests - Boot HTTP! // HandleRequests - Boot HTTP!
func HandleRequests() { func HandleRequests() {
@ -40,8 +43,8 @@ func HandleRequests() {
v1.HandleFunc("/profile/{id}", jwtMiddleware(serveEndpoint)) v1.HandleFunc("/profile/{id}", jwtMiddleware(serveEndpoint))
// No Auth // No Auth
v1.HandleFunc("/register", limitMiddleware(handleRegister)).Methods("POST") v1.HandleFunc("/register", limitMiddleware(handleRegister, heavyLimiter)).Methods("POST")
v1.HandleFunc("/login", serveEndpoint).Methods("POST") v1.HandleFunc("/login", limitMiddleware(serveEndpoint, standardLimiter)).Methods("POST")
v1.HandleFunc("/logout", serveEndpoint).Methods("POST") v1.HandleFunc("/logout", serveEndpoint).Methods("POST")
// This just prevents it serving frontend stuff over /api // This just prevents it serving frontend stuff over /api
@ -104,11 +107,13 @@ func jwtMiddleware(next http.HandlerFunc) http.HandlerFunc {
} }
// limitMiddleware - Rate limits important stuff // limitMiddleware - Rate limits important stuff
func limitMiddleware(next http.HandlerFunc) http.HandlerFunc { func limitMiddleware(next http.HandlerFunc, limiter *IPRateLimiter) http.HandlerFunc {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
limiter := limiter.GetLimiter(r.RemoteAddr) limiter := limiter.GetLimiter(r.RemoteAddr)
if !limiter.Allow() { if !limiter.Allow() {
http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests) msg := generateJsonMessage("Too many requests")
w.WriteHeader(http.StatusTooManyRequests)
w.Write(msg)
return return
} }