Update BaseSubjectSet to support caveat expressions

This is the first (massive) step in supporting caveats in LookupSubjects, as this implements all the bookkeeping and tracking associated with each subject added to the subject set

First part of authzed#931

internal/dispatch/graph/lookupsubjects_test.go

@@ -142,7 +142,7 @@ func TestSimpleLookupSubjects(t *testing.T) {
 			foundSubjectIds := []string{}
 			for _, result := range stream.Results() {
 				for _, found := range result.FoundSubjects {
-					if len(found.ExcludedSubjectIds) > 0 {
+					if len(found.ExcludedSubjects) > 0 {
 						continue
 					}
 

internal/membership/foundsubject.go

@@ -6,13 +6,14 @@ import (
 	"strings"
 
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
+	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
 
 	"github.com/authzed/spicedb/pkg/tuple"
 )
 
 // NewFoundSubject creates a new FoundSubject for a subject and a set of its resources.
 func NewFoundSubject(subject *core.ObjectAndRelation, resources ...*core.ObjectAndRelation) FoundSubject {
-	return FoundSubject{subject, nil, tuple.NewONRSet(resources...)}
+	return FoundSubject{subject, nil, nil, tuple.NewONRSet(resources...)}
 }
 
 // FoundSubject contains a single found subject and all the relationships in which that subject
@@ -22,7 +23,10 @@ type FoundSubject struct {
 	subject *core.ObjectAndRelation
 
 	// excludedSubjects are any subjects excluded. Only should be set if subject is a wildcard.
-	excludedSubjectIds []string
+	excludedSubjects []FoundSubject
+
+	// caveatExpression is the conditional expression on the found subject.
+	caveatExpression *v1.CaveatExpression
 
 	// relations are the relations under which the subject lives that informed the locating
 	// of this subject for the root ONR.
@@ -36,8 +40,12 @@ func (fs FoundSubject) GetSubjectId() string {
 	return fs.subject.ObjectId
 }
 
-func (fs FoundSubject) GetExcludedSubjectIds() []string {
-	return fs.excludedSubjectIds
+func (fs FoundSubject) GetCaveatExpression() *v1.CaveatExpression {
+	return fs.caveatExpression
+}
+
+func (fs FoundSubject) GetExcludedSubjects() []FoundSubject {
+	return fs.excludedSubjects
 }
 
 // Subject returns the Subject of the FoundSubject.
@@ -58,13 +66,14 @@ func (fs FoundSubject) WildcardType() (string, bool) {
 // If not a wildcard subject, returns false.
 func (fs FoundSubject) ExcludedSubjectsFromWildcard() ([]*core.ObjectAndRelation, bool) {
 	if fs.subject.ObjectId == tuple.PublicWildcard {
-		excludedSubjects := make([]*core.ObjectAndRelation, 0, len(fs.excludedSubjectIds))
-		for _, excludedID := range fs.excludedSubjectIds {
-			excludedSubjects = append(excludedSubjects, &core.ObjectAndRelation{
-				Namespace: fs.subject.Namespace,
-				ObjectId:  excludedID,
-				Relation:  fs.subject.Relation,
-			})
+		excludedSubjects := make([]*core.ObjectAndRelation, 0, len(fs.excludedSubjects))
+		for _, excludedSubject := range fs.excludedSubjects {
+			// TODO(jschorr): Fix once we add caveats support to debug tooling
+			if excludedSubject.caveatExpression != nil {
+				panic("not yet supported")
+			}
+
+			excludedSubjects = append(excludedSubjects, excludedSubject.subject)
 		}
 
 		return excludedSubjects, true
@@ -73,6 +82,20 @@ func (fs FoundSubject) ExcludedSubjectsFromWildcard() ([]*core.ObjectAndRelation
 	return []*core.ObjectAndRelation{}, false
 }
 
+func (fs FoundSubject) excludedSubjectIDs() []string {
+	excludedSubjects := make([]string, 0, len(fs.excludedSubjects))
+	for _, excludedSubject := range fs.excludedSubjects {
+		// TODO(jschorr): Fix once we add caveats support to debug tooling
+		if excludedSubject.caveatExpression != nil {
+			panic("not yet supported")
+		}
+
+		excludedSubjects = append(excludedSubjects, excludedSubject.subject.ObjectId)
+	}
+
+	return excludedSubjects
+}
+
 // Relationships returns all the relationships in which the subject was found as per the expand.
 func (fs FoundSubject) Relationships() []*core.ObjectAndRelation {
 	return fs.relationships.AsSlice()
@@ -81,6 +104,11 @@ func (fs FoundSubject) Relationships() []*core.ObjectAndRelation {
 // ToValidationString returns the FoundSubject in a format that is consumable by the validationfile
 // package.
 func (fs FoundSubject) ToValidationString() string {
+	if fs.caveatExpression != nil {
+		// TODO(jschorr): Implement once we have a format for this.
+		panic("conditional found subjects not yet supported")
+	}
+
 	onrString := tuple.StringONR(fs.Subject())
 	excluded, isWildcard := fs.ExcludedSubjectsFromWildcard()
 	if isWildcard && len(excluded) > 0 {

internal/membership/trackingsubjectset.go

@@ -4,7 +4,10 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/authzed/spicedb/pkg/tuple"
+
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
+	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
 
 	"github.com/authzed/spicedb/internal/util"
 )
@@ -62,24 +65,21 @@ func (tss *TrackingSubjectSet) getSetForKey(key string) util.BaseSubjectSet[Foun
 	parts := strings.Split(key, "#")
 
 	created := util.NewBaseSubjectSet[FoundSubject](
-		func(subjectID string, excludedSubjectIDs []string, sources ...FoundSubject) FoundSubject {
+		func(subjectID string, caveatExpression *v1.CaveatExpression, excludedSubjects []FoundSubject, sources ...FoundSubject) FoundSubject {
 			fs := NewFoundSubject(&core.ObjectAndRelation{
 				Namespace: parts[0],
 				ObjectId:  subjectID,
 				Relation:  parts[1],
 			})
-			fs.excludedSubjectIds = excludedSubjectIDs
+			fs.excludedSubjects = excludedSubjects
+			fs.caveatExpression = caveatExpression
 			for _, source := range sources {
-				fs.relationships.UpdateFrom(source.relationships)
+				if source.relationships != nil {
+					fs.relationships.UpdateFrom(source.relationships)
+				}
 			}
 			return fs
 		},
-		func(existing FoundSubject, added FoundSubject) FoundSubject {
-			fs := NewFoundSubject(existing.subject)
-			fs.excludedSubjectIds = existing.excludedSubjectIds
-			fs.relationships = existing.relationships.Union(added.relationships)
-			return fs
-		},
 	)
 	tss.setByType[key] = created
 	return created
@@ -151,6 +151,16 @@ func (tss TrackingSubjectSet) removeExact(subjects ...*core.ObjectAndRelation) {
 	}
 }
 
+func (tss TrackingSubjectSet) getSubjects() []string {
+	var subjects []string
+	for _, subjectSet := range tss.setByType {
+		for _, foundSubject := range subjectSet.AsSlice() {
+			subjects = append(subjects, tuple.StringONR(foundSubject.subject))
+		}
+	}
+	return subjects
+}
+
 // ToSlice returns a slice of all subjects found in the set.
 func (tss TrackingSubjectSet) ToSlice() []FoundSubject {
 	subjects := []FoundSubject{}

internal/membership/trackingsubjectset_test.go

@@ -45,10 +45,15 @@ func subtract(firstSet *TrackingSubjectSet, sets ...*TrackingSubjectSet) *Tracki
 }
 
 func fs(subjectType string, subjectID string, subjectRel string, excludedSubjectIDs ...string) FoundSubject {
+	excludedSubjects := make([]FoundSubject, 0, len(excludedSubjectIDs))
+	for _, excludedSubjectID := range excludedSubjectIDs {
+		excludedSubjects = append(excludedSubjects, FoundSubject{subject: ONR(subjectType, excludedSubjectID, subjectRel)})
+	}
+
 	return FoundSubject{
-		subject:            ONR(subjectType, subjectID, subjectRel),
-		excludedSubjectIds: excludedSubjectIDs,
-		relationships:      tuple.NewONRSet(),
+		subject:          ONR(subjectType, subjectID, subjectRel),
+		excludedSubjects: excludedSubjects,
+		relationships:    tuple.NewONRSet(),
 	}
 }
 
@@ -330,8 +335,8 @@ func TestTrackingSubjectSet(t *testing.T) {
 					found, ok := tc.set.Get(fs.subject)
 					require.True(ok, "missing expected subject %s", fs.subject)
 
-					expectedExcluded := util.NewSet[string](fs.excludedSubjectIds...)
-					foundExcluded := util.NewSet[string](found.excludedSubjectIds...)
+					expectedExcluded := util.NewSet[string](fs.excludedSubjectIDs()...)
+					foundExcluded := util.NewSet[string](found.excludedSubjectIDs()...)
 					require.Len(expectedExcluded.Subtract(foundExcluded).AsSlice(), 0, "mismatch on excluded subjects on %s: expected: %s, found: %s", fs.subject, expectedExcluded, foundExcluded)
 					require.Len(foundExcluded.Subtract(expectedExcluded).AsSlice(), 0, "mismatch on excluded subjects on %s: expected: %s, found: %s", fs.subject, expectedExcluded, foundExcluded)
 				} else {
@@ -340,7 +345,7 @@ func TestTrackingSubjectSet(t *testing.T) {
 				tc.set.removeExact(fs.subject)
 			}
 
-			require.True(tc.set.IsEmpty())
+			require.True(tc.set.IsEmpty(), "Found remaining: %v", tc.set.getSubjects())
 		})
 	}
 }

internal/services/v1/permissions.go

@@ -430,9 +430,14 @@ func (ps *permissionServer) LookupSubjects(req *v1.LookupSubjectsRequest, resp v
 
 	stream := dispatchpkg.NewHandlingDispatchStream(ctx, func(result *dispatch.DispatchLookupSubjectsResponse) error {
 		for _, foundSubject := range result.FoundSubjects {
+			excludedSubjectIDs := make([]string, 0, len(foundSubject.ExcludedSubjects))
+			for _, excludedSubject := range foundSubject.ExcludedSubjects {
+				excludedSubjectIDs = append(excludedSubjectIDs, excludedSubject.SubjectId)
+			}
+
 			err := resp.Send(&v1.LookupSubjectsResponse{
 				SubjectObjectId:    foundSubject.SubjectId,
-				ExcludedSubjectIds: foundSubject.ExcludedSubjectIds,
+				ExcludedSubjectIds: excludedSubjectIDs,
 				LookedUpAt:         revisionReadAt,
 			})
 			if err != nil {