coder/coderd/database/dbauthz/setup_test.go

package dbauthz_test

import (
	"context"
	"encoding/gob"
	"errors"
	"flag"
	"fmt"
	"reflect"
	"slices"
	"strings"
	"testing"

	"github.com/brianvoe/gofakeit/v7"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"github.com/google/uuid"
	"github.com/open-policy-agent/opa/topdown"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	"go.uber.org/mock/gomock"
	"golang.org/x/xerrors"

	"cdr.dev/slog/v3"
	"github.com/coder/coder/v2/coderd/coderdtest"
	"github.com/coder/coder/v2/coderd/database"
	"github.com/coder/coder/v2/coderd/database/dbauthz"
	"github.com/coder/coder/v2/coderd/database/dbmock"
	"github.com/coder/coder/v2/coderd/database/dbtestutil"
	"github.com/coder/coder/v2/coderd/rbac"
	"github.com/coder/coder/v2/coderd/rbac/policy"
	"github.com/coder/coder/v2/coderd/rbac/regosql"
	"github.com/coder/coder/v2/coderd/rbac/rolestore"
	"github.com/coder/coder/v2/coderd/util/slice"
)

var errMatchAny = xerrors.New("match any error")

var skipMethods = map[string]string{
	"InTx":           "Not relevant",
	"Ping":           "Not relevant",
	"PGLocks":        "Not relevant",
	"Wrappers":       "Not relevant",
	"AcquireLock":    "Not relevant",
	"TryAcquireLock": "Not relevant",
}

// TestMethodTestSuite runs MethodTestSuite.
// In order for 'go test' to run this suite, we need to create
// a normal test function and pass our suite to suite.Run
// nolint: paralleltest
func TestMethodTestSuite(t *testing.T) {
	suite.Run(t, new(MethodTestSuite))
}

// MethodTestSuite runs all methods tests for querier. We use
// a test suite so we can account for all functions tested on the querier.
// We can then assert all methods were tested and asserted for proper RBAC
// checks. This forces RBAC checks to be written for all methods.
// Additionally, the way unit tests are written allows for easily executing
// a single test for debugging.
type MethodTestSuite struct {
	suite.Suite
	// methodAccounting counts all methods called by a 'RunMethodTest'
	methodAccounting map[string]int
}

// SetupSuite sets up the suite by creating a map of all methods on querier
// and setting their count to 0.
func (s *MethodTestSuite) SetupSuite() {
	ctrl := gomock.NewController(s.T())
	mockStore := dbmock.NewMockStore(ctrl)
	// We intentionally set no expectations apart from this.
	mockStore.EXPECT().Wrappers().Return([]string{}).AnyTimes()
	az := dbauthz.New(mockStore, nil, slog.Make(), coderdtest.AccessControlStorePointer())
	// Take the underlying type of the interface.
	azt := reflect.TypeOf(az)
	require.Greater(s.T(), azt.NumMethod(), 0, "no methods found on querier")
	s.methodAccounting = make(map[string]int)
	for i := 0; i < azt.NumMethod(); i++ {
		method := azt.Method(i)
		if _, ok := skipMethods[method.Name]; ok {
			// We can't use s.T().Skip as this will skip the entire suite.
			s.T().Logf("Skipping method %q: %s", method.Name, skipMethods[method.Name])
			continue
		}
		s.methodAccounting[method.Name] = 0
	}
}

// TearDownSuite asserts that all methods were called at least once.
func (s *MethodTestSuite) TearDownSuite() {
	s.Run("Accounting", func() {
		// testify/suite's -testify.m flag filters which suite methods
		// run, but TearDownSuite still executes. Skip the Accounting
		// check when filtering to avoid misleading "method never
		// called" errors for every method that was filtered out.
		if f := flag.Lookup("testify.m"); f != nil {
			if f.Value.String() != "" {
				s.T().Skip("Skipping Accounting check: -testify.m flag is set")
			}
		}

		t := s.T()
		notCalled := []string{}
		for m, c := range s.methodAccounting {
			if c <= 0 {
				notCalled = append(notCalled, m)
			}
		}
		slices.Sort(notCalled)
		for _, m := range notCalled {
			t.Errorf("Method never called: %q", m)
		}
	})
}

var testActorID = uuid.New()

type includeSystemRolesMatcher struct{}

func (includeSystemRolesMatcher) Matches(x any) bool {
	p, ok := x.(database.CustomRolesParams)
	if !ok {
		return false
	}
	return p.IncludeSystemRoles
}

func (includeSystemRolesMatcher) String() string {
	return "CustomRolesParams with IncludeSystemRoles=true"
}

// Mocked runs a subtest with a mocked database. Removing the overhead of a real
// postgres database resulting in much faster tests.
func (s *MethodTestSuite) Mocked(testCaseF func(dmb *dbmock.MockStore, faker *gofakeit.Faker, check *expects)) func() {
	t := s.T()
	mDB := dbmock.NewMockStore(gomock.NewController(t))
	mDB.EXPECT().Wrappers().Return([]string{}).AnyTimes()
	// dbauthz now expands DB-backed system roles (e.g. organization-member)
	// during role-assignment validation, which triggers a CustomRoles lookup
	// with IncludeSystemRoles=true.
	mDB.EXPECT().CustomRoles(gomock.Any(), includeSystemRolesMatcher{}).DoAndReturn(func(_ context.Context, arg database.CustomRolesParams) ([]database.CustomRole, error) {
		if len(arg.LookupRoles) == 0 {
			return []database.CustomRole{}, nil
		}

		out := make([]database.CustomRole, 0, len(arg.LookupRoles))

		for _, pair := range arg.LookupRoles {
			// Minimal set of fields that the tested code uses.
			out = append(out, database.CustomRole{
				Name: pair.Name,
				OrganizationID: uuid.NullUUID{
					UUID:  pair.OrganizationID,
					Valid: pair.OrganizationID != uuid.Nil,
				},
				IsSystem: rolestore.IsSystemRoleName(pair.Name),
				ID:       uuid.New(),
			})
		}

		return out, nil
	}).AnyTimes()

	// Use a constant seed to prevent flakes from random data generation.
	faker := gofakeit.New(0)

	// The usual Subtest assumes the test setup will use a real database to populate
	// with data. In this mocked case, we want to pass the underlying mocked database
	// to the test case instead.
	return s.SubtestWithDB(mDB, func(_ database.Store, check *expects) {
		testCaseF(mDB, faker, check)
	})
}

// Subtest starts up a real postgres database for each test case.
// Deprecated: Use 'Mocked' instead for much faster tests.
func (s *MethodTestSuite) Subtest(testCaseF func(db database.Store, check *expects)) func() {
	t := s.T()
	db, _ := dbtestutil.NewDB(t)
	return s.SubtestWithDB(db, testCaseF)
}

// SubtestWithDB is a helper function that returns a function that can be passed to
// s.Run(). This function will run the test case for the method that is being
// tested. The check parameter is used to assert the results of the method.
// If the caller does not use the `check` parameter, the test will fail.
func (s *MethodTestSuite) SubtestWithDB(db database.Store, testCaseF func(db database.Store, check *expects)) func() {
	return func() {
		t := s.T()
		testName := s.T().Name()
		names := strings.Split(testName, "/")
		methodName := names[len(names)-1]
		// Repeated subtests get "#NN" suffixes; count them under the base method.
		if baseMethodName, _, ok := strings.Cut(methodName, "#"); ok {
			methodName = baseMethodName
		}
		s.methodAccounting[methodName]++

		fakeAuthorizer := &coderdtest.FakeAuthorizer{}
		rec := &coderdtest.RecordingAuthorizer{
			Wrapped: fakeAuthorizer,
		}
		az := dbauthz.New(db, rec, slog.Make(), coderdtest.AccessControlStorePointer())
		actor := rbac.Subject{
			ID:     testActorID.String(),
			Roles:  rbac.RoleIdentifiers{rbac.RoleOwner()},
			Groups: []string{},
			Scope:  rbac.ScopeAll,
		}
		ctx := dbauthz.As(context.Background(), actor)

		var testCase expects
		testCaseF(db, &testCase)
		// Check the developer added assertions. If there are no assertions,
		// an empty list should be passed.
		s.Require().False(testCase.assertions == nil, "rbac assertions not set, use the 'check' parameter")

		// Find the method with the name of the test.
		var callMethod func(ctx context.Context) ([]reflect.Value, error)
		azt := reflect.TypeOf(az)
	MethodLoop:
		for i := 0; i < azt.NumMethod(); i++ {
			method := azt.Method(i)
			if method.Name == methodName {
				methodF := reflect.ValueOf(az).Method(i)

				callMethod = func(ctx context.Context) ([]reflect.Value, error) {
					resp := methodF.Call(append([]reflect.Value{reflect.ValueOf(ctx)}, testCase.inputs...))
					return splitResp(t, resp)
				}
				break MethodLoop
			}
		}

		require.NotNil(t, callMethod, "method %q does not exist", methodName)

		if len(testCase.assertions) > 0 {
			// Only run these tests if we know the underlying call makes
			// rbac assertions.
			s.NotAuthorizedErrorTest(ctx, fakeAuthorizer, testCase, callMethod)
		}

		if len(testCase.assertions) > 0 ||
			slice.Contains([]string{
				"GetAuthorizedWorkspaces",
				"GetAuthorizedTemplates",
				"GetDefaultChatModelConfig",
			}, methodName) {
			// Some methods do not make RBAC assertions because they use
			// SQL. We still want to test that they return an error if the
			// actor is not set.
			s.NoActorErrorTest(callMethod)
		}

		// Always run
		s.Run("Success", func() {
			rec.Reset()
			if testCase.successAuthorizer != nil {
				fakeAuthorizer.ConditionalReturn = testCase.successAuthorizer
			} else {
				fakeAuthorizer.AlwaysReturn(nil)
			}

			outputs, err := callMethod(ctx)
			if testCase.err == nil {
				s.NoError(err, "method %q returned an error", methodName)
			} else {
				if errors.Is(testCase.err, errMatchAny) {
					// This means we do not care exactly what the error is.
					s.Error(err, "method %q returned an error", methodName)
				} else {
					s.EqualError(err, testCase.err.Error(), "method %q returned an unexpected error", methodName)
				}
			}

			// Some tests may not care about the outputs, so we only assert if
			// they are provided.
			if testCase.outputs != nil {
				// Assert the required outputs
				s.Equal(len(testCase.outputs), len(outputs), "method %q returned unexpected number of outputs", methodName)
				cmpOptions := []cmp.Option{
					// Equate nil and empty slices.
					cmpopts.EquateEmpty(),
				}
				for i := range outputs {
					a, b := testCase.outputs[i].Interface(), outputs[i].Interface()

					// To avoid the extra small overhead of gob encoding, we can
					// first check if the values are equal with regard to order.
					// If not, re-check disregarding order and show a nice diff
					// output of the two values.
					if !cmp.Equal(a, b, cmpOptions...) {
						diffOpts := append(
							append([]cmp.Option{}, cmpOptions...),
							// Allow slice order to be ignored.
							cmpopts.SortSlices(func(a, b any) bool {
								var ab, bb strings.Builder
								_ = gob.NewEncoder(&ab).Encode(a)
								_ = gob.NewEncoder(&bb).Encode(b)
								// This might seem a bit dubious, but we really
								// don't care about order and cmp doesn't provide
								// a generic less function for slices:
								// https://github.com/google/go-cmp/issues/67
								return ab.String() < bb.String()
							}),
						)
						if diff := cmp.Diff(a, b, diffOpts...); diff != "" {
							s.Failf("compare outputs failed", "method %q returned unexpected output %d (-want +got):\n%s", methodName, i, diff)
						}
					}
				}
			}

			var pairs []coderdtest.ActionObjectPair
			for _, assrt := range testCase.assertions {
				for _, action := range assrt.Actions {
					pairs = append(pairs, coderdtest.ActionObjectPair{
						Action: action,
						Object: assrt.Object,
					})
				}
			}

			if testCase.outOfOrder {
				rec.AssertOutOfOrder(s.T(), actor, pairs...)
			} else {
				rec.AssertActor(s.T(), actor, pairs...)
			}
			s.NoError(rec.AllAsserted(), "all rbac calls must be asserted")
		})
	}
}

func (s *MethodTestSuite) NoActorErrorTest(callMethod func(ctx context.Context) ([]reflect.Value, error)) {
	s.Run("AsRemoveActor", func() {
		// Call without any actor
		_, err := callMethod(context.Background())
		s.ErrorIs(err, dbauthz.ErrNoActor, "method should return NoActorError error when no actor is provided")
	})
}

// NotAuthorizedErrorTest runs the given method with an authorizer that will fail authz.
// Asserts that the error returned is a NotAuthorizedError.
func (s *MethodTestSuite) NotAuthorizedErrorTest(ctx context.Context, az *coderdtest.FakeAuthorizer, testCase expects, callMethod func(ctx context.Context) ([]reflect.Value, error)) {
	s.Run("NotAuthorized", func() {
		az.AlwaysReturn(rbac.ForbiddenWithInternal(xerrors.New("Always fail authz"), rbac.Subject{}, "", rbac.Object{}, nil))
		// Override the SQL filter to always fail.
		az.OverrideSQLFilter("FALSE")

		// If we have assertions, that means the method should FAIL
		// if RBAC will disallow the request. The returned error should
		// be expected to be a NotAuthorizedError.
		resp, err := callMethod(ctx)

		// This is unfortunate, but if we are using `Filter` the error returned will be nil. So filter out
		// any case where the error is nil and the response is an empty slice.
		if err != nil || !hasEmptyResponse(resp) {
			// Expect the default error
			if testCase.notAuthorizedExpect == "" {
				s.ErrorContainsf(err, "unauthorized", "error string should have a good message")
				s.Errorf(err, "method should an error with disallow authz")
				s.ErrorAs(err, &dbauthz.NotAuthorizedError{}, "error should be NotAuthorizedError")
			} else {
				s.ErrorContains(err, testCase.notAuthorizedExpect)
			}
		}
	})

	s.Run("Canceled", func() {
		// Pass in a canceled context
		ctx, cancel := context.WithCancel(ctx)
		cancel()
		az.AlwaysReturn(rbac.ForbiddenWithInternal(&topdown.Error{Code: topdown.CancelErr},
			rbac.Subject{}, "", rbac.Object{}, nil))

		// If we have assertions, that means the method should FAIL
		// if RBAC will disallow the request. The returned error should
		// be expected to be a NotAuthorizedError.
		resp, err := callMethod(ctx)

		// This is unfortunate, but if we are using `Filter` the error returned will be nil. So filter out
		// any case where the error is nil and the response is an empty slice or int64(0).
		if err != nil || !hasEmptyResponse(resp) {
			if testCase.cancelledCtxExpect == "" {
				s.Errorf(err, "method should an error with cancellation")
				s.ErrorIsf(err, context.Canceled, "error should match context.Canceled")
			} else {
				s.ErrorContains(err, testCase.cancelledCtxExpect)
			}
		}
	})
}

func hasEmptyResponse(values []reflect.Value) bool {
	for _, r := range values {
		if r.Kind() == reflect.Slice || r.Kind() == reflect.Array {
			if r.Len() == 0 {
				return true
			}
		}

		// Special case for int64, as it's the return type for count queries.
		if r.Kind() == reflect.Int64 {
			if r.Int() == 0 {
				return true
			}
		}
	}
	return false
}

func splitResp(t *testing.T, values []reflect.Value) ([]reflect.Value, error) {
	outputs := []reflect.Value{}
	for _, r := range values {
		if r.Type().Implements(reflect.TypeOf((*error)(nil)).Elem()) {
			if r.IsNil() {
				// Error is found, but it's nil!
				return outputs, nil
			}
			err, ok := r.Interface().(error)
			if !ok {
				t.Fatal("error is not an error?!")
			}
			return outputs, err
		}
		outputs = append(outputs, r)
	}
	t.Fatal("no expected error value found in responses (error can be nil)")
	return nil, nil // unreachable, required to compile
}

// expects is used to build a test case for a method.
// It includes the expected inputs, rbac assertions, and expected outputs.
type expects struct {
	inputs     []reflect.Value
	assertions []AssertRBAC
	// outputs is optional. Can assert non-error return values.
	outputs []reflect.Value
	err     error

	// Optional override of the default error checks.
	// By default, we search for the expected error strings.
	// If these strings are present, these strings will be searched
	// instead.
	notAuthorizedExpect string
	cancelledCtxExpect  string
	successAuthorizer   func(ctx context.Context, subject rbac.Subject, action policy.Action, obj rbac.Object) error
	outOfOrder          bool
}

// OutOfOrder is optional. It controls whether the assertions should be
// asserted in order.
func (m *expects) OutOfOrder() *expects {
	m.outOfOrder = true
	return m
}

// Asserts is required. Asserts the RBAC authorize calls that should be made.
// If no RBAC calls are expected, pass an empty list: 'm.Asserts()'
func (m *expects) Asserts(pairs ...any) *expects {
	m.assertions = asserts(pairs...)
	return m
}

// Args is required. The arguments to be provided to the method.
// If there are no arguments, pass an empty list: 'm.Args()'
// The first context argument should not be included, as the test suite
// will provide it.
func (m *expects) Args(args ...any) *expects {
	m.inputs = values(args...)
	return m
}

// Returns is optional. If it is never called, it will not be asserted.
func (m *expects) Returns(rets ...any) *expects {
	m.outputs = values(rets...)
	return m
}

// Errors is optional. If it is never called, it will not be asserted.
func (m *expects) Errors(err error) *expects {
	m.err = err
	return m
}

func (m *expects) FailSystemObjectChecks() *expects {
	return m.WithSuccessAuthorizer(func(ctx context.Context, subject rbac.Subject, action policy.Action, obj rbac.Object) error {
		if obj.Type == rbac.ResourceSystem.Type {
			return xerrors.Errorf("hard coded system authz failed")
		}
		return nil
	})
}

// WithSuccessAuthorizer is helpful when an optimization authz check is made
// to skip some RBAC checks. This check in testing would prevent the ability
// to assert the more nuanced RBAC checks.
func (m *expects) WithSuccessAuthorizer(f func(ctx context.Context, subject rbac.Subject, action policy.Action, obj rbac.Object) error) *expects {
	m.successAuthorizer = f
	return m
}

func (m *expects) WithNotAuthorized(contains string) *expects {
	m.notAuthorizedExpect = contains
	return m
}

func (m *expects) WithCancelled(contains string) *expects {
	m.cancelledCtxExpect = contains
	return m
}

// AssertRBAC contains the object and actions to be asserted.
type AssertRBAC struct {
	Object  rbac.Object
	Actions []policy.Action
}

// values is a convenience method for creating []reflect.Value.
//
// values(workspace, template, ...)
//
// is equivalent to
//
//	[]reflect.Value{
//	  reflect.ValueOf(workspace),
//	  reflect.ValueOf(template),
//	  ...
//	}
func values(ins ...any) []reflect.Value {
	out := make([]reflect.Value, 0)
	for _, input := range ins {
		out = append(out, reflect.ValueOf(input))
	}
	return out
}

// asserts is a convenience method for creating AssertRBACs.
//
// The number of inputs must be an even number.
// asserts() will panic if this is not the case.
//
// Even-numbered inputs are the objects, and odd-numbered inputs are the actions.
// Objects must implement rbac.Objecter.
// Inputs can be a single policy.Action, or a slice of policy.Action.
//
//	asserts(workspace, policy.ActionRead, template, slice(policy.ActionRead, policy.ActionWrite), ...)
//
// is equivalent to
//
//	[]AssertRBAC{
//	  {Object: workspace, Actions: []policy.Action{policy.ActionRead}},
//	  {Object: template, Actions: []policy.Action{policy.ActionRead, policy.ActionWrite)}},
//	   ...
//	}
func asserts(inputs ...any) []AssertRBAC {
	if len(inputs)%2 != 0 {
		panic(fmt.Sprintf("Must be an even length number of args, found %d", len(inputs)))
	}

	out := make([]AssertRBAC, 0)
	for i := 0; i < len(inputs); i += 2 {
		obj, ok := inputs[i].(rbac.Objecter)
		if !ok {
			panic(fmt.Sprintf("object type '%T' does not implement rbac.Objecter", inputs[i]))
		}
		rbacObj := obj.RBACObject()

		var actions []policy.Action
		actions, ok = inputs[i+1].([]policy.Action)
		if !ok {
			action, ok := inputs[i+1].(policy.Action)
			if !ok {
				// Could be the string type.
				actionAsString, ok := inputs[i+1].(string)
				if !ok {
					panic(fmt.Sprintf("action '%T' not a supported action", inputs[i+1]))
				}
				action = policy.Action(actionAsString)
			}
			actions = []policy.Action{action}
		}

		out = append(out, AssertRBAC{
			Object:  rbacObj,
			Actions: actions,
		})
	}
	return out
}

type emptyPreparedAuthorized struct{}

func (emptyPreparedAuthorized) Authorize(_ context.Context, _ rbac.Object) error { return nil }
func (emptyPreparedAuthorized) CompileToSQL(_ context.Context, _ regosql.ConvertConfig) (string, error) {
	return "", nil
}