SpiceDB calls to LookupResources using LookupResources2 with caveats may return context is missing when it is not

@@ -11,6 +11,7 @@ import (
 	"github.com/ccoveille/go-safecast"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
+	"google.golang.org/protobuf/types/known/structpb"
 
 	"github.com/authzed/spicedb/internal/datastore/memdb"
 	"github.com/authzed/spicedb/internal/dispatch"
@@ -339,12 +340,14 @@ func TestMaxDepthLookup2(t *testing.T) {
 
 func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 	testCases := []struct {
-		name                string
-		schema              string
-		relationships       []*core.RelationTuple
-		permission          *core.RelationReference
-		subject             *core.ObjectAndRelation
-		expectedResourceIDs []string
+		name                  string
+		schema                string
+		relationships         []*core.RelationTuple
+		permission            *core.RelationReference
+		subject               *core.ObjectAndRelation
+		optionalCaveatContext map[string]any
+		expectedResourceIDs   []string
+		expectedMissingFields []string
 	}{
 		{
 			"basic union",
@@ -361,7 +364,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 1510),
+			nil,
 		},
 		{
 			"basic exclusion",
@@ -375,7 +380,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			genTuples("document", "viewer", "user", "tom", 1010),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 1010),
+			nil,
 		},
 		{
 			"basic intersection",
@@ -392,7 +399,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 510),
+			nil,
 		},
 		{
 			"union and excluded union",
@@ -411,7 +420,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 2450),
+			nil,
 		},
 		{
 			"basic caveats",
@@ -428,7 +439,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			genTuplesWithCaveat("document", "viewer", "user", "tom", "somecaveat", map[string]any{"somecondition": 42}, 0, 2450),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 2450),
+			nil,
 		},
 		{
 			"excluded items",
@@ -445,7 +458,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 1210),
+			nil,
 		},
 		{
 			"basic caveats with missing field",
@@ -462,7 +477,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			genTuplesWithCaveat("document", "viewer", "user", "tom", "somecaveat", map[string]any{}, 0, 2450),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 2450),
+			[]string{"somecondition"},
 		},
 		{
 			"larger arrow dispatch",
@@ -482,7 +499,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 150),
+			nil,
 		},
 		{
 			"big",
@@ -499,7 +518,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 15100),
+			nil,
 		},
 		{
 			"arrow under intersection",
@@ -523,7 +544,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 510),
+			nil,
 		},
 		{
 			"all arrow",
@@ -563,7 +586,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			},
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			[]string{"doc0", "doc1", "doc4"},
+			nil,
 		},
 		{
 			"indirect intersection and exclusion",
@@ -583,7 +608,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 1410),
+			nil,
 		},
 		{
 			"indirect intersections",
@@ -611,7 +638,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 1410),
+			nil,
 		},
 		{
 			"indirect over arrow",
@@ -639,7 +668,9 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 1410),
+			nil,
 		},
 		{
 			"root indirect with intermediate shearing",
@@ -675,7 +706,40 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
+			nil,
 			genResourceIds("document", 2000),
+			nil,
+		},
+		{
+			"indirect caveats",
+			`caveat somecaveat(somevalue int) {
+				somevalue == 42
+			}
+
+			definition user {}
+
+			definition container {
+				relation access: user with somecaveat 
+				permission accesses = access
+			}
+
+			definition document {
+				relation container: container
+				relation viewer: user with somecaveat
+				permission view = viewer & container->accesses
+			}`,
+			joinTuples(
+				[]*core.RelationTuple{
+					tuple.MustParse("container:somecontainer#access@user:tom[somecaveat]"),
+				},
+				genTuplesWithCaveat("document", "viewer", "user", "tom", "somecaveat", map[string]any{}, 0, 2450),
+				genTuples("document", "container", "container", "somecontainer", 2450),
+			),
+			RR("document", "view"),
+			ONR("user", "tom", "..."),
+			map[string]any{"somevalue": 42},
+			genResourceIds("document", 2450),
+			nil,
 		},
 	}
 
@@ -706,6 +770,12 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 						uintPageSize, err := safecast.ToUint32(pageSize)
 						require.NoError(err)
 
+						var caveatContext *structpb.Struct
+						if tc.optionalCaveatContext != nil {
+							caveatContext, err = structpb.NewStruct(tc.optionalCaveatContext)
+							require.NoError(err)
+						}
+
 						err = dispatcher.DispatchLookupResources2(&v1.DispatchLookupResources2Request{
 							ResourceRelation: tc.permission,
 							SubjectRelation:  RR(tc.subject.Namespace, "..."),
@@ -717,6 +787,7 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 							},
 							OptionalLimit:  uintPageSize,
 							OptionalCursor: currentCursor,
+							Context:        caveatContext,
 						}, stream)
 						require.NoError(err)
 
@@ -727,6 +798,7 @@ func TestLookupResources2OverSchemaWithCursors(t *testing.T) {
 						foundChunks = append(foundChunks, stream.Results())
 
 						for _, result := range stream.Results() {
+							require.ElementsMatch(tc.expectedMissingFields, result.Resource.MissingContextParams)
 							foundResourceIDs.Insert(result.Resource.ResourceId)
 							currentCursor = result.AfterResponseCursor
 						}

@@ -612,6 +612,7 @@ func (crr *CursoredLookupResources2) redispatchOrReport(
 					},
 					OptionalCursor: ci.currentCursor,
 					OptionalLimit:  parentRequest.OptionalLimit,
+					Context:        parentRequest.Context,
 				}, stream)
 			}
 

@@ -227,6 +227,7 @@ func (rdc *checkAndDispatchRunner) runDispatch(
 		},
 		OptionalCursor: updatedCi.currentCursor,
 		OptionalLimit:  rdc.ci.limits.currentLimit,
+		Context:        rdc.parentRequest.Context,
 	}, wrappedStream)
 }
 

@@ -360,6 +360,8 @@ func requireSubsetOf(t *testing.T, found []string, expected []string) {
 // validateLookupResources ensures that a lookup resources call returns the expected objects and
 // only those expected.
 func validateLookupResources(t *testing.T, vctx validationContext) {
+	// Run a lookup resources for each resource type and ensure that the returned objects are those
+	// that are accessible to the subject.
 	testForEachResourceType(t, vctx, "validate_lookup_resources",
 		func(t *testing.T, resourceRelation *core.RelationReference) {
 			for _, subject := range vctx.accessibilitySet.AllSubjectsNoWildcards() {
@@ -375,7 +377,7 @@ func validateLookupResources(t *testing.T, vctx validationContext) {
 							var currentCursor *v1.Cursor
 							resolvedResources := map[string]*v1.LookupResourcesResponse{}
 							for i := 0; i < 100; i++ {
-								foundResources, lastCursor, err := vctx.serviceTester.LookupResources(context.Background(), resourceRelation, subject, vctx.revision, currentCursor, pageSize)
+								foundResources, lastCursor, err := vctx.serviceTester.LookupResources(context.Background(), resourceRelation, subject, vctx.revision, currentCursor, pageSize, nil)
 								require.NoError(t, err)
 
 								if pageSize > 0 {
@@ -676,45 +678,73 @@ func runAssertions(t *testing.T, vctx validationContext) {
 							require.NoError(t, err)
 							require.Equal(t, entry.expectedPermissionship, permissionship, "Assertion `%s` returned %s; expected %s", tuple.MustString(rel), permissionship, entry.expectedPermissionship)
 
-							// Ensure the assertion passes LookupResources.
-							resolvedResources, _, err := vctx.serviceTester.LookupResources(context.Background(), &core.RelationReference{
+							// Ensure the assertion passes LookupResources with context, directly.
+							resolvedDirectResources, _, err := vctx.serviceTester.LookupResources(context.Background(), &core.RelationReference{
 								Namespace: rel.ResourceAndRelation.Namespace,
 								Relation:  rel.ResourceAndRelation.Relation,
-							}, rel.Subject, vctx.revision, nil, 0)
+							}, rel.Subject, vctx.revision, nil, 0, assertion.CaveatContext)
 							require.NoError(t, err)
 
-							resolvedResourcesMap := map[string]*v1.LookupResourcesResponse{}
-							for _, resource := range resolvedResources {
-								resolvedResourcesMap[resource.ResourceObjectId] = resource
+							resolvedDirectResourcesMap := map[string]*v1.LookupResourcesResponse{}
+							for _, resource := range resolvedDirectResources {
+								resolvedDirectResourcesMap[resource.ResourceObjectId] = resource
 							}
 
-							resolvedResourceIds := maps.Keys(resolvedResourcesMap)
+							// Ensure the assertion passes LookupResources without context, indirectly.
+							resolvedIndirectResources, _, err := vctx.serviceTester.LookupResources(context.Background(), &core.RelationReference{
+								Namespace: rel.ResourceAndRelation.Namespace,
+								Relation:  rel.ResourceAndRelation.Relation,
+							}, rel.Subject, vctx.revision, nil, 0, nil)
+							require.NoError(t, err)
+
+							resolvedIndirectResourcesMap := map[string]*v1.LookupResourcesResponse{}
+							for _, resource := range resolvedIndirectResources {
+								resolvedIndirectResourcesMap[resource.ResourceObjectId] = resource
+							}
+
+							// Check the assertion was returned for a direct (with context) lookup.
+							resolvedDirectResourceIds := maps.Keys(resolvedDirectResourcesMap)
+							switch permissionship {
+							case v1.CheckPermissionResponse_PERMISSIONSHIP_NO_PERMISSION:
+								require.NotContains(t, resolvedDirectResourceIds, rel.ResourceAndRelation.ObjectId, "Found unexpected object %s in direct lookup for assertion %s", rel.ResourceAndRelation, rel)
+
+							case v1.CheckPermissionResponse_PERMISSIONSHIP_HAS_PERMISSION:
+								require.Contains(t, resolvedDirectResourceIds, rel.ResourceAndRelation.ObjectId, "Missing object %s in lookup for assertion %s", rel.ResourceAndRelation, rel)
+								require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_HAS_PERMISSION, resolvedDirectResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
+
+							case v1.CheckPermissionResponse_PERMISSIONSHIP_CONDITIONAL_PERMISSION:
+								require.Contains(t, resolvedDirectResourceIds, rel.ResourceAndRelation.ObjectId, "Missing object %s in lookup for assertion %s", rel.ResourceAndRelation, rel)
+								require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_CONDITIONAL_PERMISSION, resolvedDirectResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
+							}
+
+							// Check the assertion was returned for an indirect (without context) lookup.
+							resolvedIndirectResourceIds := maps.Keys(resolvedIndirectResourcesMap)
 							accessibility, _, _ := vctx.accessibilitySet.AccessibiliyAndPermissionshipFor(rel.ResourceAndRelation, rel.Subject)
 
 							switch permissionship {
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_NO_PERMISSION:
 								// If the caveat context given is empty, then the lookup result must not exist at all.
 								// Otherwise, it *could* be caveated or not exist, depending on the context given.
 								if len(assertion.CaveatContext) == 0 {
-									require.NotContains(t, resolvedResourceIds, rel.ResourceAndRelation.ObjectId, "Found unexpected object %s in lookup for assertion %s", rel.ResourceAndRelation, rel)
+									require.NotContains(t, resolvedIndirectResourceIds, rel.ResourceAndRelation.ObjectId, "Found unexpected object %s in indirect lookup for assertion %s", rel.ResourceAndRelation, rel)
 								} else if accessibility == consistencytestutil.NotAccessible {
-									found, ok := resolvedResourcesMap[rel.ResourceAndRelation.ObjectId]
+									found, ok := resolvedIndirectResourcesMap[rel.ResourceAndRelation.ObjectId]
 									require.True(t, !ok || found.Permissionship != v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_HAS_PERMISSION) // LookupResources can be caveated, since we didn't rerun LookupResources with the context
 								} else if accessibility != consistencytestutil.NotAccessibleDueToPrespecifiedCaveat {
-									require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_CONDITIONAL_PERMISSION, resolvedResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
+									require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_CONDITIONAL_PERMISSION, resolvedIndirectResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
 								}
 
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_HAS_PERMISSION:
-								require.Contains(t, resolvedResourceIds, rel.ResourceAndRelation.ObjectId, "Missing object %s in lookup for assertion %s", rel.ResourceAndRelation, rel)
+								require.Contains(t, resolvedIndirectResourceIds, rel.ResourceAndRelation.ObjectId, "Missing object %s in lookup for assertion %s", rel.ResourceAndRelation, rel)
 								// If the caveat context given is empty, then the lookup result must be fully permissioned.
 								// Otherwise, it *could* be caveated or fully permissioned, depending on the context given.
 								if len(assertion.CaveatContext) == 0 {
-									require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_HAS_PERMISSION, resolvedResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
+									require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_HAS_PERMISSION, resolvedIndirectResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
 								}
 
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_CONDITIONAL_PERMISSION:
-								require.Contains(t, resolvedResourceIds, rel.ResourceAndRelation.ObjectId, "Missing object %s in lookup for assertion %s", rel.ResourceAndRelation, rel)
-								require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_CONDITIONAL_PERMISSION, resolvedResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
+								require.Contains(t, resolvedIndirectResourceIds, rel.ResourceAndRelation.ObjectId, "Missing object %s in lookup for assertion %s", rel.ResourceAndRelation, rel)
+								require.Equal(t, v1.LookupPermissionship_LOOKUP_PERMISSIONSHIP_CONDITIONAL_PERMISSION, resolvedIndirectResourcesMap[rel.ResourceAndRelation.ObjectId].Permissionship)
 
 							default:
 								panic("unknown permissionship")

@@ -28,7 +28,7 @@ type ServiceTester interface {
 	Expand(ctx context.Context, resource *core.ObjectAndRelation, atRevision datastore.Revision) (*core.RelationTupleTreeNode, error)
 	Write(ctx context.Context, relationship *core.RelationTuple) error
 	Read(ctx context.Context, namespaceName string, atRevision datastore.Revision) ([]*core.RelationTuple, error)
-	LookupResources(ctx context.Context, resourceRelation *core.RelationReference, subject *core.ObjectAndRelation, atRevision datastore.Revision, cursor *v1.Cursor, limit uint32) ([]*v1.LookupResourcesResponse, *v1.Cursor, error)
+	LookupResources(ctx context.Context, resourceRelation *core.RelationReference, subject *core.ObjectAndRelation, atRevision datastore.Revision, cursor *v1.Cursor, limit uint32, caveatContext map[string]any) ([]*v1.LookupResourcesResponse, *v1.Cursor, error)
 	LookupSubjects(ctx context.Context, resource *core.ObjectAndRelation, subjectRelation *core.RelationReference, atRevision datastore.Revision, caveatContext map[string]any) (map[string]*v1.LookupSubjectsResponse, error)
 	// NOTE: ExperimentalService/BulkCheckPermission has been promoted to PermissionsService/CheckBulkPermissions
 	BulkCheck(ctx context.Context, items []*v1.BulkCheckPermissionRequestItem, atRevision datastore.Revision) ([]*v1.BulkCheckPermissionPair, error)
@@ -153,7 +153,16 @@ func (v1st v1ServiceTester) Read(_ context.Context, namespaceName string, atRevi
 	return tuples, nil
 }
 
-func (v1st v1ServiceTester) LookupResources(_ context.Context, resourceRelation *core.RelationReference, subject *core.ObjectAndRelation, atRevision datastore.Revision, cursor *v1.Cursor, limit uint32) ([]*v1.LookupResourcesResponse, *v1.Cursor, error) {
+func (v1st v1ServiceTester) LookupResources(_ context.Context, resourceRelation *core.RelationReference, subject *core.ObjectAndRelation, atRevision datastore.Revision, cursor *v1.Cursor, limit uint32, caveatContext map[string]any) ([]*v1.LookupResourcesResponse, *v1.Cursor, error) {
+	var builtContext *structpb.Struct
+	if caveatContext != nil {
+		built, err := structpb.NewStruct(caveatContext)
+		if err != nil {
+			return nil, nil, err
+		}
+		builtContext = built
+	}
+
 	lookupResp, err := v1st.permClient.LookupResources(context.Background(), &v1.LookupResourcesRequest{
 		ResourceObjectType: resourceRelation.Namespace,
 		Permission:         resourceRelation.Relation,
@@ -171,6 +180,7 @@ func (v1st v1ServiceTester) LookupResources(_ context.Context, resourceRelation
 		},
 		OptionalLimit:  limit,
 		OptionalCursor: cursor,
+		Context:        builtContext,
 	})
 	if err != nil {
 		return nil, nil, err

@@ -0,0 +1,29 @@
+---
+schema: |+
+  caveat unexpired(expires_at timestamp, now timestamp) {
+    now < expires_at
+  }
+
+  definition user {}
+
+  definition container {
+    relation access: user with unexpired 
+    permission accesses = access
+  }
+
+  definition resource {
+    relation container: container
+    relation viewer: user with unexpired
+    permission viewers = viewer & container->accesses
+  }
+relationships: |-
+  resource:someresource#container@container:somecontainer
+  resource:someresource#viewer@user:tom[unexpired:{"expires_at":"2024-10-03T14:16:30.776734Z"}]
+  container:somecontainer#access@user:fred[unexpired:{"expires_at":"2024-01-01T00:00:00Z"}]
+  container:somecontainer#access@user:tom[unexpired:{"expires_at":"2024-10-03T07:16:28.068632-07:00"}]
+  container:somecontainer#access@user:sarah[unexpired:{"expires_at":"2024-10-03T07:16:28.068632-07:00"}]
+assertions:
+  assertTrue:
+    - 'resource:someresource#viewers@user:tom with {"now": "2023-01-01T00:00:00Z"}'
+  assertFalse:
+    - 'resource:someresource#viewers@user:tom with {"now": "2025-01-01T00:00:00Z"}'

@@ -0,0 +1,25 @@
+---
+schema: |+
+  definition user {}
+
+  caveat some_caveat(somevalue int) {
+    somevalue == 42
+  }
+
+  definition team {
+    relation member: user with some_caveat
+  }
+
+  definition document {
+    relation viewer: team#member with some_caveat
+    permission view = viewer->member
+  }
+
+relationships: |-
+  team:first#member@user:tom[some_caveat]
+  document:firstdoc#viewer@team:first#member[some_caveat]
+assertions:
+  assertTrue:
+    - 'document:firstdoc#view@user:tom with {"somevalue": 42}'
+  assertFalse:
+    - 'document:firstdoc#view@user:tom with {"somevalue": 41}'

@@ -0,0 +1,33 @@
+---
+schema: |+
+  definition user {}
+
+  caveat some_caveat(somevalue int) {
+    somevalue == 42
+  }
+
+  definition team {
+    relation member: user with some_caveat
+  }
+
+  definition document {
+    relation viewer: team#member with some_caveat
+    relation editor: team#member with some_caveat
+    relation admin: team#member with some_caveat
+
+    permission inner_view = viewer & editor
+    permission view = inner_view & admin
+  }
+
+relationships: |-
+  team:first#member@user:tom[some_caveat]
+  team:second#member@user:tom[some_caveat]
+  team:third#member@user:tom[some_caveat]
+  document:firstdoc#viewer@team:first#member[some_caveat]
+  document:firstdoc#editor@team:second#member[some_caveat]
+  document:firstdoc#admin@team:third#member[some_caveat]
+assertions:
+  assertTrue:
+    - 'document:firstdoc#view@user:tom with {"somevalue": 42}'
+  assertFalse:
+    - 'document:firstdoc#view@user:tom with {"somevalue": 41}'