If the fleet/forgot_password endpoint is rate limited, it should return the proper status code (#12323)

Return proper HTTP status code if endpoint is rate limited.

changes/bug-10720-ratelimits-should-return-proper-status-code

@@ -0,0 +1,3 @@
+- If the `fleet/forgot_password` endpoint is rate limited it should return the proper HTTP status
+  code.
+- Fixed MaxBurst limit parameter for `fleet/forgot_password` endpoint.

server/service/handler.go

@@ -244,10 +244,11 @@ func addMetrics(r *mux.Router) {
 	r.Walk(walkFn) //nolint:errcheck
 }
 
-// desktopRateLimitMaxBurst is the max burst used for device request rate limiting.
-//
-// Defined as const to be used in tests.
-const desktopRateLimitMaxBurst = 100
+// These are defined as const so that they can be used in tests.
+const (
+	desktopRateLimitMaxBurst        = 100 // Max burst used for device request rate limiting.
+	forgotPasswordRateLimitMaxBurst = 9   // Max burst used for rate limiting on the the forgot_password endpoint.
+)
 
 func attachFleetAPIRoutes(r *mux.Router, svc fleet.Service, config config.FleetConfig,
 	logger kitlog.Logger, limitStore throttled.GCRAStore, opts []kithttp.ServerOption,
@@ -613,7 +614,7 @@ func attachFleetAPIRoutes(r *mux.Router, svc fleet.Service, config config.FleetC
 	// the handler.
 	ne.UsePathPrefix().PathHandler("GET", "/api/_version_/fleet/results/", makeStreamDistributedQueryCampaignResultsHandler(config.Server, svc, logger))
 
-	quota := throttled.RateQuota{MaxRate: throttled.PerHour(10), MaxBurst: 90}
+	quota := throttled.RateQuota{MaxRate: throttled.PerHour(10), MaxBurst: forgotPasswordRateLimitMaxBurst}
 	limiter := ratelimit.NewMiddleware(limitStore)
 	ne.
 		WithCustomMiddleware(limiter.Limit("forgot_password", quota)).

server/service/integration_desktop_test.go

@@ -225,12 +225,40 @@ func (s *integrationTestSuite) TestDefaultTransparencyURL() {
 	require.Equal(t, fleet.DefaultTransparencyURL, rawResp.Header.Get("Location"))
 }
 
-func (s *integrationTestSuite) TestDesktopRateLimit() {
+func (s *integrationTestSuite) TestRateLimitOfEndpoints() {
 	headers := map[string]string{
 		"X-Forwarded-For": "1.2.3.4",
 	}
-	for i := 0; i < desktopRateLimitMaxBurst+1; i++ { // rate limiting off-by-one
-		s.DoRawWithHeaders("GET", "/api/latest/fleet/device/"+uuid.NewString(), nil, http.StatusUnauthorized, headers).Body.Close()
+
+	testCases := []struct {
+		endpoint string
+		verb     string
+		payload  interface{}
+		burst    int
+		status   int
+	}{
+		{
+			endpoint: "/api/latest/fleet/forgot_password",
+			verb:     "POST",
+			payload:  forgotPasswordRequest{Email: "some@one.com"},
+			burst:    forgotPasswordRateLimitMaxBurst - 1,
+			status:   http.StatusAccepted,
+		},
+		{
+			endpoint: "/api/latest/fleet/device/" + uuid.NewString(),
+			verb:     "GET",
+			burst:    desktopRateLimitMaxBurst + 1,
+			status:   http.StatusUnauthorized,
+		},
+	}
+
+	for _, tCase := range testCases {
+		b, err := json.Marshal(tCase.payload)
+		require.NoError(s.T(), err)
+
+		for i := 0; i < tCase.burst; i++ {
+			s.DoRawWithHeaders(tCase.verb, tCase.endpoint, b, tCase.status, headers).Body.Close()
+		}
+		s.DoRawWithHeaders(tCase.verb, tCase.endpoint, b, http.StatusTooManyRequests, headers).Body.Close()
 	}
-	s.DoRawWithHeaders("GET", "/api/latest/fleet/device/"+uuid.NewString(), nil, http.StatusTooManyRequests, headers).Body.Close()
 }

server/service/middleware/ratelimit/ratelimit.go

@@ -40,7 +40,13 @@ func (m *Middleware) Limit(keyName string, quota throttled.RateQuota) endpoint.M
 			if err != nil {
 				return nil, ctxerr.Wrap(ctx, err, "check rate limit")
 			}
+
 			if limited {
+				// We need to set authentication as checked, otherwise we end up returning HTTP 500
+				// errors.
+				if az, ok := authz_ctx.FromContext(ctx); ok {
+					az.SetChecked()
+				}
 				return nil, ctxerr.Wrap(ctx, &ratelimitError{result: result})
 			}
 

server/service/middleware/ratelimit/ratelimit_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"testing"
 
+	authz_ctx "github.com/fleetdm/fleet/v4/server/contexts/authz"
 	"github.com/stretchr/testify/assert"
 	"github.com/throttled/throttled/v2"
 	"github.com/throttled/throttled/v2/store/memstore"
@@ -24,14 +25,19 @@ func TestLimit(t *testing.T) {
 		throttled.RateQuota{MaxRate: throttled.PerHour(1), MaxBurst: 0},
 	)(endpoint)
 
-	_, err := wrapped(context.Background(), struct{}{})
+	authzCtx := &authz_ctx.AuthorizationContext{}
+	ctx := authz_ctx.NewContext(context.Background(), authzCtx)
+
+	_, err := wrapped(ctx, struct{}{})
 	assert.NoError(t, err)
 
 	// Hits rate limit
-	_, err = wrapped(context.Background(), struct{}{})
+	_, err = wrapped(ctx, struct{}{})
 	assert.Error(t, err)
 	var rle Error
+
 	assert.True(t, errors.As(err, &rle))
+	assert.True(t, authzCtx.Checked())
 }
 
 func TestNewErrorMiddlewarePanics(t *testing.T) {