tenant: run healthz check with system priviledges (#877)

This is an alternative to #875.

We run the health check with system priviledges. This way we run an
actual search, just like we do if tenant enforcement is off.

I also make sure we don't log system searches as "missing_tenant".

internal/tenant/context.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
+	"github.com/sourcegraph/zoekt/internal/tenant/systemtenant"
 	"github.com/sourcegraph/zoekt/trace"
 )
 
@@ -26,6 +27,10 @@ func FromContext(ctx context.Context) (*tenanttype.Tenant, error) {
 // Log logs the tenant ID to the trace. If tenant logging is enabled, it also
 // logs a stack trace to a pprof profile.
 func Log(ctx context.Context, tr *trace.Trace) {
+	if systemtenant.Is(ctx) {
+		tr.LazyPrintf("tenant: system")
+		return
+	}
 	tnt, ok := tenanttype.GetTenant(ctx)
 	if !ok {
 		if profile := pprofMissingTenant(); profile != nil {

internal/tenant/systemtenant/systemtenant.go

@@ -10,9 +10,11 @@ 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)
+// WithUnsafeContext taints the context to allow queries across all tenants.
+// Never use this for user requests.
+func WithUnsafeContext(ctx context.Context) context.Context {
+	return context.WithValue(ctx, systemTenantKey, systemTenantKey)
+}
 
 // Is returns true if the context has been marked to allow queries across all
 // tenants.

internal/tenant/systemtenant/systemtenant_test.go

@@ -7,9 +7,9 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestSystemtenantRoundtrip(t *testing.T) {
+func TestSystemTenantRoundTrip(t *testing.T) {
 	if Is(context.Background()) {
 		t.Fatal()
 	}
-	require.True(t, Is(UnsafeCtx))
+	require.True(t, Is(WithUnsafeContext(context.Background())))
 }

shards/shards.go

@@ -1083,10 +1083,10 @@ func (s *shardedSearcher) getLoaded() loaded {
 
 func mkRankedShard(s zoekt.Searcher) *rankedShard {
 	q := query.Const{Value: true}
-	// We need to use UnsafeCtx here, otherwise we cannot return a proper
+	// We need to use WithUnsafeContext 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)
+	result, err := s.List(systemtenant.WithUnsafeContext(context.Background()), &q, nil)
 	if err != nil {
 		log.Printf("[ERROR] mkRankedShard(%s): failed to cache repository list: %v", s, err)
 		return &rankedShard{Searcher: s}

web/server.go

@@ -32,7 +32,9 @@ import (
 	"time"
 
 	"github.com/grafana/regexp"
+
 	"github.com/sourcegraph/zoekt"
+	"github.com/sourcegraph/zoekt/internal/tenant/systemtenant"
 	zjson "github.com/sourcegraph/zoekt/json"
 	"github.com/sourcegraph/zoekt/query"
 )
@@ -206,7 +208,10 @@ func (s *Server) serveHealthz(w http.ResponseWriter, r *http.Request) {
 	q := &query.Const{Value: true}
 	opts := &zoekt.SearchOptions{ShardMaxMatchCount: 1, TotalMaxMatchCount: 1, MaxDocDisplayCount: 1}
 
-	result, err := s.Searcher.Search(r.Context(), q, opts)
+	// We need to use WithUnsafeContext here because we want to perform a full
+	// search returning results. The result of this search is not used for anything
+	// other than determining if the server is healthy.
+	result, err := s.Searcher.Search(systemtenant.WithUnsafeContext(r.Context()), q, opts)
 	if err != nil {
 		http.Error(w, fmt.Sprintf("not ready: %v", err), http.StatusInternalServerError)
 		return