Merge pull request #2553 from authzed/barakmich/wildcard

feat: add wildcard support to the plan

Author: barakmich

internal/services/integrationtesting/query_plan_consistency_test.go

@@ -127,13 +127,13 @@ func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
 
 							switch entry.expectedPermissionship {
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_CONDITIONAL_PERMISSION:
-								require.Equal(len(rels), 1)
+								require.Len(rels, 1)
 								require.NotNil(rels[0].OptionalCaveat)
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_HAS_PERMISSION:
-								require.Equal(len(rels), 1)
+								require.Len(rels, 1)
 								require.Nil(rels[0].OptionalCaveat)
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_NO_PERMISSION:
-								require.Equal(len(rels), 0)
+								require.Len(rels, 0)
 							}
 						})
 					}

pkg/query/build_tree.go

@@ -147,6 +147,11 @@ func (b *iteratorBuilder) buildBaseRelationIterator(br *schema.BaseRelation, wit
 		return base, nil
 	}
 
+	// Wildcards represent direct access, so no subrelation processing needed
+	if br.Wildcard {
+		return base, nil
+	}
+
 	// We must check the effective arrow of a subrelation if we have one and subrelations are enabled
 	// (subrelations are disabled in cases of actual arrows)
 	union := NewUnion()

pkg/query/build_tree_test.go

@@ -724,3 +724,100 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		require.NoError(err)
 	})
 }
+
+func TestBuildTreeWildcardIterator(t *testing.T) {
+	t.Parallel()
+
+	require := require.New(t)
+
+	// Create a simple schema with a wildcard relation using core types directly
+	docDef := &corev1.NamespaceDefinition{
+		Name: "document",
+		Relation: []*corev1.Relation{
+			{
+				Name: "viewer",
+				TypeInformation: &corev1.TypeInformation{
+					AllowedDirectRelations: []*corev1.AllowedRelation{
+						{
+							Namespace: "user",
+							RelationOrWildcard: &corev1.AllowedRelation_PublicWildcard_{
+								PublicWildcard: &corev1.AllowedRelation_PublicWildcard{},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	userDef := &corev1.NamespaceDefinition{
+		Name: "user",
+	}
+
+	objectDefs := []*corev1.NamespaceDefinition{userDef, docDef}
+	dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
+	require.NoError(err)
+
+	// Verify the schema has the wildcard BaseRelation
+	documentDef := dsSchema.Definitions["document"]
+	require.NotNil(documentDef)
+	viewerRelation := documentDef.Relations["viewer"]
+	require.NotNil(viewerRelation)
+	require.Len(viewerRelation.BaseRelations, 1)
+	baseRel := viewerRelation.BaseRelations[0]
+	require.True(baseRel.Wildcard, "BaseRelation should have Wildcard: true")
+	require.Equal("user", baseRel.Type)
+
+	// Print debug info
+	t.Logf("BaseRelation: Type=%s, Subrelation=%s, Wildcard=%v", baseRel.Type, baseRel.Subrelation, baseRel.Wildcard)
+
+	t.Run("Schema with wildcard creates WildcardIterator", func(t *testing.T) {
+		t.Parallel()
+		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		require.NoError(err)
+		require.NotNil(it)
+
+		// Verify it's an Alias wrapping a RelationIterator with wildcard support
+		require.IsType(&Alias{}, it)
+		alias := it.(*Alias)
+		require.IsType(&RelationIterator{}, alias.subIt)
+
+		// Check the explain output contains wildcard information
+		explain := it.Explain()
+		explainStr := explain.String()
+		require.Contains(explainStr, "Relation")
+		require.Contains(explainStr, "user:*")
+	})
+
+	t.Run("Mixed wildcard and regular relations", func(t *testing.T) {
+		t.Parallel()
+		// Create a schema with both wildcard and regular relations
+		mixedDocDef := namespace.Namespace(
+			"document",
+			namespace.MustRelation("viewer", nil,
+				namespace.AllowedRelation("user", ""),    // Regular relation
+				namespace.AllowedPublicNamespace("user"), // Wildcard relation
+			),
+		)
+
+		mixedObjectDefs := []*corev1.NamespaceDefinition{userDef, mixedDocDef}
+		mixedSchema, err := schema.BuildSchemaFromDefinitions(mixedObjectDefs, nil)
+		require.NoError(err)
+
+		it, err := BuildIteratorFromSchema(mixedSchema, "document", "viewer")
+		require.NoError(err)
+		require.NotNil(it)
+
+		// Should create an alias with a union containing both regular and wildcard iterators
+		require.IsType(&Alias{}, it)
+		alias := it.(*Alias)
+		require.IsType(&Union{}, alias.subIt)
+
+		// Check explain contains both relation types (regular and wildcard)
+		explain := it.Explain()
+		explainStr := explain.String()
+		require.Contains(explainStr, "Union")
+		require.Contains(explainStr, "user:...", "should contain regular relation")
+		require.Contains(explainStr, "user:*", "should contain wildcard relation")
+	})
+}

pkg/query/datastore.go

@@ -36,6 +36,20 @@ func (r *RelationIterator) buildSubjectRelationFilter() datastore.SubjectRelatio
 }
 
 func (r *RelationIterator) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRelation) (RelationSeq, error) {
+	// If the subject type doesn't match the base relation type, return no results
+	if subject.ObjectType != r.base.Type {
+		return func(yield func(Relation, error) bool) {
+			// Empty sequence
+		}, nil
+	}
+
+	if r.base.Wildcard {
+		return r.checkWildcardImpl(ctx, resources, subject)
+	}
+	return r.checkNormalImpl(ctx, resources, subject)
+}
+
+func (r *RelationIterator) checkNormalImpl(ctx *Context, resources []Object, subject ObjectAndRelation) (RelationSeq, error) {
 	resourceIDs := make([]string, len(resources))
 	for i, res := range resources {
 		resourceIDs[i] = res.ObjectID
@@ -68,14 +82,101 @@ func (r *RelationIterator) CheckImpl(ctx *Context, resources []Object, subject O
 	return RelationSeq(relIter), nil
 }
 
+func (r *RelationIterator) checkWildcardImpl(ctx *Context, resources []Object, subject ObjectAndRelation) (RelationSeq, error) {
+	// Query the datastore for wildcard relationships (subject ObjectID = "*")
+	resourceIDs := make([]string, len(resources))
+	for i, res := range resources {
+		resourceIDs[i] = res.ObjectID
+	}
+
+	filter := datastore.RelationshipsFilter{
+		OptionalResourceType:     r.base.DefinitionName(),
+		OptionalResourceIds:      resourceIDs,
+		OptionalResourceRelation: r.base.RelationName(),
+		OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+			{
+				OptionalSubjectType: r.base.Type,
+				OptionalSubjectIds:  []string{tuple.PublicWildcard}, // Look for "*" subjects
+				RelationFilter:      r.buildSubjectRelationFilter(),
+			},
+		},
+	}
+
+	reader := ctx.Datastore.SnapshotReader(ctx.Revision)
+
+	relIter, err := reader.QueryRelationships(ctx, filter,
+		options.WithSkipCaveats(r.base.Caveat == ""),
+		options.WithSkipExpiration(!r.base.Expiration),
+		options.WithQueryShape(queryshape.CheckPermissionSelectDirectSubjects),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Transform the wildcard relationships to use the concrete subject
+	return func(yield func(Relation, error) bool) {
+		for rel, err := range relIter {
+			if err != nil {
+				if !yield(rel, err) {
+					return
+				}
+				continue
+			}
+
+			// Replace the wildcard subject with the concrete subject
+			concreteRel := rel
+			concreteRel.Subject = subject
+
+			if !yield(concreteRel, nil) {
+				return
+			}
+		}
+	}, nil
+}
+
 func (r *RelationIterator) IterSubjectsImpl(ctx *Context, resource Object) (RelationSeq, error) {
+	if r.base.Wildcard {
+		return r.iterSubjectsWildcardImpl(ctx, resource)
+	}
+	return r.iterSubjectsNormalImpl(ctx, resource)
+}
+
+func (r *RelationIterator) iterSubjectsNormalImpl(ctx *Context, resource Object) (RelationSeq, error) {
+	filter := datastore.RelationshipsFilter{
+		OptionalResourceType:     r.base.DefinitionName(),
+		OptionalResourceIds:      []string{resource.ObjectID},
+		OptionalResourceRelation: r.base.RelationName(),
+		OptionalSubjectsSelectors: []datastore.SubjectsSelector{
+			{
+				OptionalSubjectType: r.base.Type,
+				RelationFilter:      r.buildSubjectRelationFilter(),
+			},
+		},
+	}
+
+	reader := ctx.Datastore.SnapshotReader(ctx.Revision)
+
+	relIter, err := reader.QueryRelationships(ctx, filter,
+		options.WithSkipCaveats(r.base.Caveat == ""),
+		options.WithSkipExpiration(!r.base.Expiration),
+		options.WithQueryShape(queryshape.AllSubjectsForResources),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return RelationSeq(relIter), nil
+}
+
+func (r *RelationIterator) iterSubjectsWildcardImpl(ctx *Context, resource Object) (RelationSeq, error) {
 	filter := datastore.RelationshipsFilter{
 		OptionalResourceType:     r.base.DefinitionName(),
 		OptionalResourceIds:      []string{resource.ObjectID},
 		OptionalResourceRelation: r.base.RelationName(),
 		OptionalSubjectsSelectors: []datastore.SubjectsSelector{
 			{
 				OptionalSubjectType: r.base.Type,
+				OptionalSubjectIds:  []string{tuple.PublicWildcard}, // Look for "*" subjects
 				RelationFilter:      r.buildSubjectRelationFilter(),
 			},
 		},
@@ -106,7 +207,13 @@ func (r *RelationIterator) Clone() Iterator {
 }
 
 func (r *RelationIterator) Explain() Explain {
+	relationName := r.base.Subrelation
+	if r.base.Wildcard {
+		relationName = "*"
+	}
 	return Explain{
-		Info: fmt.Sprintf("Relation(%s:%s, \"%s\", caveat: %v, expiration: %v)", r.base.DefinitionName(), r.base.RelationName(), r.base.Subrelation, r.base.Caveat != "", r.base.Expiration),
+		Info: fmt.Sprintf("Relation(%s:%s -> %s:%s, caveat: %v, expiration: %v)",
+			r.base.DefinitionName(), r.base.RelationName(), r.base.Type, relationName,
+			r.base.Caveat != "", r.base.Expiration),
 	}
 }