1 files changed, 205 insertions, 0 deletions
diff --git a/vendor/github.com/authzed/spicedb/internal/graph/computed/computecheck.go b/vendor/github.com/authzed/spicedb/internal/graph/computed/computecheck.go
new file mode 100644
index 0000000..0bf20b5
--- /dev/null
+++ b/vendor/github.com/authzed/spicedb/internal/graph/computed/computecheck.go
@@ -0,0 +1,205 @@
+package computed
+
+import (
+	"context"
+
+	cexpr "github.com/authzed/spicedb/internal/caveats"
+	"github.com/authzed/spicedb/internal/dispatch"
+	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
+	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/genutil/slicez"
+	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
+	"github.com/authzed/spicedb/pkg/spiceerrors"
+	"github.com/authzed/spicedb/pkg/tuple"
+)
+
+// DebugOption defines the various debug level options for Checks.
+type DebugOption int
+
+const (
+	// NoDebugging indicates that debug information should be retained
+	// while performing the Check.
+	NoDebugging DebugOption = 0
+
+	// BasicDebuggingEnabled indicates that basic debug information, such
+	// as which steps were taken, should be retained while performing the
+	// Check and returned to the caller.
+	//
+	// NOTE: This has a minor performance impact.
+	BasicDebuggingEnabled DebugOption = 1
+
+	// TraceDebuggingEnabled indicates that the Check is being issued for
+	// tracing the exact calls made for debugging, which means that not only
+	// should debug information be recorded and returned, but that optimizations
+	// such as batching should be disabled.
+	//
+	// WARNING: This has a fairly significant performance impact and should only
+	// be used in tooling!
+	TraceDebuggingEnabled DebugOption = 2
+)
+
+// CheckParameters are the parameters for the ComputeCheck call. *All* are required.
+type CheckParameters struct {
+	ResourceType  tuple.RelationReference
+	Subject       tuple.ObjectAndRelation
+	CaveatContext map[string]any
+	AtRevision    datastore.Revision
+	MaximumDepth  uint32
+	DebugOption   DebugOption
+	CheckHints    []*v1.CheckHint
+}
+
+// ComputeCheck computes a check result for the given resource and subject, computing any
+// caveat expressions found.
+func ComputeCheck(
+	ctx context.Context,
+	d dispatch.Check,
+	ts *caveattypes.TypeSet,
+	params CheckParameters,
+	resourceID string,
+	dispatchChunkSize uint16,
+) (*v1.ResourceCheckResult, *v1.ResponseMeta, error) {
+	resultsMap, meta, di, err := computeCheck(ctx, d, ts, params, []string{resourceID}, dispatchChunkSize)
+	if err != nil {
+		return nil, meta, err
+	}
+
+	spiceerrors.DebugAssert(func() bool {
+		return (len(di) == 0 && meta.DebugInfo == nil) || (len(di) == 1 && meta.DebugInfo != nil)
+	}, "mismatch in debug information returned from computeCheck")
+
+	return resultsMap[resourceID], meta, err
+}
+
+// ComputeBulkCheck computes a check result for the given resources and subject, computing any
+// caveat expressions found.
+func ComputeBulkCheck(
+	ctx context.Context,
+	d dispatch.Check,
+	ts *caveattypes.TypeSet,
+	params CheckParameters,
+	resourceIDs []string,
+	dispatchChunkSize uint16,
+) (map[string]*v1.ResourceCheckResult, *v1.ResponseMeta, []*v1.DebugInformation, error) {
+	return computeCheck(ctx, d, ts, params, resourceIDs, dispatchChunkSize)
+}
+
+func computeCheck(ctx context.Context,
+	d dispatch.Check,
+	ts *caveattypes.TypeSet,
+	params CheckParameters,
+	resourceIDs []string,
+	dispatchChunkSize uint16,
+) (map[string]*v1.ResourceCheckResult, *v1.ResponseMeta, []*v1.DebugInformation, error) {
+	debugging := v1.DispatchCheckRequest_NO_DEBUG
+	if params.DebugOption == BasicDebuggingEnabled {
+		debugging = v1.DispatchCheckRequest_ENABLE_BASIC_DEBUGGING
+	} else if params.DebugOption == TraceDebuggingEnabled {
+		debugging = v1.DispatchCheckRequest_ENABLE_TRACE_DEBUGGING
+	}
+
+	setting := v1.DispatchCheckRequest_REQUIRE_ALL_RESULTS
+	if len(resourceIDs) == 1 {
+		setting = v1.DispatchCheckRequest_ALLOW_SINGLE_RESULT
+	}
+
+	// Ensure that the number of resources IDs given to each dispatch call is not in excess of the maximum.
+	results := make(map[string]*v1.ResourceCheckResult, len(resourceIDs))
+	metadata := &v1.ResponseMeta{}
+
+	bf, err := v1.NewTraversalBloomFilter(uint(params.MaximumDepth))
+	if err != nil {
+		return nil, nil, nil, spiceerrors.MustBugf("failed to create new traversal bloom filter")
+	}
+
+	caveatRunner := cexpr.NewCaveatRunner(ts)
+
+	// TODO(jschorr): Should we make this run in parallel via the preloadedTaskRunner?
+	debugInfo := make([]*v1.DebugInformation, 0)
+	_, err = slicez.ForEachChunkUntil(resourceIDs, dispatchChunkSize, func(resourceIDsToCheck []string) (bool, error) {
+		checkResult, err := d.DispatchCheck(ctx, &v1.DispatchCheckRequest{
+			ResourceRelation: params.ResourceType.ToCoreRR(),
+			ResourceIds:      resourceIDsToCheck,
+			ResultsSetting:   setting,
+			Subject:          params.Subject.ToCoreONR(),
+			Metadata: &v1.ResolverMeta{
+				AtRevision:     params.AtRevision.String(),
+				DepthRemaining: params.MaximumDepth,
+				TraversalBloom: bf,
+			},
+			Debug:      debugging,
+			CheckHints: params.CheckHints,
+		})
+
+		if checkResult.Metadata.DebugInfo != nil {
+			debugInfo = append(debugInfo, checkResult.Metadata.DebugInfo)
+		}
+
+		if len(resourceIDs) == 1 {
+			metadata = checkResult.Metadata
+		} else {
+			metadata = &v1.ResponseMeta{
+				DispatchCount:       metadata.DispatchCount + checkResult.Metadata.DispatchCount,
+				DepthRequired:       max(metadata.DepthRequired, checkResult.Metadata.DepthRequired),
+				CachedDispatchCount: metadata.CachedDispatchCount + checkResult.Metadata.CachedDispatchCount,
+				DebugInfo:           nil,
+			}
+		}
+
+		if err != nil {
+			return false, err
+		}
+
+		for _, resourceID := range resourceIDsToCheck {
+			computed, err := computeCaveatedCheckResult(ctx, caveatRunner, params, resourceID, checkResult)
+			if err != nil {
+				return false, err
+			}
+			results[resourceID] = computed
+		}
+
+		return true, nil
+	})
+	return results, metadata, debugInfo, err
+}
+
+func computeCaveatedCheckResult(ctx context.Context, runner *cexpr.CaveatRunner, params CheckParameters, resourceID string, checkResult *v1.DispatchCheckResponse) (*v1.ResourceCheckResult, error) {
+	result, ok := checkResult.ResultsByResourceId[resourceID]
+	if !ok {
+		return &v1.ResourceCheckResult{
+			Membership: v1.ResourceCheckResult_NOT_MEMBER,
+		}, nil
+	}
+
+	if result.Membership == v1.ResourceCheckResult_MEMBER {
+		return result, nil
+	}
+
+	ds := datastoremw.MustFromContext(ctx)
+	reader := ds.SnapshotReader(params.AtRevision)
+
+	caveatResult, err := runner.RunCaveatExpression(ctx, result.Expression, params.CaveatContext, reader, cexpr.RunCaveatExpressionNoDebugging)
+	if err != nil {
+		return nil, err
+	}
+
+	if caveatResult.IsPartial() {
+		missingFields, _ := caveatResult.MissingVarNames()
+		return &v1.ResourceCheckResult{
+			Membership:        v1.ResourceCheckResult_CAVEATED_MEMBER,
+			Expression:        result.Expression,
+			MissingExprFields: missingFields,
+		}, nil
+	}
+
+	if caveatResult.Value() {
+		return &v1.ResourceCheckResult{
+			Membership: v1.ResourceCheckResult_MEMBER,
+		}, nil
+	}
+
+	return &v1.ResourceCheckResult{
+		Membership: v1.ResourceCheckResult_NOT_MEMBER,
+	}, nil
+}