Update tooling and testing to properly handle wildcards

1) Remove Simplify
2) Move membership package into internal and implement trackingSubjectSet there
3) Add a lot of additional testing for wildcards with both
4) Update consistency test suite to handle additional wildcard cases
5) Add additional developer tests
6) Disable ability to call Check or Lookup on subjects of type `*`, since it can be confusing

internal/graph/check.go

@@ -45,7 +45,9 @@ type ValidatedCheckRequest struct {
 func (cc *ConcurrentChecker) Check(ctx context.Context, req ValidatedCheckRequest, relation *v0.Relation) (*v1.DispatchCheckResponse, error) {
 	var directFunc ReduceableCheckFunc
 
-	if onrEqual(req.Subject, req.ObjectAndRelation) {
+	if req.Subject.ObjectId == tuple.PublicWildcard {
+		directFunc = checkError(NewErrInvalidArgument(errors.New("cannot perform check on wildcard")))
+	} else if onrEqual(req.Subject, req.ObjectAndRelation) {
 		// If we have found the goal's ONR, then we know that the ONR is a member.
 		directFunc = alwaysMember()
 	} else if relation.UsersetRewrite == nil {

internal/graph/errors.go

@@ -122,3 +122,15 @@ func NewRelationMissingTypeInfoErr(nsName string, relationName string) error {
 		relationName:  relationName,
 	}
 }
+
+// ErrInvalidArgument occurs when a request sent has an invalid argument.
+type ErrInvalidArgument struct {
+	error
+}
+
+// NewErrInvalidArgument constructs a request sent has an invalid argument.
+func NewErrInvalidArgument(baseErr error) error {
+	return ErrInvalidArgument{
+		error: baseErr,
+	}
+}

internal/graph/lookup.go

@@ -2,6 +2,7 @@ package graph
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	v0 "github.com/authzed/authzed-go/proto/authzed/api/v0"
@@ -46,6 +47,10 @@ const (
 // Lookup performs a lookup request with the provided request and context.
 func (cl *ConcurrentLookup) Lookup(ctx context.Context, req ValidatedLookupRequest) (*v1.DispatchLookupResponse, error) {
 	funcToResolve := cl.lookupInternal(ctx, req)
+	if req.Subject.ObjectId == tuple.PublicWildcard {
+		funcToResolve = returnResult(lookupResultError(req, NewErrInvalidArgument(errors.New("cannot perform lookup on wildcard")), emptyMetadata))
+	}
+
 	resolved := lookupOne(ctx, req, funcToResolve)
 
 	// Remove the resolved relation reference from the excluded direct list to mark that it was completely resolved.

internal/membership/foundsubject.go

@@ -0,0 +1,142 @@
+package membership
+
+import (
+	"fmt"
+	"sort"
+	"strings"
+
+	v0 "github.com/authzed/authzed-go/proto/authzed/api/v0"
+
+	"github.com/authzed/spicedb/pkg/tuple"
+)
+
+// NewFoundSubject creates a new FoundSubject for a subject and a set of its resources.
+func NewFoundSubject(subject *v0.ObjectAndRelation, resources ...*v0.ObjectAndRelation) FoundSubject {
+	return FoundSubject{subject, tuple.NewONRSet(), tuple.NewONRSet(resources...)}
+}
+
+// FoundSubject contains a single found subject and all the relationships in which that subject
+// is a member which were found via the ONRs expansion.
+type FoundSubject struct {
+	// subject is the subject found.
+	subject *v0.ObjectAndRelation
+
+	// excludedSubjects are any subjects excluded. Only should be set if subject is a wildcard.
+	excludedSubjects *tuple.ONRSet
+
+	// relations are the relations under which the subject lives that informed the locating
+	// of this subject for the root ONR.
+	relationships *tuple.ONRSet
+}
+
+// Subject returns the Subject of the FoundSubject.
+func (fs FoundSubject) Subject() *v0.ObjectAndRelation {
+	return fs.subject
+}
+
+// WildcardType returns the object type for the wildcard subject, if this is a wildcard subject.
+func (fs FoundSubject) WildcardType() (string, bool) {
+	if fs.subject.ObjectId == tuple.PublicWildcard {
+		return fs.subject.Namespace, true
+	}
+
+	return "", false
+}
+
+// ExcludedSubjectsFromWildcard returns those subjects excluded from the wildcard subject.
+// If not a wildcard subject, returns false.
+func (fs FoundSubject) ExcludedSubjectsFromWildcard() ([]*v0.ObjectAndRelation, bool) {
+	if fs.subject.ObjectId == tuple.PublicWildcard {
+		return fs.excludedSubjects.AsSlice(), true
+	}
+
+	return []*v0.ObjectAndRelation{}, false
+}
+
+// Relationships returns all the relationships in which the subject was found as per the expand.
+func (fs FoundSubject) Relationships() []*v0.ObjectAndRelation {
+	return fs.relationships.AsSlice()
+}
+
+// ToValidationString returns the FoundSubject in a format that is consumable by the validationfile
+// package.
+func (fs FoundSubject) ToValidationString() string {
+	onrString := tuple.StringONR(fs.Subject())
+	excluded, isWildcard := fs.ExcludedSubjectsFromWildcard()
+	if isWildcard && len(excluded) > 0 {
+		excludedONRStrings := make([]string, 0, len(excluded))
+		for _, excludedONR := range excluded {
+			excludedONRStrings = append(excludedONRStrings, tuple.StringONR(excludedONR))
+		}
+
+		sort.Strings(excludedONRStrings)
+		return fmt.Sprintf("%s - {%s}", onrString, strings.Join(excludedONRStrings, ", "))
+	}
+
+	return onrString
+}
+
+// union performs merging of two FoundSubject's with the same subject.
+func (fs FoundSubject) union(other FoundSubject) FoundSubject {
+	if toKey(fs.subject) != toKey(other.subject) {
+		panic("Got wrong found subject to union")
+	}
+
+	relationships := fs.relationships.Union(other.relationships)
+	var excludedSubjects *tuple.ONRSet
+
+	// If a wildcard, then union together excluded subjects.
+	_, isWildcard := fs.WildcardType()
+	if isWildcard {
+		excludedSubjects = fs.excludedSubjects.Union(other.excludedSubjects)
+	}
+
+	return FoundSubject{
+		subject:          fs.subject,
+		excludedSubjects: excludedSubjects,
+		relationships:    relationships,
+	}
+}
+
+// intersect performs intersection between two FoundSubject's with the same subject.
+func (fs FoundSubject) intersect(other FoundSubject) FoundSubject {
+	if toKey(fs.subject) != toKey(other.subject) {
+		panic("Got wrong found subject to intersect")
+	}
+
+	relationships := fs.relationships.Union(other.relationships)
+	var excludedSubjects *tuple.ONRSet
+
+	// If a wildcard, then union together excluded subjects.
+	_, isWildcard := fs.WildcardType()
+	if isWildcard {
+		excludedSubjects = fs.excludedSubjects.Union(other.excludedSubjects)
+	}
+
+	return FoundSubject{
+		subject:          fs.subject,
+		excludedSubjects: excludedSubjects,
+		relationships:    relationships,
+	}
+}
+
+// FoundSubjects contains the subjects found for a specific ONR.
+type FoundSubjects struct {
+	// subjects is a map from the Subject ONR (as a string) to the FoundSubject information.
+	subjects map[string]FoundSubject
+}
+
+// ListFound returns a slice of all the FoundSubject's.
+func (fs FoundSubjects) ListFound() []FoundSubject {
+	found := []FoundSubject{}
+	for _, sub := range fs.subjects {
+		found = append(found, sub)
+	}
+	return found
+}
+
+// LookupSubject returns the FoundSubject for a matching subject, if any.
+func (fs FoundSubjects) LookupSubject(subject *v0.ObjectAndRelation) (FoundSubject, bool) {
+	found, ok := fs.subjects[toKey(subject)]
+	return found, ok
+}

internal/membership/foundsubject_test.go

@@ -0,0 +1,62 @@
+package membership
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/authzed/spicedb/pkg/validationfile"
+	"github.com/stretchr/testify/require"
+)
+
+func TestToValidationString(t *testing.T) {
+	testCases := []struct {
+		name     string
+		fs       FoundSubject
+		expected string
+	}{
+		{
+			"basic",
+			fs("user", "user1", "..."),
+			"user:user1",
+		},
+		{
+			"with exclusion",
+			fs("user", "*", "...", ONR("user", "user1", "...")),
+			"user:* - {user:user1}",
+		},
+		{
+			"with some exclusion",
+			fs("user", "*", "...",
+				ONR("user", "user1", "..."),
+				ONR("user", "user2", "..."),
+				ONR("user", "user3", "..."),
+				ONR("user", "user4", "..."),
+				ONR("user", "user5", "..."),
+			),
+			"user:* - {user:user1, user:user2, user:user3, user:user4, user:user5}",
+		},
+		{
+			"with many exclusion",
+			fs("user", "*", "...",
+				ONR("user", "user1", "..."),
+				ONR("user", "user2", "..."),
+				ONR("user", "user3", "..."),
+				ONR("user", "user4", "..."),
+				ONR("user", "user5", "..."),
+				ONR("user", "user6", "..."),
+			),
+			"user:* - {user:user1, user:user2, user:user3, user:user4, user:user5, user:user6}",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require := require.New(t)
+			require.Equal(tc.expected, tc.fs.ToValidationString())
+
+			sub, err := validationfile.ValidationString(fmt.Sprintf("[%s]", tc.expected)).Subject()
+			require.Nil(err)
+			require.NotNil(sub)
+		})
+	}
+}

internal/membership/membership.go

@@ -0,0 +1,140 @@
+package membership
+
+import (
+	"fmt"
+
+	v0 "github.com/authzed/authzed-go/proto/authzed/api/v0"
+
+	"github.com/authzed/spicedb/pkg/tuple"
+)
+
+// Set represents the set of membership for one or more ONRs, based on expansion
+// trees.
+type Set struct {
+	// objectsAndRelations is a map from an ONR (as a string) to the subjects found for that ONR.
+	objectsAndRelations map[string]FoundSubjects
+}
+
+// SubjectsByONR returns a map from ONR (as a string) to the FoundSubjects for that ONR.
+func (ms *Set) SubjectsByONR() map[string]FoundSubjects {
+	return ms.objectsAndRelations
+}
+
+// NewMembershipSet constructs a new membership set.
+//
+// NOTE: This is designed solely for the developer API and should *not* be used in any performance
+// sensitive code.
+func NewMembershipSet() *Set {
+	return &Set{
+		objectsAndRelations: map[string]FoundSubjects{},
+	}
+}
+
+// AddExpansion adds the expansion of an ONR to the membership set. Returns false if the ONR was already added.
+//
+// NOTE: The expansion tree *should* be the fully recursive expansion.
+func (ms *Set) AddExpansion(onr *v0.ObjectAndRelation, expansion *v0.RelationTupleTreeNode) (FoundSubjects, bool, error) {
+	onrString := tuple.StringONR(onr)
+	existing, ok := ms.objectsAndRelations[onrString]
+	if ok {
+		return existing, false, nil
+	}
+
+	tss, err := populateFoundSubjects(onr, expansion)
+	if err != nil {
+		return FoundSubjects{}, false, err
+	}
+
+	fs := tss.ToFoundSubjects()
+	ms.objectsAndRelations[onrString] = fs
+	return fs, true, nil
+}
+
+// AccessibleExpansionSubjects returns a TrackingSubjectSet representing the set of accessible subjects in the expansion.
+func AccessibleExpansionSubjects(treeNode *v0.RelationTupleTreeNode) (TrackingSubjectSet, error) {
+	return populateFoundSubjects(treeNode.Expanded, treeNode)
+}
+
+func populateFoundSubjects(rootONR *v0.ObjectAndRelation, treeNode *v0.RelationTupleTreeNode) (TrackingSubjectSet, error) {
+	resource := rootONR
+	if treeNode.Expanded != nil {
+		resource = treeNode.Expanded
+	}
+
+	switch typed := treeNode.NodeType.(type) {
+	case *v0.RelationTupleTreeNode_IntermediateNode:
+		switch typed.IntermediateNode.Operation {
+		case v0.SetOperationUserset_UNION:
+			toReturn := NewTrackingSubjectSet()
+			for _, child := range typed.IntermediateNode.ChildNodes {
+				tss, err := populateFoundSubjects(resource, child)
+				if err != nil {
+					return nil, err
+				}
+
+				toReturn.AddFrom(tss)
+			}
+			return toReturn, nil
+
+		case v0.SetOperationUserset_INTERSECTION:
+			if len(typed.IntermediateNode.ChildNodes) == 0 {
+				return nil, fmt.Errorf("Found intersection with no children")
+			}
+
+			firstChildSet, err := populateFoundSubjects(rootONR, typed.IntermediateNode.ChildNodes[0])
+			if err != nil {
+				return nil, err
+			}
+
+			toReturn := NewTrackingSubjectSet()
+			toReturn.AddFrom(firstChildSet)
+
+			for _, child := range typed.IntermediateNode.ChildNodes[1:] {
+				childSet, err := populateFoundSubjects(rootONR, child)
+				if err != nil {
+					return nil, err
+				}
+				toReturn = toReturn.Intersect(childSet)
+			}
+			return toReturn, nil
+
+		case v0.SetOperationUserset_EXCLUSION:
+			if len(typed.IntermediateNode.ChildNodes) == 0 {
+				return nil, fmt.Errorf("Found exclusion with no children")
+			}
+
+			firstChildSet, err := populateFoundSubjects(rootONR, typed.IntermediateNode.ChildNodes[0])
+			if err != nil {
+				return nil, err
+			}
+
+			toReturn := NewTrackingSubjectSet()
+			toReturn.AddFrom(firstChildSet)
+
+			for _, child := range typed.IntermediateNode.ChildNodes[1:] {
+				childSet, err := populateFoundSubjects(rootONR, child)
+				if err != nil {
+					return nil, err
+				}
+				toReturn = toReturn.Exclude(childSet)
+			}
+
+			return toReturn, nil
+
+		default:
+			panic("unknown expand operation")
+		}
+
+	case *v0.RelationTupleTreeNode_LeafNode:
+		toReturn := NewTrackingSubjectSet()
+		for _, user := range typed.LeafNode.Users {
+			fs := NewFoundSubject(user.GetUserset())
+			toReturn.Add(fs)
+			fs.relationships.Add(resource)
+		}
+		return toReturn, nil
+
+	default:
+		panic("unknown TreeNode type")
+	}
+}