diff --git a/cmd/api/errors.go b/cmd/api/errors.go index 9a87e52..d5c4892 100644 --- a/cmd/api/errors.go +++ b/cmd/api/errors.go @@ -51,3 +51,8 @@ func (app *application) editConflictResponse(w http.ResponseWriter, r *http.Requ message := "unable to update the record due to an edit conflict, please try again" app.errorResponse(w, r, http.StatusConflict, message) } + +func (app *application) rateLimitExceededResponse(w http.ResponseWriter, r *http.Request) { + message := "rate limit exceeded" + app.errorResponse(w, r, http.StatusTooManyRequests, message) +} diff --git a/cmd/api/main.go b/cmd/api/main.go index a6a3613..bccab22 100644 --- a/cmd/api/main.go +++ b/cmd/api/main.go @@ -32,6 +32,11 @@ type config struct { maxIdleConns int maxIdleTime string } + limiter struct { + rps float64 + burst int + enabled bool + } } func openDB(cfg config) (*sql.DB, error) { @@ -69,6 +74,11 @@ func main() { flag.IntVar(&cfg.db.maxOpenConns, "db-max-open-conns", 25, "PostgreSQL max open connections") flag.IntVar(&cfg.db.maxIdleConns, "db-max-idle-conns", 25, "PostgreSQL max idle connections") flag.StringVar(&cfg.db.maxIdleTime, "db-max-idle-time", "15m", "PostgreSQL max connection idle time") + + flag.Float64Var(&cfg.limiter.rps, "limiter-rps", 2, "Rate limiter maximum requests per second") + flag.IntVar(&cfg.limiter.burst, "limiter-burst", 4, "Rate limiter maximum burst") + flag.BoolVar(&cfg.limiter.enabled, "limiter-enabled", true, "Enable rate limiter") + flag.Parse() logger := jsonlog.New(os.Stdout, jsonlog.LevelInfo) diff --git a/cmd/api/middleware.go b/cmd/api/middleware.go index 5052f87..a959ebb 100644 --- a/cmd/api/middleware.go +++ b/cmd/api/middleware.go @@ -2,7 +2,12 @@ package main import ( "fmt" + "net" "net/http" + "sync" + "time" + + "golang.org/x/time/rate" ) func (app *application) recoverPanic(next http.Handler) http.Handler { @@ -16,3 +21,57 @@ func (app *application) recoverPanic(next http.Handler) http.Handler { next.ServeHTTP(w, r) }) } + +func (app *application) rateLimit(next http.Handler) http.Handler { + type client struct { + limiter *rate.Limiter + lastSeen time.Time + } + var ( + mu sync.Mutex + clients = make(map[string]*client) + ) + + go func() { + for { + time.Sleep(time.Minute) + + mu.Lock() + + for ip, client := range clients { + if time.Since(client.lastSeen) > 3*time.Minute { + delete(clients, ip) + } + } + + mu.Unlock() + } + }() + + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if app.config.limiter.enabled { + ip, _, err := net.SplitHostPort(r.RemoteAddr) + if err != nil { + app.serverErrorResponse(w, r, err) + return + } + + mu.Lock() + + if _, found := clients[ip]; !found { + clients[ip] = &client{limiter: rate.NewLimiter(rate.Limit(app.config.limiter.rps), app.config.limiter.burst)} + } + clients[ip].lastSeen = time.Now() + + if !clients[ip].limiter.Allow() { + mu.Unlock() + app.rateLimitExceededResponse(w, r) + return + } + + mu.Unlock() + } + + next.ServeHTTP(w, r) + }) +} diff --git a/cmd/api/routes.go b/cmd/api/routes.go index 5f79f26..28b6270 100644 --- a/cmd/api/routes.go +++ b/cmd/api/routes.go @@ -20,5 +20,5 @@ func (app *application) routes() http.Handler { router.HandlerFunc(http.MethodPatch, "/v1/movies/:id", app.updateMovieHandler) router.HandlerFunc(http.MethodDelete, "/v1/movies/:id", app.deleteMovieHandler) - return app.recoverPanic(router) + return app.recoverPanic(app.rateLimit(router)) } diff --git a/go.mod b/go.mod index 683af40..699c043 100644 --- a/go.mod +++ b/go.mod @@ -5,4 +5,5 @@ go 1.25.0 require ( github.com/julienschmidt/httprouter v1.3.0 // indirect github.com/lib/pq v1.11.2 // indirect + golang.org/x/time v0.15.0 // indirect ) diff --git a/go.sum b/go.sum index 7fd2db1..4d218cc 100644 --- a/go.sum +++ b/go.sum @@ -2,3 +2,5 @@ github.com/julienschmidt/httprouter v1.3.0 h1:U0609e9tgbseu3rBINet9P48AI/D3oJs4d github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8IZAc4RVcycCCAKdM= github.com/lib/pq v1.11.2 h1:x6gxUeu39V0BHZiugWe8LXZYZ+Utk7hSJGThs8sdzfs= github.com/lib/pq v1.11.2/go.mod h1:/p+8NSbOcwzAEI7wiMXFlgydTwcgTr3OSKMsD2BitpA= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=