coder · Emyrk · Apr 18, 2023 · Apr 18, 2023 · Apr 18, 2023
diff --git a/coderd/rbac/authz_internal_test.go b/coderd/rbac/authz_internal_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/google/uuid"
@@ -12,6 +13,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/coderd/rbac/regosql"
 	"github.com/coder/coder/testutil"
 )
 
@@ -63,11 +65,14 @@ func TestFilterError(t *testing.T) {
 
 	t.Run("CancelledContext", func(t *testing.T) {
 		t.Parallel()
-		t.Skipf("This test is racy as rego eval checks the ctx canceled in a go routine. " +
-			"It is a coin flip if the query finishes before the 'cancel' is checked. " +
-			"So sometimes the 'Authorize' call succeeds even if ctx is canceled.")
 
-		auth := NewAuthorizer(prometheus.NewRegistry())
+		auth := &MockAuthorizer{
+			AuthorizeFunc: func(ctx context.Context, subject Subject, action Action, object Object) error {
+				// Authorize func always returns nil, unless the context is cancelled.
+				return ctx.Err()
+			},
+		}
+
 		subject := Subject{
 			ID: uuid.NewString(),
 			Roles: RoleNames{
@@ -77,20 +82,44 @@ func TestFilterError(t *testing.T) {
 			Scope:  ScopeAll,
 		}
 
-		ctx, cancel := context.WithCancel(context.Background())
-		defer cancel()
-		objects := []Objecter{
-			ResourceUser,
-			ResourceUser,
-			&objectBomb{
+		t.Run("SmallSet", func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			objects := []Objecter{
+				ResourceUser,
+				ResourceUser,
+				&objectBomb{
+					Objecter: ResourceUser,
+					bomb:     cancel,
+				},
+				ResourceUser,
+			}
+
+			_, err := Filter(ctx, auth, subject, ActionRead, objects)
+			require.ErrorIs(t, err, context.Canceled)
+		})
+
+		// Triggers Prepared Authorize
+		t.Run("LargeSet", func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			objects := make([]Objecter, 100)
+			for i := 0; i < 100; i++ {
+				objects[i] = ResourceUser
+			}
+			objects[20] = &objectBomb{
 				Objecter: ResourceUser,
 				bomb:     cancel,
-			},
-			ResourceUser,
-		}
+			}
 
-		_, err := Filter(ctx, auth, subject, ActionRead, objects)
-		require.ErrorIs(t, err, context.Canceled)
+			_, err := Filter(ctx, auth, subject, ActionRead, objects)
+			require.ErrorIs(t, err, context.Canceled)
+		})
 	})
 }
 
@@ -1095,3 +1124,42 @@ func must[T any](value T, err error) T {
 	}
 	return value
 }
+
+type MockAuthorizer struct {
+	AuthorizeFunc func(context.Context, Subject, Action, Object) error
+}
+
+var _ Authorizer = (*MockAuthorizer)(nil)
+
+func (d *MockAuthorizer) Authorize(ctx context.Context, s Subject, a Action, o Object) error {
+	return d.AuthorizeFunc(ctx, s, a, o)
+}
+
+func (d *MockAuthorizer) Prepare(_ context.Context, subject Subject, action Action, _ string) (PreparedAuthorized, error) {
+	return &mockPreparedAuthorizer{
+		Original: d,
+		Subject:  subject,
+		Action:   action,
+	}, nil
+}
+
+var _ PreparedAuthorized = (*mockPreparedAuthorizer)(nil)
+
+// fakePreparedAuthorizer is the prepared version of a FakeAuthorizer. It will
+// return the same error as the original FakeAuthorizer.
+type mockPreparedAuthorizer struct {
+	sync.RWMutex
+	Original *MockAuthorizer
+	Subject  Subject
+	Action   Action
+}
+
+func (f *mockPreparedAuthorizer) Authorize(ctx context.Context, object Object) error {
+	return f.Original.Authorize(ctx, f.Subject, f.Action, object)
+}
+
+// CompileToSQL returns a compiled version of the authorizer that will work for
+// in memory databases. This fake version will not work against a SQL database.
+func (*mockPreparedAuthorizer) CompileToSQL(_ context.Context, _ regosql.ConvertConfig) (string, error) {
+	return "not a valid sql string", nil
+}