coder · ammario · May 8, 2023 · May 5, 2023 · May 5, 2023 · May 5, 2023
diff --git a/.gitignore b/.gitignore
@@ -52,3 +52,4 @@ site/stats/
 
 # direnv
 .envrc
+*.test
diff --git a/.prettierignore b/.prettierignore
@@ -55,6 +55,7 @@ site/stats/
 
 # direnv
 .envrc
+*.test
 # .prettierignore.include:
 # Helm templates contain variables that are invalid YAML and can't be formatted
 # by Prettier.

diff --git a/coderd/coderd.go b/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),

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{}
 	}
@@ -190,8 +190,14 @@ func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.Can
 	}
 
 	if options.Authorizer == nil {
-		options.Authorizer = &RecordingAuthorizer{
-			Wrapped: rbac.NewCachingAuthorizer(prometheus.NewRegistry()),
+		defAuth := 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()

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/httpmw/authz.go b/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))
-	})
-}
diff --git a/coderd/rbac/authz.go b/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 {
-	cache := cacheFromContext(ctx)
+func (c *authCache) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
+	authorizeCacheKey := hashAuthorizeCall(subject, action, object)
 
-	resp, ok := cache.Load(subject, action, object)
-	if ok {
-		return resp
+	var err error
+	err, _, ok := c.cache.Get(authorizeCacheKey)
+	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 {