tenant: introduce systemtenant

eval.go

@@ -231,7 +231,7 @@ nextFileMatch:
 
 			// 🚨 SECURITY: Skip documents that don't belong to the tenant. This check is
 			// necessary to prevent leaking data across tenants.
-			if !tenant.EqualsID(ctx, repoMetadata.TenantID) {
+			if !tenant.HasAccess(ctx, repoMetadata.TenantID) {
 				continue
 			}
 
@@ -624,7 +624,7 @@ func (d *indexData) List(ctx context.Context, q query.Q, opts *ListOptions) (rl
 		}
 		// 🚨 SECURITY: Skip documents that don't belong to the tenant. This check is
 		// necessary to prevent leaking data across tenants.
-		if !tenant.EqualsID(ctx, d.repoMetaData[i].TenantID) {
+		if !tenant.HasAccess(ctx, d.repoMetaData[i].TenantID) {
 			continue
 		}
 		rle := &d.repoListEntry[i]

internal/tenant/query.go

@@ -2,14 +2,19 @@ package tenant
 
 import (
 	"context"
+
+	"github.com/sourcegraph/zoekt/internal/tenant/systemtenant"
 )
 
-// EqualsID returns true if the tenant ID in the context matches the
+// HasAccess returns true if the tenant ID in the context matches the
 // given ID. If tenant enforcement is disabled, it always returns true.
-func EqualsID(ctx context.Context, id int) bool {
+func HasAccess(ctx context.Context, id int) bool {
 	if !EnforceTenant() {
 		return true
 	}
+	if systemtenant.Is(ctx) {
+		return true
+	}
 	t, err := FromContext(ctx)
 	if err != nil {
 		return false

internal/tenant/systemtenant/systemtenant.go

@@ -0,0 +1,22 @@
+// Package systemtenant exports UnsafeCtx which allows to access shards across
+// all tenants. This must only be used for tasks that are not request specific.
+package systemtenant
+
+import (
+	"context"
+)
+
+type contextKey int
+
+const systemTenantKey contextKey = iota
+
+// UnsafeCtx is a context that allows queries across all tenants. Don't use this
+// for user requests.
+var UnsafeCtx = context.WithValue(context.Background(), systemTenantKey, systemTenantKey)
+
+// Is returns true if the context has been marked to allow queries across all
+// tenants.
+func Is(ctx context.Context) bool {
+	_, ok := ctx.Value(systemTenantKey).(contextKey)
+	return ok
+}

internal/tenant/systemtenant/systemtenant_test.go

@@ -0,0 +1,15 @@
+package systemtenant
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSystemtenantRoundtrip(t *testing.T) {
+	if Is(context.Background()) {
+		t.Fatal()
+	}
+	require.True(t, Is(UnsafeCtx))
+}

shards/shards.go

@@ -35,6 +35,7 @@ import (
 	"go.uber.org/atomic"
 
 	"github.com/sourcegraph/zoekt"
+	"github.com/sourcegraph/zoekt/internal/tenant/systemtenant"
 	"github.com/sourcegraph/zoekt/query"
 	"github.com/sourcegraph/zoekt/trace"
 )
@@ -1064,7 +1065,10 @@ func (s *shardedSearcher) getLoaded() loaded {
 
 func mkRankedShard(s zoekt.Searcher) *rankedShard {
 	q := query.Const{Value: true}
-	result, err := s.List(context.Background(), &q, nil)
+	// We need to use UnsafeCtx here, otherwise we cannot return a proper
+	// rankedShard. On the user request path we use selectRepoSet which relies on
+	// rankedShard.repos being set.
+	result, err := s.List(systemtenant.UnsafeCtx, &q, nil)
 	if err != nil {
 		return &rankedShard{Searcher: s}
 	}