chore: use global cache for RBAC

coder · ammario · May 8, 2023 · May 5, 2023 · May 5, 2023 · May 5, 2023
commit 032338f918c1d2329018f7fb6be801e725861818
diff --git a/coderd.test b/coderd.test
diff --git a/coderd/coderdtest/coderdtest.go b/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{}
 	}
@@ -359,7 +359,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 +384,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 +465,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 +1111,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()

diff --git a/coderd/database/dbtestutil/db.go b/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()

diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go
@@ -3,10 +3,12 @@ package rbac
 import (
 	"context"
 	_ "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"
@@ -310,6 +312,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 +608,13 @@ func (a *authorizedSQLFilter) SQLString() string {
 	return a.sqlString
 }
 
-type cachedCalls struct {
+type authCache struct {
+	// cache is a cache of JSON marshaled inputs to "authorize" to
+	// corresponding errors. When the Authorizer is immutable (e.g. in the case
+	// of Rego)"
+	// determistic function.
+	cache *tlru.Cache[string, error]
+
 	authz Authorizer
 }
 
@@ -617,94 +626,37 @@ type cachedCalls struct {
 //
 // Cacher is safe for multiple actors.
 func Cacher(authz Authorizer) Authorizer {
-	return &cachedCalls{authz: authz}
+	return &authCache{
+		authz: authz,
+		cache: tlru.New[string](tlru.ConstantCost[error], 4096),
+	}
 }
 
-func (c *cachedCalls) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
-	cache := cacheFromContext(ctx)
+func (c *authCache) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
+	var authorizeCacheKey strings.Builder
+	enc := json.NewEncoder(&authorizeCacheKey)
+	_ = enc.Encode(subject)
+	_ = enc.Encode(action)
+	_ = enc.Encode(object)
 
-	resp, ok := cache.Load(subject, action, object)
-	if ok {
-		return resp
+	var err error
+	err, _, ok := c.cache.Get(authorizeCacheKey.String())
+	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.String(), 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 {

diff --git a/coderd/rbac/authz_test.go b/coderd/rbac/authz_test.go
@@ -291,7 +291,7 @@ func BenchmarkCacher(b *testing.B) {
 	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())
+			ctx := context.Background()
 			authz := rbac.Cacher(&coderdtest.FakeAuthorizer{AlwaysReturn: nil})
 			for i := 0; i < size; i++ {
 				// Preload the cache of a given size

diff --git a/coderd/users_test.go b/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)
+	}
+}
diff --git a/go.mod b/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

diff --git a/go.sum b/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=