coder · Emyrk · Apr 3, 2023 · Mar 31, 2023 · Mar 31, 2023 · Mar 31, 2023
diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go
@@ -3,6 +3,7 @@ package rbac
 import (
 	"context"
 	_ "embed"
+	"strings"
 	"sync"
 	"time"
 
@@ -15,6 +16,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/coderd/rbac/regosql"
+	"github.com/coder/coder/coderd/rbac/regosql/sqltypes"
 	"github.com/coder/coder/coderd/tracing"
 	"github.com/coder/coder/coderd/util/slice"
 )
@@ -34,6 +36,12 @@ func AllActions() []Action {
 	return []Action{ActionCreate, ActionRead, ActionUpdate, ActionDelete}
 }
 
+type AuthCall struct {
+	Actor  Subject
+	Action Action
+	Object Object
+}
+
 // Subject is a struct that contains all the elements of a subject in an rbac
 // authorize.
 type Subject struct {
@@ -519,6 +527,160 @@ func (a RegoAuthorizer) newPartialAuthorizer(ctx context.Context, subject Subjec
 	return pAuth, nil
 }
 
+// AuthorizeFilter is a compiled partial query that can be converted to SQL.
+// This allows enforcing the policy on the database side in a WHERE clause.
+type AuthorizeFilter interface {
+	SQLString() string
+}
+
+type authorizedSQLFilter struct {
+	sqlString string
+	auth      *PartialAuthorizer
+}
+
+// ConfigWithACL is the basic configuration for converting rego to SQL when
+// the object has group and user ACL fields.
+func ConfigWithACL() regosql.ConvertConfig {
+	return regosql.ConvertConfig{
+		VariableConverter: regosql.DefaultVariableConverter(),
+	}
+}
+
+// ConfigWithoutACL is the basic configuration for converting rego to SQL when
+// the object has no ACL fields.
+func ConfigWithoutACL() regosql.ConvertConfig {
+	return regosql.ConvertConfig{
+		VariableConverter: regosql.NoACLConverter(),
+	}
+}
+
+func Compile(cfg regosql.ConvertConfig, pa *PartialAuthorizer) (AuthorizeFilter, error) {
+	root, err := regosql.ConvertRegoAst(cfg, pa.partialQueries)
+	if err != nil {
+		return nil, xerrors.Errorf("convert rego ast: %w", err)
+	}
+
+	// Generate the SQL
+	gen := sqltypes.NewSQLGenerator()
+	sqlString := root.SQLString(gen)
+	if len(gen.Errors()) > 0 {
+		var errStrings []string
+		for _, err := range gen.Errors() {
+			errStrings = append(errStrings, err.Error())
+		}
+		return nil, xerrors.Errorf("sql generation errors: %v", strings.Join(errStrings, ", "))
+	}
+
+	return &authorizedSQLFilter{
+		sqlString: sqlString,
+		auth:      pa,
+	}, nil
+}
+
+func (a *authorizedSQLFilter) SQLString() string {
+	return a.sqlString
+}
+
+type cachedCalls struct {
+	authz Authorizer
+}
+
+// Cacher returns an Authorizer that can use a cache stored on a context
+// to short circuit duplicate calls to the Authorizer. This is useful when
+// multiple calls are made to the Authorizer for the same subject, action, and
+// object. The cache is on each `ctx` and is not shared between requests.
+// If no cache is found on the context, the Authorizer is called as normal.
+//
+// Cacher is safe for multiple actors.
+func Cacher(authz Authorizer) Authorizer {
+	return &cachedCalls{authz: authz}
+}
+
+func (c *cachedCalls) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
+	cache := cacheFromContext(ctx)
+
+	resp, ok := cache.Load(subject, action, object)
+	if ok {
+		return resp
+	}
+
+	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) {
+	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
@@ -2,12 +2,14 @@ package rbac_test
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 
+	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/coderd/rbac"
 )
 
@@ -233,3 +235,105 @@ func benchmarkSetup(orgs []uuid.UUID, users []uuid.UUID, size int, opts ...func(
 
 	return objectList
 }
+
+// BenchmarkCacher benchmarks the performance of the cacher with a given
+// 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()
+	// Size of the cache.
+	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.
+			subj, obj, action := coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_ = authz.Authorize(ctx, subj, action, obj)
+			}
+		})
+	}
+}
+
+func TestCacher(t *testing.T) {
+	t.Parallel()
+
+	t.Run("EmptyCacheCtx", 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)
+		_ = authz.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.Parallel()
+
+		ctx := rbac.WithCacheCtx(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)
+		_ = authz.Authorize(ctx, subj, action, obj)
+
+		// Yields only 1 call to the wrapped Authorizer for that subject
+		rec.AssertActor(t, subj, rec.Pair(action, obj))
+		require.NoError(t, rec.AllAsserted(), "all assertions should have been made")
+	})
+
+	t.Run("MultipleSubjects", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := rbac.WithCacheCtx(context.Background())
+		rec := &coderdtest.RecordingAuthorizer{
+			Wrapped: &coderdtest.FakeAuthorizer{AlwaysReturn: nil},
+		}
+		authz := rbac.Cacher(rec)
+		subj1, obj1, action1 := coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
+
+		// Two identical calls
+		_ = authz.Authorize(ctx, subj1, action1, obj1)
+		_ = authz.Authorize(ctx, subj1, action1, obj1)
+
+		// Extra unique calls
+		var pairs []coderdtest.ActionObjectPair
+		subj2, obj2, action2 := coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
+		_ = authz.Authorize(ctx, subj2, action2, obj2)
+		pairs = append(pairs, rec.Pair(action2, obj2))
+
+		obj3, action3 := coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
+		_ = authz.Authorize(ctx, subj2, action3, obj3)
+		pairs = append(pairs, rec.Pair(action3, obj3))
+
+		// Extra identical call after some unique calls
+		_ = authz.Authorize(ctx, subj1, action1, obj1)
+
+		// Yields 3 calls, 1 for the first subject, 2 for the unique subjects
+		rec.AssertActor(t, subj1, rec.Pair(action1, obj1))
+		rec.AssertActor(t, subj2, pairs...)
+		require.NoError(t, rec.AllAsserted(), "all assertions should have been made")
+	})
+}