test: Handle Filter flake with ctx errors (#7119)

Emyrk · web-flow · commit 2137db044585 · 2023-04-14T12:30:35.000-05:00
* test: Handle Fitler flake with ctx errors
* Add unit test to check filter for proper error
* Correctly return category of errors
diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go
@@ -137,6 +137,10 @@ func Filter[O Objecter](ctx context.Context, auth Authorizer, subject Subject, a
 			err := auth.Authorize(ctx, subject, action, o.RBACObject())
 			if err == nil {
 				filtered = append(filtered, o)
+			} else if !IsUnauthorizedError(err) {
+				// If the error is not the expected "Unauthorized" error, then
+				// it is something unexpected.
+				return nil, err
 			}
 		}
 		return filtered, nil
@@ -155,6 +159,10 @@ func Filter[O Objecter](ctx context.Context, auth Authorizer, subject Subject, a
 		err := prepared.Authorize(ctx, rbacObj)
 		if err == nil {
 			filtered = append(filtered, object)
+		} else if !IsUnauthorizedError(err) {
+			// If the error is not the expected "Unauthorized" error, then
+			// it is something unexpected.
+			return nil, err
 		}
 	}
 
@@ -319,7 +327,8 @@ func (a RegoAuthorizer) authorize(ctx context.Context, subject Subject, action A
 
 	results, err := a.query.Eval(ctx, rego.EvalParsedInput(astV))
 	if err != nil {
-		return ForbiddenWithInternal(xerrors.Errorf("eval rego: %w", err), subject, action, object, results)
+		err = correctCancelError(err)
+		return xerrors.Errorf("evaluate rego: %w", err)
 	}
 
 	if !results.Allowed() {
@@ -430,7 +439,8 @@ EachQueryLoop:
 		// We need to eval each query with the newly known fields.
 		results, err := q.Eval(ctx, rego.EvalParsedInput(parsed))
 		if err != nil {
-			continue EachQueryLoop
+			err = correctCancelError(err)
+			return xerrors.Errorf("eval error: %w", err)
 		}
 
 		// If there are no results, then the query is false. This is because rego
diff --git a/coderd/rbac/authz_internal_test.go b/coderd/rbac/authz_internal_test.go
@@ -30,18 +30,68 @@ func (w fakeObject) RBACObject() Object {
 	}
 }
 
+// objectBomb is a wrapper around an Objecter that calls a function when
+// RBACObject is called.
+type objectBomb struct {
+	Objecter
+	bomb func()
+}
+
+func (o *objectBomb) RBACObject() Object {
+	o.bomb()
+	return o.Objecter.RBACObject()
+}
+
 func TestFilterError(t *testing.T) {
 	t.Parallel()
-	auth := NewAuthorizer(prometheus.NewRegistry())
-	subject := Subject{
-		ID:     uuid.NewString(),
-		Roles:  RoleNames{},
-		Groups: []string{},
-		Scope:  ScopeAll,
-	}
+	_ = objectBomb{}
 
-	_, err := Filter(context.Background(), auth, subject, ActionRead, []Object{ResourceUser, ResourceWorkspace})
-	require.ErrorContains(t, err, "object types must be uniform")
+	t.Run("DifferentResourceTypes", func(t *testing.T) {
+		t.Parallel()
+
+		auth := NewAuthorizer(prometheus.NewRegistry())
+		subject := Subject{
+			ID:     uuid.NewString(),
+			Roles:  RoleNames{},
+			Groups: []string{},
+			Scope:  ScopeAll,
+		}
+
+		_, err := Filter(context.Background(), auth, subject, ActionRead, []Object{ResourceUser, ResourceWorkspace})
+		require.ErrorContains(t, err, "object types must be uniform")
+	})
+
+	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())
+		subject := Subject{
+			ID: uuid.NewString(),
+			Roles: RoleNames{
+				RoleOwner(),
+			},
+			Groups: []string{},
+			Scope:  ScopeAll,
+		}
+
+		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)
+	})
 }
 
 // TestFilter ensures the filter acts the same as an individual authorize.
@@ -170,7 +220,7 @@ func TestFilter(t *testing.T) {
 			localObjects := make([]fakeObject, len(objects))
 			copy(localObjects, objects)
 
-			ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+			ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
 			defer cancel()
 			auth := NewAuthorizer(prometheus.NewRegistry())
 
diff --git a/coderd/rbac/error.go b/coderd/rbac/error.go
@@ -1,11 +1,14 @@
 package rbac
 
 import (
+	"context"
 	"errors"
 	"flag"
 	"fmt"
 
 	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/topdown"
+	"golang.org/x/xerrors"
 )
 
 const (
@@ -97,3 +100,17 @@ func (*UnauthorizedError) As(target interface{}) bool {
 	}
 	return false
 }
+
+// correctCancelError will return the correct error for a canceled context. This
+// is because rego changes a canceled context to a topdown.CancelErr. This error
+// is not helpful if the code is "canceled". To make the error conform with the
+// rest of our canceled errors, we will convert the error to a context.Canceled
+// error. No good information is lost, as the topdown.CancelErr provides the
+// location of the query that was canceled, which does not matter.
+func correctCancelError(err error) error {
+	e := new(topdown.Error)
+	if xerrors.As(err, &e) || e.Code == topdown.CancelErr {
+		return context.Canceled
+	}
+	return err
+}
