Add ability to do rbac checks in unit tests

coder · Emyrk · Mar 21, 2023 · Mar 16, 2023 · Mar 17, 2023 · Mar 17, 2023
commit 7118bf09a332faf477b77de011950f163cd50b65
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -164,6 +164,9 @@ func New(options *Options) *API {
 		options = &Options{}
 	}
 
+	if options.Authorizer == nil {
+		options.Authorizer = rbac.NewCachingAuthorizer(options.PrometheusRegistry)
+	}
 	options.Database = dbauthz.New(
 		options.Database,
 		options.Authorizer,
@@ -204,9 +207,6 @@ func New(options *Options) *API {
 	if options.PrometheusRegistry == nil {
 		options.PrometheusRegistry = prometheus.NewRegistry()
 	}
-	if options.Authorizer == nil {
-		options.Authorizer = rbac.NewCachingAuthorizer(options.PrometheusRegistry)
-	}
 	if options.TailnetCoordinator == nil {
 		options.TailnetCoordinator = tailnet.NewCoordinator()
 	}

diff --git a/coderd/coderdtest/authorize.go b/coderd/coderdtest/authorize.go
@@ -19,16 +19,92 @@ import (
 	"github.com/stretchr/testify/require"
 	"golang.org/x/xerrors"
 
-	"github.com/coder/coder/cryptorand"
-
 	"github.com/coder/coder/coderd"
+	"github.com/coder/coder/coderd/database/dbauthz"
 	"github.com/coder/coder/coderd/rbac"
 	"github.com/coder/coder/coderd/rbac/regosql"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/cryptorand"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
 )
 
+type CoderSDKObject interface {
+	codersdk.User
+}
+
+func RBACObject[C CoderSDKObject](o C) rbac.Object {
+	switch ro := any(o).(type) {
+	case codersdk.User:
+		return rbac.ResourceUser.WithID(ro.ID)
+	default:
+		panic("unknown object type")
+	}
+}
+
+// RBACAsserter is a helper for asserting that the correct RBAC checks are
+// performed. This struct is tied to a given user, and only authorizes calls
+// for this user are checked.
+type RBACAsserter struct {
+	Subject rbac.Subject
+
+	Recorder *RecordingAuthorizer
+}
+
+func (a RBACAsserter) AllCalls() []rbac.AuthCall {
+	return a.Recorder.AllCalls(&a.Subject)
+}
+
+// AssertChecked will assert a given rbac check was performed. It does not care
+// about order of checks, or any other checks. This is useful when you do not
+// care about asserting every check that was performed.
+func (a RBACAsserter) AssertChecked(t *testing.T, action rbac.Action, objects ...rbac.Object) {
+	pairs := make([]ActionObjectPair, 0, len(objects))
+	for _, obj := range objects {
+		pairs = append(pairs, a.Recorder.Pair(action, obj))
+	}
+	a.Recorder.AssertOutOfOrder(t, a.Subject, pairs...)
+}
+
+// AssertInOrder must be called in the correct order of authz checks. If the objects
+// or actions are not in the correct order, the test will fail.
+func (a RBACAsserter) AssertInOrder(t *testing.T, action rbac.Action, objects ...rbac.Object) {
+	pairs := make([]ActionObjectPair, 0, len(objects))
+	for _, obj := range objects {
+		pairs = append(pairs, a.Recorder.Pair(action, obj))
+	}
+	a.Recorder.AssertActor(t, a.Subject, pairs...)
+}
+
+func AssertRBAC(t *testing.T, api *coderd.API, client *codersdk.Client) RBACAsserter {
+	recorder, ok := api.Authorizer.(*RecordingAuthorizer)
+	if !ok {
+		t.Fatal("expected RecordingAuthorizer")
+	}
+
+	// We use the database directly to not cause additional auth checks on behalf
+	// of the user. This does add authz checks on behalf of the system user, but
+	// it is hard to avoid that.
+	ctx := dbauthz.AsSystemRestricted(context.Background())
+	token := client.SessionToken()
+	parts := strings.Split(token, "-")
+	key, err := api.Database.GetAPIKeyByID(ctx, parts[0])
+	require.NoError(t, err, "fetch client api key")
+
+	roles, err := api.Database.GetAuthorizationUserRoles(ctx, key.UserID)
+	require.NoError(t, err, "fetch user roles")
+
+	return RBACAsserter{
+		Subject: rbac.Subject{
+			ID:     key.UserID.String(),
+			Roles:  rbac.RoleNames(roles.Roles),
+			Groups: roles.Groups,
+			Scope:  rbac.ScopeName(key.Scope),
+		},
+		Recorder: recorder,
+	}
+}
+
 func AGPLRoutes(a *AuthTester) (map[string]string, map[string]RouteCheck) {
 	// Some quick reused objects
 	workspaceRBACObj := rbac.ResourceWorkspace.WithID(a.Workspace.ID).InOrg(a.Organization.ID).WithOwner(a.Workspace.OwnerID.String())
@@ -598,11 +674,48 @@ func (r *RecordingAuthorizer) AllAsserted() error {
 	return nil
 }
 
+// AllCalls is useful for debugging.
+func (r *RecordingAuthorizer) AllCalls(actor *rbac.Subject) []rbac.AuthCall {
+	r.RLock()
+	defer r.RUnlock()
+
+	called := make([]rbac.AuthCall, 0, len(r.Called))
+	for _, c := range r.Called {
+		if actor != nil && !c.Actor.Equal(*actor) {
+			continue
+		}
+		called = append(called, c.AuthCall)
+	}
+	return called
+}
+
+// AssertOutOfOrder asserts that the given actor performed the given action
+// on the given objects. It does not care about the order of the calls.
+// When marking authz calls as asserted, it will mark the first matching
+// calls first.
+func (r *RecordingAuthorizer) AssertOutOfOrder(t *testing.T, actor rbac.Subject, did ...ActionObjectPair) {
+	r.Lock()
+	defer r.Unlock()
+
+	for _, do := range did {
+		found := false
+		// Find the first non-asserted call that matches the actor, action, and object.
+		for i, call := range r.Called {
+			if !call.asserted && call.Actor.Equal(actor) && call.Action == do.Action && call.Object.Equal(do.Object) {
+				r.Called[i].asserted = true
+				found = true
+				break
+			}
+		}
+		require.True(t, found, "assertion missing: %s %s %s", actor, do.Action, do.Object)
+	}
+}
+
 // AssertActor asserts in order. If the order of authz calls does not match,
 // this will fail.
 func (r *RecordingAuthorizer) AssertActor(t *testing.T, actor rbac.Subject, did ...ActionObjectPair) {
-	r.RLock()
-	defer r.RUnlock()
+	r.Lock()
+	defer r.Unlock()
 	ptr := 0
 	for i, call := range r.Called {
 		if ptr == len(did) {

diff --git a/coderd/coderdtest/authorize_test.go b/coderd/coderdtest/authorize_test.go
@@ -2,6 +2,7 @@ package coderdtest_test
 
 import (
 	"context"
+	"math/rand"
 	"os"
 	"strings"
 	"testing"
@@ -81,6 +82,42 @@ func TestAuthzRecorder(t *testing.T) {
 		rec.AssertActor(t, a, aPairs...)
 		require.NoError(t, rec.AllAsserted(), "all assertions should have been made")
 	})
+
+	t.Run("AuthorizeOutOfOrder", func(t *testing.T) {
+		t.Parallel()
+
+		rec := &coderdtest.RecordingAuthorizer{
+			Wrapped: &coderdtest.FakeAuthorizer{},
+		}
+		sub := coderdtest.RandomRBACSubject()
+		pairs := fuzzAuthz(t, sub, rec, 10)
+		rand.Shuffle(len(pairs), func(i, j int) {
+			pairs[i], pairs[j] = pairs[j], pairs[i]
+		})
+
+		rec.AssertOutOfOrder(t, sub, pairs...)
+		require.NoError(t, rec.AllAsserted(), "all assertions should have been made")
+	})
+
+	t.Run("AllCalls", func(t *testing.T) {
+		t.Parallel()
+
+		rec := &coderdtest.RecordingAuthorizer{
+			Wrapped: &coderdtest.FakeAuthorizer{},
+		}
+		sub := coderdtest.RandomRBACSubject()
+		calls := rec.AllCalls(&sub)
+		pairs := make([]coderdtest.ActionObjectPair, 0, len(calls))
+		for _, call := range calls {
+			pairs = append(pairs, coderdtest.ActionObjectPair{
+				Action: call.Action,
+				Object: call.Object,
+			})
+		}
+
+		rec.AssertActor(t, sub, pairs...)
+		require.NoError(t, rec.AllAsserted(), "all assertions should have been made")
+	})
 }
 
 // fuzzAuthzPrep has same action and object types for all calls.

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -180,17 +180,18 @@ func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.Can
 			close(options.AutobuildStats)
 		})
 	}
-	if options.Database == nil {
-		options.Database, options.Pubsub = dbtestutil.NewDB(t)
-	}
-	options.Database = dbauthz.New(options.Database, options.Authorizer, slogtest.Make(t, nil).Leveled(slog.LevelDebug))
 
 	if options.Authorizer == nil {
 		options.Authorizer = &RecordingAuthorizer{
 			Wrapped: rbac.NewCachingAuthorizer(prometheus.NewRegistry()),
 		}
 	}
 
+	if options.Database == nil {
+		options.Database, options.Pubsub = dbtestutil.NewDB(t)
+		options.Database = dbauthz.New(options.Database, options.Authorizer, slogtest.Make(t, nil).Leveled(slog.LevelDebug))
+	}
+
 	if options.DeploymentValues == nil {
 		options.DeploymentValues = DeploymentValues(t)
 	}