Function TestRateLimitByAuthToken

coderd/httpmw/ratelimit_test.go:151–249

(t *testing.T)

}

func TestRateLimitByAuthToken(t *testing.T) {
	t.Parallel()

	t.Run("LimitsByAuthHeader", func(t *testing.T) {
		t.Parallel()

		tests := []struct {
			name       string
			headerName string
			headerVal  string
		}{
			{
				name:       "BearerToken",
				headerName: "Authorization",
				headerVal:  "Bearer test-token-123",
			},
			{
				name:       "XApiKey",
				headerName: "X-Api-Key",
				headerVal:  "test-api-key-456",
			},
			{
				name:       "NoToken",
				headerName: "",
				headerVal:  "",
			},
		}

		for _, tt := range tests {
			t.Run(tt.name, func(t *testing.T) {
				t.Parallel()
				rtr := chi.NewRouter()
				rtr.Use(httpmw.RateLimitByAuthToken(2, time.Hour))
				rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
					rw.WriteHeader(http.StatusOK)
				})

				// Same token (or IP if no token) should be rate limited after 2 requests.
				for i := 0; i < 5; i++ {
					req := httptest.NewRequest("GET", "/", nil)
					if tt.headerName != "" {
						req.Header.Set(tt.headerName, tt.headerVal)
					}
					rec := httptest.NewRecorder()
					rtr.ServeHTTP(rec, req)
					resp := rec.Result()
					_ = resp.Body.Close()
					if i < 2 {
						require.Equal(t, http.StatusOK, resp.StatusCode, "request %d should succeed", i)
					} else {
						require.Equal(t, http.StatusTooManyRequests, resp.StatusCode, "request %d should be rate limited", i)
						// Verify Retry-After header is set.
						require.NotEmpty(t, resp.Header.Get("Retry-After"), "Retry-After header should be set")
					}
				}
			})
		}
	})