From 9ef47414f83a8fd60d0a0a4f56c74b52d2d6149d Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Mon, 4 Apr 2022 14:25:38 +0000
Subject: [PATCH] feat: Add rate-limits to the API

Closes #285.
---
 coderd/coderd.go                | 13 ++++++++++--
 coderd/httpmw/ratelimit.go      | 35 +++++++++++++++++++++++++++++++++
 coderd/httpmw/ratelimit_test.go | 32 ++++++++++++++++++++++++++++++
 go.mod                          |  2 ++
 go.sum                          |  2 ++
 5 files changed, 82 insertions(+), 2 deletions(-)
 create mode 100644 coderd/httpmw/ratelimit.go
 create mode 100644 coderd/httpmw/ratelimit_test.go

diff --git a/coderd/coderd.go b/coderd/coderd.go
index 0eac01793eb76..d60e975a726dd 100644
--- a/coderd/coderd.go
+++ b/coderd/coderd.go
@@ -47,14 +47,23 @@ func New(options *Options) (http.Handler, func()) {
 
 	r := chi.NewRouter()
 	r.Route("/api/v2", func(r chi.Router) {
-		r.Use(chitrace.Middleware())
+		r.Use(
+			chitrace.Middleware(),
+			// Specific routes can specify smaller limits.
+			httpmw.RateLimitPerMinute(512),
+		)
 		r.Get("/", func(w http.ResponseWriter, r *http.Request) {
 			httpapi.Write(w, http.StatusOK, httpapi.Response{
 				Message: "👋",
 			})
 		})
 		r.Route("/files", func(r chi.Router) {
-			r.Use(httpmw.ExtractAPIKey(options.Database, nil))
+			r.Use(
+				httpmw.ExtractAPIKey(options.Database, nil),
+				// This number is arbitrary, but reading/writing
+				// file content is expensive so it should be small.
+				httpmw.RateLimitPerMinute(12),
+			)
 			r.Get("/{hash}", api.fileByHash)
 			r.Post("/", api.postFile)
 		})
diff --git a/coderd/httpmw/ratelimit.go b/coderd/httpmw/ratelimit.go
new file mode 100644
index 0000000000000..293d71efae3e1
--- /dev/null
+++ b/coderd/httpmw/ratelimit.go
@@ -0,0 +1,35 @@
+package httpmw
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/go-chi/httprate"
+	"github.com/go-chi/render"
+
+	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/coderd/httpapi"
+)
+
+// RateLimitPerMinute returns a handler that limits requests per-minute based
+// on IP, endpoint, and user ID (if available).
+func RateLimitPerMinute(count int) func(http.Handler) http.Handler {
+	return httprate.Limit(
+		count,
+		1*time.Minute,
+		httprate.WithKeyFuncs(func(r *http.Request) (string, error) {
+			// Prioritize by user, but fallback to IP.
+			apiKey, ok := r.Context().Value(apiKeyContextKey{}).(database.APIKey)
+			if ok {
+				return apiKey.UserID.String(), nil
+			}
+			return httprate.KeyByIP(r)
+		}, httprate.KeyByEndpoint),
+		httprate.WithLimitHandler(func(w http.ResponseWriter, r *http.Request) {
+			render.Status(r, http.StatusTooManyRequests)
+			render.JSON(w, r, httpapi.Response{
+				Message: "You've been rate limited for sending too many requests!",
+			})
+		}),
+	)
+}
diff --git a/coderd/httpmw/ratelimit_test.go b/coderd/httpmw/ratelimit_test.go
new file mode 100644
index 0000000000000..c2e4ebd41fd78
--- /dev/null
+++ b/coderd/httpmw/ratelimit_test.go
@@ -0,0 +1,32 @@
+package httpmw_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/coderd/httpmw"
+)
+
+func TestRateLimit(t *testing.T) {
+	t.Parallel()
+	t.Run("NoUser", func(t *testing.T) {
+		t.Parallel()
+		rtr := chi.NewRouter()
+		rtr.Use(httpmw.RateLimitPerMinute(5))
+		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+			rw.WriteHeader(http.StatusOK)
+		})
+
+		require.Eventually(t, func() bool {
+			req := httptest.NewRequest("GET", "/", nil)
+			rec := httptest.NewRecorder()
+			rtr.ServeHTTP(rec, req)
+			return rec.Result().StatusCode == http.StatusTooManyRequests
+		}, 5*time.Second, time.Millisecond)
+	})
+}
diff --git a/go.mod b/go.mod
index 231f3612cac52..8972c6e241fe1 100644
--- a/go.mod
+++ b/go.mod
@@ -94,6 +94,8 @@ require (
 	storj.io/drpc v0.0.30
 )
 
+require github.com/go-chi/httprate v0.5.3
+
 require (
 	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
 	github.com/BurntSushi/toml v1.0.0 // indirect
diff --git a/go.sum b/go.sum
index ef7dc0b6fee6b..8f93749646d20 100644
--- a/go.sum
+++ b/go.sum
@@ -598,6 +598,8 @@ github.com/go-chi/chi/v4 v4.0.0-rc1/go.mod h1:Yfiy+5nynjDc7IMJiguACIro1KxlGW2dLU
 github.com/go-chi/chi/v5 v5.0.0/go.mod h1:BBug9lr0cqtdAhsu6R4AAdvufI0/XBzAQSsUqJpoZOs=
 github.com/go-chi/chi/v5 v5.0.7 h1:rDTPXLDHGATaeHvVlLcR4Qe0zftYethFucbjVQ1PxU8=
 github.com/go-chi/chi/v5 v5.0.7/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8=
+github.com/go-chi/httprate v0.5.3 h1:5HPWb0N6ymIiuotMtCfOGpQKiKeqXVzMexHh1W1yXPc=
+github.com/go-chi/httprate v0.5.3/go.mod h1:kYR4lorHX3It9tTh4eTdHhcF2bzrYnCrRNlv5+IBm2M=
 github.com/go-chi/render v1.0.1 h1:4/5tis2cKaNdnv9zFLfXzcquC9HbeZgCnxGnKrltBS8=
 github.com/go-chi/render v1.0.1/go.mod h1:pq4Rr7HbnsdaeHagklXub+p6Wd16Af5l9koip1OvJns=
 github.com/go-errors/errors v1.0.1/go.mod h1:f4zRHt4oKfwPJE5k8C9vpYG+aDHdBFUsgrm6/TyX73Q=