chore: add global caching to rbac (#7439)

Co-authored-by: Steven Masley <stevenmasley@coder.com>

.gitignore

@@ -53,3 +53,4 @@ site/stats/
 
 # direnv
 .envrc
+*.test

.prettierignore

@@ -56,6 +56,7 @@ site/stats/
 
 # direnv
 .envrc
+*.test
 # .prettierignore.include:
 # Helm templates contain variables that are invalid YAML and can't be formatted
 # by Prettier.

coderd/coderd.go

@@ -398,7 +398,6 @@ func New(options *Options) *API {
 		tracing.StatusWriterMiddleware,
 		tracing.Middleware(api.TracerProvider),
 		httpmw.AttachRequestID,
-		httpmw.AttachAuthzCache,
 		httpmw.ExtractRealIP(api.RealIPConfig),
 		httpmw.Logger(api.Logger),
 		httpmw.Prometheus(options.PrometheusRegistry),

coderd/coderdtest/coderdtest.go

@@ -137,7 +137,7 @@ type Options struct {
 }
 
 // New constructs a codersdk client connected to an in-memory API instance.
-func New(t *testing.T, options *Options) *codersdk.Client {
+func New(t testing.TB, options *Options) *codersdk.Client {
 	client, _ := newWithCloser(t, options)
 	return client
 }
@@ -162,12 +162,12 @@ func NewWithProvisionerCloser(t *testing.T, options *Options) (*codersdk.Client,
 // upon thee. Even the io.Closer that is exposed here shouldn't be exposed
 // and is a temporary measure while the API to register provisioners is ironed
 // out.
-func newWithCloser(t *testing.T, options *Options) (*codersdk.Client, io.Closer) {
+func newWithCloser(t testing.TB, options *Options) (*codersdk.Client, io.Closer) {
 	client, closer, _ := NewWithAPI(t, options)
 	return client, closer
 }
 
-func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.CancelFunc, *url.URL, *coderd.Options) {
+func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.CancelFunc, *url.URL, *coderd.Options) {
 	if options == nil {
 		options = &Options{}
 	}
@@ -190,8 +190,14 @@ func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.Can
 	}
 
 	if options.Authorizer == nil {
-		options.Authorizer = &RecordingAuthorizer{
+		defAuth := rbac.NewCachingAuthorizer(prometheus.NewRegistry())
-			Wrapped: rbac.NewCachingAuthorizer(prometheus.NewRegistry()),
+		if _, ok := t.(*testing.T); ok {
+			options.Authorizer = &RecordingAuthorizer{
+				Wrapped: defAuth,
+			}
+		} else {
+			// In benchmarks, the recording authorizer greatly skews results.
+			options.Authorizer = defAuth
 		}
 	}
 
@@ -359,7 +365,7 @@ func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.Can
 // NewWithAPI constructs an in-memory API instance and returns a client to talk to it.
 // Most tests never need a reference to the API, but AuthorizationTest in this module uses it.
 // Do not expose the API or wrath shall descend upon thee.
-func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *coderd.API) {
+func NewWithAPI(t testing.TB, options *Options) (*codersdk.Client, io.Closer, *coderd.API) {
 	if options == nil {
 		options = &Options{}
 	}
@@ -384,7 +390,7 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
 // NewProvisionerDaemon launches a provisionerd instance configured to work
 // well with coderd testing. It registers the "echo" provisioner for
 // quick testing.
-func NewProvisionerDaemon(t *testing.T, coderAPI *coderd.API) io.Closer {
+func NewProvisionerDaemon(t testing.TB, coderAPI *coderd.API) io.Closer {
 	echoClient, echoServer := provisionersdk.MemTransportPipe()
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	t.Cleanup(func() {
@@ -465,7 +471,7 @@ var FirstUserParams = codersdk.CreateFirstUserRequest{
 
 // CreateFirstUser creates a user with preset credentials and authenticates
 // with the passed in codersdk client.
-func CreateFirstUser(t *testing.T, client *codersdk.Client) codersdk.CreateFirstUserResponse {
+func CreateFirstUser(t testing.TB, client *codersdk.Client) codersdk.CreateFirstUserResponse {
 	resp, err := client.CreateFirstUser(context.Background(), FirstUserParams)
 	require.NoError(t, err)
 
@@ -1111,7 +1117,7 @@ sz9Di8sGIaUbLZI2rd0CQQCzlVwEtRtoNCyMJTTrkgUuNufLP19RZ5FpyXxBO5/u
 QastnN77KfUwdj3SJt44U/uh1jAIv4oSLBr8HYUkbnI8
 -----END RSA PRIVATE KEY-----`
 
-func DeploymentValues(t *testing.T) *codersdk.DeploymentValues {
+func DeploymentValues(t testing.TB) *codersdk.DeploymentValues {
 	var cfg codersdk.DeploymentValues
 	opts := cfg.Options()
 	err := opts.SetDefaults()

coderd/database/dbtestutil/db.go

@@ -13,7 +13,7 @@ import (
 	"github.com/coder/coder/coderd/database/postgres"
 )
 
-func NewDB(t *testing.T) (database.Store, database.Pubsub) {
+func NewDB(t testing.TB) (database.Store, database.Pubsub) {
 	t.Helper()
 
 	db := dbfake.New()

coderd/httpmw/authz.go

@@ -6,7 +6,6 @@ import (
 	"github.com/go-chi/chi/v5"
 
 	"github.com/coder/coder/coderd/database/dbauthz"
-	"github.com/coder/coder/coderd/rbac"
 )
 
 // AsAuthzSystem is a chained handler that temporarily sets the dbauthz context
@@ -36,16 +35,3 @@ func AsAuthzSystem(mws ...func(http.Handler) http.Handler) func(http.Handler) ht
 		})
 	}
 }
-
-// AttachAuthzCache enables the authz cache for the authorizer. All rbac checks will
-// run against the cache, meaning duplicate checks will not be performed.
-//
-// Note the cache is safe for multiple actors. So mixing user and system checks
-// is ok.
-func AttachAuthzCache(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
-		ctx := rbac.WithCacheCtx(r.Context())
-
-		next.ServeHTTP(rw, r.WithContext(ctx))
-	})
-}

coderd/rbac/authz.go

@@ -2,11 +2,14 @@ package rbac
 
 import (
 	"context"
+	"crypto/sha256"
 	_ "embed"
+	"encoding/json"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/ammario/tlru"
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/prometheus/client_golang/prometheus"
@@ -42,6 +45,31 @@ type AuthCall struct {
 	Object Object
 }
 
+// hashAuthorizeCall guarantees a unique hash for a given auth call.
+// If two hashes are equal, then the result of a given authorize() call
+// will be the same.
+//
+// Note that this ignores some fields such as the permissions within a given
+// role, as this assumes all roles are static to a given role name.
+func hashAuthorizeCall(actor Subject, action Action, object Object) [32]byte {
+	var hashOut [32]byte
+	hash := sha256.New()
+
+	// We use JSON for the forward security benefits if the rbac structs are
+	// modified without consideration for the caching layer.
+	enc := json.NewEncoder(hash)
+	_ = enc.Encode(actor)
+	_ = enc.Encode(action)
+	_ = enc.Encode(object)
+
+	// We might be able to avoid this extra copy?
+	// sha256.Sum256() returns a [32]byte. We need to return
+	// an array vs a slice so we can use it as a key in the cache.
+	image := hash.Sum(nil)
+	copy(hashOut[:], image)
+	return hashOut
+}
+
 // Subject is a struct that contains all the elements of a subject in an rbac
 // authorize.
 type Subject struct {
@@ -101,6 +129,9 @@ func (s Subject) SafeRoleNames() []string {
 }
 
 type Authorizer interface {
+	// Authorize will authorize the given subject to perform the given action
+	// on the given object. Authorize is pure and deterministic with respect to
+	// its arguments and the surrounding object.
 	Authorize(ctx context.Context, subject Subject, action Action, object Object) error
 	Prepare(ctx context.Context, subject Subject, action Action, objectType string) (PreparedAuthorized, error)
 }
@@ -310,6 +341,7 @@ func (a RegoAuthorizer) Authorize(ctx context.Context, subject Subject, action A
 	defer span.End()
 
 	err := a.authorize(ctx, subject, action, object)
+
 	span.SetAttributes(attribute.Bool("authorized", err == nil))
 
 	dur := time.Since(start)
@@ -605,7 +637,12 @@ func (a *authorizedSQLFilter) SQLString() string {
 	return a.sqlString
 }
 
-type cachedCalls struct {
+type authCache struct {
+	// cache is a cache of hashed Authorize inputs to the result of the Authorize
+	// call.
+	// determistic function.
+	cache *tlru.Cache[[32]byte, error]
+
 	authz Authorizer
 }
 
@@ -617,94 +654,35 @@ type cachedCalls struct {
 //
 // Cacher is safe for multiple actors.
 func Cacher(authz Authorizer) Authorizer {
-	return &cachedCalls{authz: authz}
+	return &authCache{
+		authz: authz,
+		// In practice, this cache should never come close to filling since the
+		// authorization calls are kept for a minute at most.
+		cache: tlru.New[[32]byte](tlru.ConstantCost[error], 64*1024),
+	}
 }
 
-func (c *cachedCalls) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
+func (c *authCache) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
-	cache := cacheFromContext(ctx)
+	authorizeCacheKey := hashAuthorizeCall(subject, action, object)
 
-	resp, ok := cache.Load(subject, action, object)
+	var err error
-	if ok {
+	err, _, ok := c.cache.Get(authorizeCacheKey)
-		return resp
+	if !ok {
+		err = c.authz.Authorize(ctx, subject, action, object)
+		// In case there is a caching bug, bound the TTL to 1 minute.
+		c.cache.Set(authorizeCacheKey, err, time.Minute)
 	}
 
-	err := c.authz.Authorize(ctx, subject, action, object)
-	cache.Save(subject, action, object, err)
 	return err
 }
 
 // Prepare returns the underlying PreparedAuthorized. The cache does not apply
 // to prepared authorizations. These should be using a SQL filter, and
 // therefore the cache is not needed.
-func (c *cachedCalls) Prepare(ctx context.Context, subject Subject, action Action, objectType string) (PreparedAuthorized, error) {
+func (c *authCache) Prepare(ctx context.Context, subject Subject, action Action, objectType string) (PreparedAuthorized, error) {
 	return c.authz.Prepare(ctx, subject, action, objectType)
 }
 
-// authorizeCache enabled caching of Authorizer calls for a given request. This
-// prevents the cost of running the same rbac checks multiple times.
-// A cache hit must match on all 3 values: subject, action, and object.
-type authorizeCache struct {
-	sync.Mutex
-	// calls is a list of all calls made to the Authorizer.
-	// This list is cached per request context. The size of this list is expected
-	// to be incredibly small. Often 1 or 2 calls.
-	calls []cachedAuthCall
-}
-
-type cachedAuthCall struct {
-	AuthCall
-	Err error
-}
-
-// cacheContextKey is a context key used to store the cache in the context.
-type cacheContextKey struct{}
-
-// cacheFromContext returns the cache from the context.
-// If there is no cache, a nil value is returned.
-// The nil cache can still be called as a normal cache, but will not cache or
-// return any values.
-func cacheFromContext(ctx context.Context) *authorizeCache {
-	cache, _ := ctx.Value(cacheContextKey{}).(*authorizeCache)
-	return cache
-}
-
-func WithCacheCtx(ctx context.Context) context.Context {
-	return context.WithValue(ctx, cacheContextKey{}, &authorizeCache{})
-}
-
-//nolint:revive
-func (c *authorizeCache) Load(subject Subject, action Action, object Object) (error, bool) {
-	if c == nil {
-		return nil, false
-	}
-	c.Lock()
-	defer c.Unlock()
-
-	for _, call := range c.calls {
-		if call.Action == action && call.Object.Equal(object) && call.Actor.Equal(subject) {
-			return call.Err, true
-		}
-	}
-	return nil, false
-}
-
-func (c *authorizeCache) Save(subject Subject, action Action, object Object, err error) {
-	if c == nil {
-		return
-	}
-	c.Lock()
-	defer c.Unlock()
-
-	c.calls = append(c.calls, cachedAuthCall{
-		AuthCall: AuthCall{
-			Actor:  subject,
-			Action: action,
-			Object: object,
-		},
-		Err: err,
-	})
-}
-
 // rbacTraceAttributes are the attributes that are added to all spans created by
 // the rbac package. These attributes should help to debug slow spans.
 func rbacTraceAttributes(actor Subject, action Action, objectType string, extra ...attribute.KeyValue) trace.SpanStartOption {

coderd/rbac/authz_test.go

@@ -282,27 +282,29 @@ func benchmarkSetup(orgs []uuid.UUID, users []uuid.UUID, size int, opts ...func(
 	return objectList
 }
 
-// BenchmarkCacher benchmarks the performance of the cacher with a given
+// BenchmarkCacher benchmarks the performance of the cacher.
-// cache size. The expected cache size in prod will usually be 1-2. In Filter
-// cases it can get as high as 10.
 func BenchmarkCacher(b *testing.B) {
-	b.ResetTimer()
+	ctx := context.Background()
-	// Size of the cache.
+	authz := rbac.Cacher(&coderdtest.FakeAuthorizer{AlwaysReturn: nil})
-	sizes := []int{1, 10, 100, 1000}
-	for _, size := range sizes {
-		b.Run(fmt.Sprintf("Size%d", size), func(b *testing.B) {
-			ctx := rbac.WithCacheCtx(context.Background())
-			authz := rbac.Cacher(&coderdtest.FakeAuthorizer{AlwaysReturn: nil})
-			for i := 0; i < size; i++ {
-				// Preload the cache of a given size
-				subj, obj, action := coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
-				_ = authz.Authorize(ctx, subj, action, obj)
-			}
 
-			// Cache is loaded as a slice, so this cache hit is always the last element.
+	rats := []int{1, 10, 100}
-			subj, obj, action := coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
+
-			b.ResetTimer()
+	for _, rat := range rats {
+		b.Run(fmt.Sprintf("%v:1", rat), func(b *testing.B) {
+			b.ReportAllocs()
+			var (
+				subj   rbac.Subject
+				obj    rbac.Object
+				action rbac.Action
+			)
 			for i := 0; i < b.N; i++ {
+				if i%rat == 0 {
+					// Cache miss
+					b.StopTimer()
+					subj, obj, action = coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
+					b.StartTimer()
+				}
+
 				_ = authz.Authorize(ctx, subj, action, obj)
 			}
 		})
@@ -312,29 +314,28 @@ func BenchmarkCacher(b *testing.B) {
 func TestCacher(t *testing.T) {
 	t.Parallel()
 
-	t.Run("EmptyCacheCtx", func(t *testing.T) {
+	t.Run("NoCache", func(t *testing.T) {
 		t.Parallel()
 
 		ctx := context.Background()
 		rec := &coderdtest.RecordingAuthorizer{
 			Wrapped: &coderdtest.FakeAuthorizer{AlwaysReturn: nil},
 		}
-		authz := rbac.Cacher(rec)
 		subj, obj, action := coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
 
 		// Two identical calls
-		_ = authz.Authorize(ctx, subj, action, obj)
+		_ = rec.Authorize(ctx, subj, action, obj)
-		_ = authz.Authorize(ctx, subj, action, obj)
+		_ = rec.Authorize(ctx, subj, action, obj)
 
 		// Yields two calls to the wrapped Authorizer
 		rec.AssertActor(t, subj, rec.Pair(action, obj), rec.Pair(action, obj))
 		require.NoError(t, rec.AllAsserted(), "all assertions should have been made")
 	})
 
-	t.Run("CacheCtx", func(t *testing.T) {
+	t.Run("Cache", func(t *testing.T) {
 		t.Parallel()
 
-		ctx := rbac.WithCacheCtx(context.Background())
+		ctx := context.Background()
 		rec := &coderdtest.RecordingAuthorizer{
 			Wrapped: &coderdtest.FakeAuthorizer{AlwaysReturn: nil},
 		}
@@ -353,7 +354,7 @@ func TestCacher(t *testing.T) {
 	t.Run("MultipleSubjects", func(t *testing.T) {
 		t.Parallel()
 
-		ctx := rbac.WithCacheCtx(context.Background())
+		ctx := context.Background()
 		rec := &coderdtest.RecordingAuthorizer{
 			Wrapped: &coderdtest.FakeAuthorizer{AlwaysReturn: nil},
 		}

coderd/users_test.go

@@ -1671,3 +1671,18 @@ func sortUsers(users []codersdk.User) {
 		return users[i].CreatedAt.Before(users[j].CreatedAt)
 	})
 }
+
+func BenchmarkUsersMe(b *testing.B) {
+	client := coderdtest.New(b, nil)
+	_ = coderdtest.CreateFirstUser(b, client)
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := client.User(ctx, codersdk.Me)
+		require.NoError(b, err)
+	}
+}

go.mod

@@ -174,6 +174,8 @@ require (
 	tailscale.com v1.32.2
 )
 
+require github.com/armon/go-radix v1.0.0 // indirect
+
 require (
 	cloud.google.com/go/compute v1.18.0 // indirect
 	cloud.google.com/go/logging v1.6.1 // indirect
@@ -189,6 +191,7 @@ require (
 	github.com/akutz/memconn v0.1.0 // indirect
 	github.com/alecthomas/chroma v0.10.0 // indirect
 	github.com/alexbrainman/sspi v0.0.0-20210105120005-909beea2cc74 // indirect
+	github.com/ammario/tlru v0.3.0
 	github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
 	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
 	github.com/aymanbagabas/go-osc52 v1.2.1 // indirect

go.sum

@@ -170,6 +170,8 @@ github.com/alexbrainman/sspi v0.0.0-20210105120005-909beea2cc74 h1:Kk6a4nehpJ3Uu
 github.com/alexbrainman/sspi v0.0.0-20210105120005-909beea2cc74/go.mod h1:cEWa1LVoE5KvSD9ONXsZrj0z6KqySlCCNKHlLzbqAt4=
 github.com/alexflint/go-filemutex v0.0.0-20171022225611-72bdc8eae2ae/go.mod h1:CgnQgUtFrFz9mxFNtED3jI5tLDjKlOM+oUF/sTk6ps0=
 github.com/alexflint/go-filemutex v1.1.0/go.mod h1:7P4iRhttt/nUvUOrYIhcpMzv2G6CY9UnI16Z+UJqRyk=
+github.com/ammario/tlru v0.3.0 h1:yK8ESoFlEyz/BVVL8yZQKAUzJwFJR/j9EfxjnKxtR/Q=
+github.com/ammario/tlru v0.3.0/go.mod h1:aYzRFu0XLo4KavE9W8Lx7tzjkX+pAApz+NgcKYIFUBQ=
 github.com/andybalholm/brotli v1.0.4 h1:V7DdXeJtZscaqfNuAdSRuRFzuiKlHSC/Zh3zl9qY3JY=
 github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
 github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8=
@@ -191,6 +193,8 @@ github.com/armon/circbuf v0.0.0-20190214190532-5111143e8da2/go.mod h1:3U/XgcO3hC
 github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
 github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY=
 github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
+github.com/armon/go-radix v1.0.0 h1:F4z6KzEeeQIMeLFa97iZU6vupzoecKdU5TX24SNppXI=
+github.com/armon/go-radix v1.0.0/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
 github.com/asaskevich/govalidator v0.0.0-20190424111038-f61b66f89f4a/go.mod h1:lB+ZfQJz7igIIfQNfa7Ml4HSf2uFQQRzpGGRXenZAgY=
 github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
 github.com/awalterschulze/gographviz v2.0.3+incompatible h1:9sVEXJBJLwGX7EQVhLm2elIKCm7P2YHFC8v6096G09E=

site/.eslintignore

@@ -56,6 +56,7 @@ stats/
 
 # direnv
 .envrc
+*.test
 # .prettierignore.include:
 # Helm templates contain variables that are invalid YAML and can't be formatted
 # by Prettier.

site/.prettierignore

@@ -56,6 +56,7 @@ stats/
 
 # direnv
 .envrc
+*.test
 # .prettierignore.include:
 # Helm templates contain variables that are invalid YAML and can't be formatted
 # by Prettier.