fix: improve wildcard support when doing IterSubjects (authzed#2843)

barakmich · web-flow · commit 9c4c183fa92d · 2026-01-23T19:51:30.000Z
diff --git a/pkg/query/datastore.go b/pkg/query/datastore.go
@@ -177,11 +177,23 @@ func (r *RelationIterator) iterSubjectsNormalImpl(ctx *Context, resource Object)
 		return nil, err
 	}
 
-	return convertRelationSeqToPathSeq(iter.Seq2[tuple.Relationship, error](relIter)), nil
+	// Filter out wildcard subjects to match the behavior of LookupSubjects. Wildcards are not
+	// concrete enumerable subjects. They will be expanded to concrete subjects by the wildcard branch.
+	return FilterWildcardSubjects(convertRelationSeqToPathSeq(iter.Seq2[tuple.Relationship, error](relIter))), nil
 }
 
 func (r *RelationIterator) iterSubjectsWildcardImpl(ctx *Context, resource Object) (PathSeq, error) {
-	filter := datastore.RelationshipsFilter{
+	// When a relation contains a wildcard (e.g., user:*), it means "all subjects of that type"
+	// that have ANY relationship with this resource. We enumerate concrete subjects by:
+	// 1. First checking if a wildcard relationship actually exists for this resource
+	// 2. If yes, querying for all concrete subjects with relationships to this resource
+	//
+	// This avoids doing a full subject enumeration when no wildcard exists (the common case).
+	// When wildcards do exist, we do 2 queries in this branch, but that's the correct semantic
+	// behavior - we only enumerate when there's actually a wildcard to expand.
+
+	// First, check if there's actually a wildcard relationship for this resource
+	wildcardFilter := datastore.RelationshipsFilter{
 		OptionalResourceType:     r.base.DefinitionName(),
 		OptionalResourceIds:      []string{resource.ObjectID},
 		OptionalResourceRelation: r.base.RelationName(),
@@ -194,16 +206,57 @@ func (r *RelationIterator) iterSubjectsWildcardImpl(ctx *Context, resource Objec
 		},
 	}
 
-	relIter, err := ctx.Reader.QueryRelationships(ctx, filter,
+	wildcardIter, err := ctx.Reader.QueryRelationships(ctx, wildcardFilter,
 		options.WithSkipCaveats(r.base.Caveat() == ""),
 		options.WithSkipExpiration(!r.base.Expiration()),
 		options.WithQueryShape(queryshape.AllSubjectsForResources),
+		options.WithLimit(options.LimitOne), // We only need to know if one exists
 	)
 	if err != nil {
 		return nil, err
 	}
 
-	return convertRelationSeqToPathSeq(iter.Seq2[tuple.Relationship, error](relIter)), nil
+	// Check if any wildcard relationship exists
+	hasWildcard := false
+	for _, err := range wildcardIter {
+		if err != nil {
+			return nil, err
+		}
+		hasWildcard = true
+		break
+	}
+
+	// If no wildcard relationship exists, return empty - nothing to enumerate
+	if !hasWildcard {
+		return EmptyPathSeq(), nil
+	}
+
+	// Wildcard exists, so enumerate all concrete subjects for this resource.
+	// Note: This may return some of the same subjects as the non-wildcard branch
+	// (when both wildcard and concrete relationships exist), but the Union will
+	// deduplicate them.
+	allSubjectsFilter := datastore.RelationshipsFilter{
+		OptionalResourceType: r.base.DefinitionName(),
+		OptionalResourceIds:  []string{resource.ObjectID},
+		OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+			{
+				OptionalSubjectType: r.base.Type(),
+				RelationFilter:      r.buildSubjectRelationFilter(),
+			},
+		},
+	}
+
+	relIter, err := ctx.Reader.QueryRelationships(ctx, allSubjectsFilter,
+		options.WithSkipCaveats(r.base.Caveat() == ""),
+		options.WithSkipExpiration(!r.base.Expiration()),
+		options.WithQueryShape(queryshape.AllSubjectsForResources),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Filter out wildcard subjects from the results - we only want concrete subjects
+	return FilterWildcardSubjects(convertRelationSeqToPathSeq(iter.Seq2[tuple.Relationship, error](relIter))), nil
 }
 
 func (r *RelationIterator) IterResourcesImpl(ctx *Context, subject ObjectAndRelation) (PathSeq, error) {
diff --git a/pkg/query/path.go b/pkg/query/path.go
@@ -383,3 +383,26 @@ func RewriteSubject(seq PathSeq, subject ObjectAndRelation) PathSeq {
 		}
 	}
 }
+
+// FilterWildcardSubjects filters out any paths with wildcard subjects.
+func FilterWildcardSubjects(seq PathSeq) PathSeq {
+	return func(yield func(Path, error) bool) {
+		for path, err := range seq {
+			if err != nil {
+				if !yield(Path{}, err) {
+					return
+				}
+				continue
+			}
+
+			// Skip wildcard subjects
+			if path.Subject.ObjectID == tuple.PublicWildcard {
+				continue
+			}
+
+			if !yield(path, nil) {
+				return
+			}
+		}
+	}
+}
diff --git a/pkg/query/wildcard_subjects_test.go b/pkg/query/wildcard_subjects_test.go
@@ -0,0 +1,209 @@
+package query
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/internal/datastore/common"
+	"github.com/authzed/spicedb/internal/datastore/dsfortesting"
+	"github.com/authzed/spicedb/internal/datastore/memdb"
+	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/genutil/slicez"
+	"github.com/authzed/spicedb/pkg/schema/v2"
+	"github.com/authzed/spicedb/pkg/schemadsl/compiler"
+	"github.com/authzed/spicedb/pkg/schemadsl/input"
+	"github.com/authzed/spicedb/pkg/tuple"
+)
+
+// TestIterSubjectsWithWildcard tests that wildcards are properly filtered and expanded
+func TestIterSubjectsWithWildcard(t *testing.T) {
+	t.Parallel()
+
+	require := require.New(t)
+	rawDS, err := dsfortesting.NewMemDBDatastoreForTesting(t, 0, 0, memdb.DisableGC)
+	require.NoError(err)
+
+	// Create a schema with wildcards: relation viewer: user | user:*
+	schemaText := `
+		definition user {}
+
+		definition resource {
+			relation viewer: user | user:*
+		}
+	`
+
+	ctx := context.Background()
+
+	// Compile the schema
+	compiled, err := compiler.Compile(compiler.InputSchema{
+		Source:       input.Source("test"),
+		SchemaString: schemaText,
+	}, compiler.AllowUnprefixedObjectType())
+	require.NoError(err)
+
+	// Write the schema
+	_, err = rawDS.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		return rwt.LegacyWriteNamespaces(ctx, compiled.ObjectDefinitions...)
+	})
+	require.NoError(err)
+
+	// Write test data:
+	// - resource:first#viewer@user:* (wildcard)
+	// - resource:first#viewer@user:concrete (concrete user)
+	revision, err := common.WriteRelationships(ctx, rawDS, tuple.UpdateOperationCreate,
+		tuple.MustParse("resource:first#viewer@user:*"),
+		tuple.MustParse("resource:first#viewer@user:concrete"),
+	)
+	require.NoError(err)
+
+	// Build schema for querying
+	dsSchema, err := schema.BuildSchemaFromDefinitions(compiled.ObjectDefinitions, nil)
+	require.NoError(err)
+
+	resourceDef, _ := dsSchema.GetTypeDefinition("resource")
+	viewerRel, _ := resourceDef.GetRelation("viewer")
+
+	// Test the non-wildcard branch (user)
+	t.Run("NonWildcardBranch", func(t *testing.T) {
+		t.Parallel()
+		// The non-wildcard branch should only return concrete subjects, filtering out wildcards
+		nonWildcardBranch := NewRelationIterator(viewerRel.BaseRelations()[0]) // user (non-wildcard)
+
+		queryCtx := NewLocalContext(ctx, WithReader(rawDS.SnapshotReader(revision)))
+		subjects, err := queryCtx.IterSubjects(nonWildcardBranch, NewObject("resource", "first"))
+		require.NoError(err)
+
+		paths, err := CollectAll(subjects)
+		require.NoError(err)
+
+		// Should only get the concrete user, not the wildcard
+		require.Len(paths, 1)
+		require.Equal("user", paths[0].Subject.ObjectType)
+		require.Equal("concrete", paths[0].Subject.ObjectID)
+	})
+
+	// Test the wildcard branch (user:*)
+	t.Run("WildcardBranch", func(t *testing.T) {
+		t.Parallel()
+		// The wildcard branch should enumerate concrete subjects when a wildcard exists
+		wildcardBranch := NewRelationIterator(viewerRel.BaseRelations()[1]) // user:* (wildcard)
+
+		queryCtx := NewLocalContext(ctx, WithReader(rawDS.SnapshotReader(revision)))
+		subjects, err := queryCtx.IterSubjects(wildcardBranch, NewObject("resource", "first"))
+		require.NoError(err)
+
+		paths, err := CollectAll(subjects)
+		require.NoError(err)
+
+		// Should only get the concrete user, not the wildcard itself
+		require.Len(paths, 1)
+		require.Equal("user", paths[0].Subject.ObjectType)
+		require.Equal("concrete", paths[0].Subject.ObjectID)
+	})
+
+	// Test the Union (combined behavior)
+	t.Run("UnionDeduplication", func(t *testing.T) {
+		t.Parallel()
+		// The Union of both branches should deduplicate the concrete user
+		union := NewUnion()
+		union.addSubIterator(NewRelationIterator(viewerRel.BaseRelations()[0])) // user
+		union.addSubIterator(NewRelationIterator(viewerRel.BaseRelations()[1])) // user:*
+
+		queryCtx := NewLocalContext(ctx, WithReader(rawDS.SnapshotReader(revision)))
+		subjects, err := queryCtx.IterSubjects(union, NewObject("resource", "first"))
+		require.NoError(err)
+
+		paths, err := CollectAll(subjects)
+		require.NoError(err)
+
+		// Should get exactly one concrete user (deduplicated)
+		require.Len(paths, 1)
+		require.Equal("user", paths[0].Subject.ObjectType)
+		require.Equal("concrete", paths[0].Subject.ObjectID)
+	})
+}
+
+// TestIterSubjectsWildcardWithoutWildcardRelationship tests that the wildcard branch
+// returns empty when no wildcard relationship exists
+func TestIterSubjectsWildcardWithoutWildcardRelationship(t *testing.T) {
+	t.Parallel()
+
+	require := require.New(t)
+	rawDS, err := dsfortesting.NewMemDBDatastoreForTesting(t, 0, 0, memdb.DisableGC)
+	require.NoError(err)
+
+	// Create a schema with wildcards: relation viewer: user | user:*
+	schemaText := `
+		definition user {}
+
+		definition resource {
+			relation viewer: user | user:*
+		}
+	`
+
+	ctx := context.Background()
+
+	// Compile the schema
+	compiled, err := compiler.Compile(compiler.InputSchema{
+		Source:       input.Source("test"),
+		SchemaString: schemaText,
+	}, compiler.AllowUnprefixedObjectType())
+	require.NoError(err)
+
+	// Write the schema
+	_, err = rawDS.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		return rwt.LegacyWriteNamespaces(ctx, compiled.ObjectDefinitions...)
+	})
+	require.NoError(err)
+
+	// Write test data with ONLY concrete users, NO wildcard
+	revision, err := common.WriteRelationships(ctx, rawDS, tuple.UpdateOperationCreate,
+		tuple.MustParse("resource:second#viewer@user:alice"),
+		tuple.MustParse("resource:second#viewer@user:bob"),
+	)
+	require.NoError(err)
+
+	// Build schema for querying
+	dsSchema, err := schema.BuildSchemaFromDefinitions(compiled.ObjectDefinitions, nil)
+	require.NoError(err)
+
+	resourceDef, _ := dsSchema.GetTypeDefinition("resource")
+	viewerRel, _ := resourceDef.GetRelation("viewer")
+
+	// Test the wildcard branch when no wildcard relationship exists
+	t.Run("WildcardBranchWithoutWildcard", func(t *testing.T) {
+		t.Parallel()
+		// The wildcard branch should return empty because there's no wildcard relationship
+		wildcardBranch := NewRelationIterator(viewerRel.BaseRelations()[1]) // user:* (wildcard)
+
+		queryCtx := NewLocalContext(ctx, WithReader(rawDS.SnapshotReader(revision)))
+		subjects, err := queryCtx.IterSubjects(wildcardBranch, NewObject("resource", "second"))
+		require.NoError(err)
+
+		paths, err := CollectAll(subjects)
+		require.NoError(err)
+
+		// Should be empty - no wildcard relationship exists
+		require.Empty(paths)
+	})
+
+	// Test the non-wildcard branch still works
+	t.Run("NonWildcardBranchWorksNormally", func(t *testing.T) {
+		t.Parallel()
+		nonWildcardBranch := NewRelationIterator(viewerRel.BaseRelations()[0]) // user (non-wildcard)
+
+		queryCtx := NewLocalContext(ctx, WithReader(rawDS.SnapshotReader(revision)))
+		subjects, err := queryCtx.IterSubjects(nonWildcardBranch, NewObject("resource", "second"))
+		require.NoError(err)
+
+		paths, err := CollectAll(subjects)
+		require.NoError(err)
+
+		// Should get both concrete users
+		require.Len(paths, 2)
+		subjectIDs := slicez.Map(paths, func(p Path) string { return p.Subject.ObjectID })
+		require.ElementsMatch([]string{"alice", "bob"}, subjectIDs)
+	})
+}
