test: Add unit test to verify partial query support

Emyrk · Emyrk · commit 75e3a120b00d · 2022-08-09T23:02:23.000-05:00
diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go
@@ -33,7 +33,8 @@ func Filter[O Objecter](ctx context.Context, auth Authorizer, subjID string, sub
 
 // RegoAuthorizer will use a prepared rego query for performing authorize()
 type RegoAuthorizer struct {
-	query rego.PreparedEvalQuery
+	query   rego.PreparedEvalQuery
+	partial rego.PreparedPartialQuery
 }
 
 // Load the policy from policy.rego in this directory.
@@ -49,10 +50,19 @@ func NewAuthorizer() (*RegoAuthorizer, error) {
 		rego.Module("policy.rego", policy),
 	).PrepareForEval(ctx)
 
+	partial, err := rego.New(
+		rego.Query("allowed = data.authz.allow"),
+		rego.Module("policy.rego", policy),
+		rego.Unknowns([]string{
+			"input.object.owner",
+			"input.object.org_owner",
+		}),
+	).PrepareForPartial(ctx)
+
 	if err != nil {
 		return nil, xerrors.Errorf("prepare query: %w", err)
 	}
-	return &RegoAuthorizer{query: query}, nil
+	return &RegoAuthorizer{query: query, partial: partial}, nil
 }
 
 type authSubject struct {
@@ -107,3 +117,23 @@ func (a RegoAuthorizer) Authorize(ctx context.Context, subjectID string, roles [
 
 	return nil
 }
+
+// CheckPartial will not authorize the request. This function is to be used for unit testing to verify the rego policy
+// can be converted into ONLY queries. This ensures we can convert the queries into SQL WHERE clauses in the future.
+// If this function returns an error, then there is a set of inputs that also returns support rules, which cannot
+// be converted.
+// This function will not be used to actually perform authorization on partial queries.
+func (a RegoAuthorizer) CheckPartial(ctx context.Context, subjectID string, roles []Role, action Action, objectType string) (*rego.PartialQueries, interface{}, error) {
+	input := map[string]interface{}{
+		"subject": authSubject{
+			ID:    subjectID,
+			Roles: roles,
+		},
+		"object": map[string]string{
+			"type": objectType,
+		},
+		"action": action,
+	}
+	result, err := a.partial.Partial(ctx, rego.EvalInput(input))
+	return result, input, err
+}
diff --git a/coderd/rbac/authz_test.go b/coderd/rbac/authz_test.go
@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"testing"
 
+	"github.com/coder/coder/testutil"
+
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/xerrors"
@@ -570,7 +572,9 @@ func testAuthorize(t *testing.T, name string, subject subject, sets ...[]authTes
 		for _, c := range cases {
 			t.Run(name, func(t *testing.T) {
 				for _, a := range c.actions {
-					err := authorizer.Authorize(context.Background(), subject.UserID, subject.Roles, a, c.resource)
+					ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+					defer cancel()
+					err := authorizer.Authorize(ctx, subject.UserID, subject.Roles, a, c.resource)
 					if c.allow {
 						if err != nil {
 							var uerr *rbac.UnauthorizedError
@@ -593,6 +597,20 @@ func testAuthorize(t *testing.T, name string, subject subject, sets ...[]authTes
 						t.Log(string(d))
 					}
 					require.Error(t, err, "expected unauthorized")
+
+					// Also check the rego policy can form a valid partial query result.
+					result, input, err := authorizer.CheckPartial(ctx, subject.UserID, subject.Roles, a, c.resource.Type)
+					require.NoError(t, err, "check partial")
+					if len(result.Support) > 0 {
+						d, _ := json.Marshal(input)
+						t.Logf("input: %s", string(d))
+						for _, q := range result.Queries {
+							t.Logf("query: %+v", q.String())
+						}
+						for _, s := range result.Support {
+							t.Logf("support: %+v", s.String())
+						}
+					}
 				}
 			})
 		}
