feat(ratelimit): Add http middleware handler function

internal/ratelimit/handler.go

@@ -0,0 +1,153 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package ratelimit
+
+import (
+	"fmt"
+	"math"
+	"net/http"
+	"regexp"
+
+	"github.com/hashicorp/boundary/internal/types/action"
+	"github.com/hashicorp/boundary/internal/types/resource"
+	"github.com/hashicorp/go-rate"
+)
+
+var (
+	errUnknownResource   = &extractResourceActionErr{http.StatusNotFound, "unknown resource"}
+	errUnknownAction     = &extractResourceActionErr{http.StatusBadRequest, "unknown action"}
+	errUnsupportedAction = &extractResourceActionErr{http.StatusMethodNotAllowed, "invalid action"}
+)
+
+type extractResourceActionErr struct {
+	statusCode int
+	msg        string
+}
+
+func (e *extractResourceActionErr) Error() string {
+	return e.msg
+}
+
+var pathRegex = regexp.MustCompile(`/v1/(?P<resource>[\w-]+)((/(?P<id>[\w/]+))?(:(?P<action>[\w-:]+)?)?)?`)
+
+func extractResourceAction(path, method string) (res, act string, err error) {
+	var id string
+
+	var r resource.Type
+	var ok bool
+	var actionSet action.ActionSet
+
+	// TODO: replace regex with lexer
+	match := pathRegex.FindStringSubmatch(path)
+	for i, name := range pathRegex.SubexpNames() {
+		switch name {
+		case "resource":
+			res = match[i]
+			r, ok = resource.FromPlural(res)
+			if !ok {
+				res = resource.Unknown.String()
+				err = errUnknownResource
+				return
+			}
+			res = r.String()
+		case "action":
+			act = match[i]
+			if act != "" {
+				actionSet, err = action.ActionSetForResource(r)
+				if err != nil {
+					act = action.Unknown.String()
+					err = errUnknownAction
+					return
+				}
+				at, ok := action.Map[act]
+				if !ok {
+					act = action.Unknown.String()
+					err = errUnknownAction
+					return
+				}
+				if !actionSet.HasAction(at) {
+					err = errUnsupportedAction
+					return
+				}
+			}
+		case "id":
+			id = match[i]
+		}
+	}
+
+	switch act {
+	case "":
+		switch id {
+		case "":
+			switch method {
+			case http.MethodGet:
+				act = action.List.String()
+			case http.MethodPost:
+				act = action.Create.String()
+			default:
+				act = action.Unknown.String()
+				err = errUnsupportedAction
+				return
+			}
+		default:
+			switch method {
+			case http.MethodGet:
+				act = action.Read.String()
+			case http.MethodDelete:
+				act = action.Delete.String()
+			case http.MethodPatch:
+				act = action.Update.String()
+			default:
+				act = action.Unknown.String()
+				err = errUnsupportedAction
+				return
+			}
+		}
+	}
+	return
+}
+
+// Handler is an http middleware handler that checks if a request is allowed
+// using the provided rate limiter. If the request is allowed, the next handler
+// is called. Otherwise a 429 is returned with the Retry-After response header
+// set to the number of seconds the client should wait to make it's next request.
+func Handler(l *rate.Limiter, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		path := req.URL.Path
+		method := req.Method
+
+		res, a, err := extractResourceAction(path, method)
+		if err != nil {
+			if extractErr, ok := err.(*extractResourceActionErr); ok {
+				rw.WriteHeader(extractErr.statusCode)
+				return
+			}
+			rw.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+
+		allowed, quota, err := l.Allow(res, a)
+		if err != nil {
+			if errFull, ok := err.(*rate.ErrLimiterFull); ok {
+				rw.Header().Add("Retry-After", fmt.Sprintf("%.0f", math.Ceil(errFull.RetryIn.Seconds())))
+				rw.WriteHeader(http.StatusServiceUnavailable)
+				return
+			}
+
+			// The only other error here should be rate.ErrLimitNotFound, which
+			// shouldn't be possible given how we initialize the limiter and
+			// the checks done by extractResourceAction.
+			rw.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+
+		if !allowed {
+			rw.Header().Add("Retry-After", fmt.Sprintf("%.0f", quota.ResetsIn().Seconds()))
+			rw.WriteHeader(http.StatusTooManyRequests)
+			return
+		}
+
+		next.ServeHTTP(rw, req)
+	})
+}

internal/ratelimit/handler_test.go

@@ -0,0 +1,405 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package ratelimit_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/boundary/internal/ratelimit"
+	"github.com/hashicorp/boundary/internal/types/action"
+	"github.com/hashicorp/boundary/internal/types/resource"
+	"github.com/hashicorp/go-rate"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHandler(t *testing.T) {
+	cases := []struct {
+		name          string
+		limiter       *rate.Limiter
+		req           func(uri string) *http.Request
+		expectCode    int
+		expectHeaders http.Header
+	}{
+		{
+			"AllowedList",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 10,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodGet, uri+"/v1/targets", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusOK,
+			http.Header{},
+		},
+		{
+			"AllowedCreate",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.Create.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 10,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodPost, uri+"/v1/targets", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusOK,
+			http.Header{},
+		},
+		{
+			"AllowedRead",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.Read.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 10,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodGet, uri+"/v1/targets/ttcp_123456789", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusOK,
+			http.Header{},
+		},
+		{
+			"AllowedUpdate",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.Update.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 10,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodPatch, uri+"/v1/targets/ttcp_123456789", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusOK,
+			http.Header{},
+		},
+		{
+			"AllowedDelete",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.Delete.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 10,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodDelete, uri+"/v1/targets/ttcp_123456789", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusOK,
+			http.Header{},
+		},
+		{
+			"AllowedAuthorizeSession",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.AuthorizeSession.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 10,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodDelete, uri+"/v1/targets/ttcp_123456789:authorize-session", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusOK,
+			http.Header{},
+		},
+		{
+			"Limited",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				_, q, err := r.Allow("target", "list")
+				require.NoError(t, err)
+
+				// consume all of the quota
+				remaining := int(q.Remaining())
+				for i := 0; i < remaining; i++ {
+					q.Consume()
+				}
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodGet, uri+"/v1/targets", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusTooManyRequests,
+			http.Header{
+				"Retry-After": []string{"60"},
+			},
+		},
+		{
+			"LimiterFull",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.Read.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 1)
+				require.NoError(t, err)
+				// use up all of the capacity for the limiter
+				_, _, err = r.Allow("target", "read")
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodGet, uri+"/v1/targets", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusServiceUnavailable,
+			http.Header{
+				"Retry-After": []string{"1"},
+			},
+		},
+		{
+			"MissingLimit",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodPost, uri+"/v1/targets", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusInternalServerError,
+			http.Header{},
+		},
+		{
+			"UnknownResource",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodPost, uri+"/v1/foo", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusNotFound,
+			http.Header{},
+		},
+		{
+			"UnknownAction",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodPut, uri+"/v1/target/ttcp_123456789:foo", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusBadRequest,
+			http.Header{},
+		},
+		{
+			"UnknownActionInvalidIdHttpMethod",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodPost, uri+"/v1/target/ttcp_123456789", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusMethodNotAllowed,
+			http.Header{},
+		},
+		{
+			"UnknownActionInvalidHttpMethod",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodDelete, uri+"/v1/target", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusMethodNotAllowed,
+			http.Header{},
+		},
+		{
+			"UnsupportedAction",
+			func() *rate.Limiter {
+				r, err := rate.NewLimiter([]*rate.Limit{
+					{
+						Resource:    resource.Target.String(),
+						Action:      action.List.String(),
+						Per:         rate.LimitPerTotal,
+						Unlimited:   false,
+						MaxRequests: 2,
+						Period:      time.Minute,
+					},
+				}, 10)
+				require.NoError(t, err)
+				return r
+			}(),
+			func(uri string) *http.Request {
+				r, err := http.NewRequest(http.MethodPut, uri+"/v1/host/hst_123456789:authorize-session", nil)
+				require.NoError(t, err)
+				return r
+			},
+			http.StatusMethodNotAllowed,
+			http.Header{},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			server := httptest.NewServer(ratelimit.Handler(tc.limiter, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusOK)
+				return
+			})))
+			req := tc.req(server.URL)
+			client := &http.Client{}
+			res, err := client.Do(req)
+			require.NoError(t, err)
+			assert.Equal(t, tc.expectCode, res.StatusCode)
+			for key, value := range tc.expectHeaders {
+				assert.Equal(t, value[0], res.Header.Get(key))
+			}
+		})
+	}
+}