From 46af1c2edcfb5be07abcb296bfa027a3fdaa16c9 Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 09:38:05 +0100 Subject: [PATCH 1/4] tenant: introduce systemtenant Some tasks, such as loading shards, require priviledged access on startup. Here I introduce a systemtenant we can use for these things. This is motivated by bug where the symbol sidebar in multi-tenant node wouldn't work because ranked shards were not loaded correctly which in turn caused `SelectRepoSet` to return 0 shards always. Test plan: - added unit test - manual testing: symbol sidebar works now --- internal/tenant/query.go | 5 +++ internal/tenant/systemtenant/systemtenant.go | 33 +++++++++++++++++++ .../tenant/systemtenant/systemtenant_test.go | 17 ++++++++++ shards/shards.go | 7 +++- 4 files changed, 61 insertions(+), 1 deletion(-) create mode 100644 internal/tenant/systemtenant/systemtenant.go create mode 100644 internal/tenant/systemtenant/systemtenant_test.go diff --git a/internal/tenant/query.go b/internal/tenant/query.go index 3dbfd48fd..8a80d989d 100644 --- a/internal/tenant/query.go +++ b/internal/tenant/query.go @@ -2,6 +2,8 @@ package tenant import ( "context" + + "github.com/sourcegraph/zoekt/internal/tenant/systemtenant" ) // EqualsID returns true if the tenant ID in the context matches the @@ -10,6 +12,9 @@ func EqualsID(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 diff --git a/internal/tenant/systemtenant/systemtenant.go b/internal/tenant/systemtenant/systemtenant.go new file mode 100644 index 000000000..8e1a619fd --- /dev/null +++ b/internal/tenant/systemtenant/systemtenant.go @@ -0,0 +1,33 @@ +// Package systemtenant contains function to mark a context as allowed to +// access shards across all tenants. This must only be used for tasks that are +// not request specific. +package systemtenant + +import ( + "context" + "fmt" + + "github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype" +) + +type contextKey int + +const systemTenantKey contextKey = iota + +// With marks a ctx to be allowed to access shards across all tenants. This MUST +// NOT BE USED on the user request path. +func With(ctx context.Context) (context.Context, error) { + // We don't want to allow setting the system tenant on a context that already + // has a user tenant set. + if _, ok := tenanttype.GetTenant(ctx); ok { + return nil, fmt.Errorf("tenant context already set") + } + return context.WithValue(ctx, systemTenantKey, systemTenantKey), nil +} + +// 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 +} diff --git a/internal/tenant/systemtenant/systemtenant_test.go b/internal/tenant/systemtenant/systemtenant_test.go new file mode 100644 index 000000000..d83761551 --- /dev/null +++ b/internal/tenant/systemtenant/systemtenant_test.go @@ -0,0 +1,17 @@ +package systemtenant + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSystemtenantRoundtrip(t *testing.T) { + if Is(context.Background()) { + t.Fatal() + } + ctx, err := With(context.Background()) + require.NoError(t, err) + require.True(t, Is(ctx)) +} diff --git a/shards/shards.go b/shards/shards.go index 309ac17e9..4963cf4b3 100644 --- a/shards/shards.go +++ b/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,11 @@ func (s *shardedSearcher) getLoaded() loaded { func mkRankedShard(s zoekt.Searcher) *rankedShard { q := query.Const{Value: true} - result, err := s.List(context.Background(), &q, nil) + ctx, err := systemtenant.With(context.Background()) + if err != nil { + return &rankedShard{Searcher: s} + } + result, err := s.List(ctx, &q, nil) if err != nil { return &rankedShard{Searcher: s} } From 934161ccfa59f48cf7588e2943392af50584ec60 Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 12:44:04 +0100 Subject: [PATCH 2/4] rename EqualsID, simplify systemtenant --- eval.go | 4 ++-- internal/tenant/query.go | 4 ++-- internal/tenant/systemtenant/systemtenant.go | 16 +++------------- .../tenant/systemtenant/systemtenant_test.go | 4 +--- shards/shards.go | 6 +----- 5 files changed, 9 insertions(+), 25 deletions(-) diff --git a/eval.go b/eval.go index bdc172399..0788fe35d 100644 --- a/eval.go +++ b/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] diff --git a/internal/tenant/query.go b/internal/tenant/query.go index 8a80d989d..3f56cb721 100644 --- a/internal/tenant/query.go +++ b/internal/tenant/query.go @@ -6,9 +6,9 @@ import ( "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 } diff --git a/internal/tenant/systemtenant/systemtenant.go b/internal/tenant/systemtenant/systemtenant.go index 8e1a619fd..46f3e7319 100644 --- a/internal/tenant/systemtenant/systemtenant.go +++ b/internal/tenant/systemtenant/systemtenant.go @@ -5,25 +5,15 @@ package systemtenant import ( "context" - "fmt" - - "github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype" ) type contextKey int const systemTenantKey contextKey = iota -// With marks a ctx to be allowed to access shards across all tenants. This MUST -// NOT BE USED on the user request path. -func With(ctx context.Context) (context.Context, error) { - // We don't want to allow setting the system tenant on a context that already - // has a user tenant set. - if _, ok := tenanttype.GetTenant(ctx); ok { - return nil, fmt.Errorf("tenant context already set") - } - return context.WithValue(ctx, systemTenantKey, systemTenantKey), nil -} +// Ctx is a context that allows queries across all tenants. This must only be +// used for tasks that are not user request specific. +var Ctx = context.WithValue(context.Background(), systemTenantKey, systemTenantKey) // Is returns true if the context has been marked to allow queries across all // tenants. diff --git a/internal/tenant/systemtenant/systemtenant_test.go b/internal/tenant/systemtenant/systemtenant_test.go index d83761551..83ae40c2e 100644 --- a/internal/tenant/systemtenant/systemtenant_test.go +++ b/internal/tenant/systemtenant/systemtenant_test.go @@ -11,7 +11,5 @@ func TestSystemtenantRoundtrip(t *testing.T) { if Is(context.Background()) { t.Fatal() } - ctx, err := With(context.Background()) - require.NoError(t, err) - require.True(t, Is(ctx)) + require.True(t, Is(Ctx)) } diff --git a/shards/shards.go b/shards/shards.go index 4963cf4b3..4d11baaf7 100644 --- a/shards/shards.go +++ b/shards/shards.go @@ -1065,11 +1065,7 @@ func (s *shardedSearcher) getLoaded() loaded { func mkRankedShard(s zoekt.Searcher) *rankedShard { q := query.Const{Value: true} - ctx, err := systemtenant.With(context.Background()) - if err != nil { - return &rankedShard{Searcher: s} - } - result, err := s.List(ctx, &q, nil) + result, err := s.List(systemtenant.Ctx, &q, nil) if err != nil { return &rankedShard{Searcher: s} } From 12a12cdd7d066c471f9a57458b390cb3da552a20 Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 12:50:11 +0100 Subject: [PATCH 3/4] Ctx -> UnafeCtx --- internal/tenant/systemtenant/systemtenant.go | 11 +++++------ internal/tenant/systemtenant/systemtenant_test.go | 2 +- shards/shards.go | 2 +- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/internal/tenant/systemtenant/systemtenant.go b/internal/tenant/systemtenant/systemtenant.go index 46f3e7319..a7957e6b6 100644 --- a/internal/tenant/systemtenant/systemtenant.go +++ b/internal/tenant/systemtenant/systemtenant.go @@ -1,6 +1,5 @@ -// Package systemtenant contains function to mark a context as allowed to -// access shards across all tenants. This must only be used for tasks that are -// not request specific. +// 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 ( @@ -11,9 +10,9 @@ type contextKey int const systemTenantKey contextKey = iota -// Ctx is a context that allows queries across all tenants. This must only be -// used for tasks that are not user request specific. -var Ctx = context.WithValue(context.Background(), systemTenantKey, systemTenantKey) +// 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. diff --git a/internal/tenant/systemtenant/systemtenant_test.go b/internal/tenant/systemtenant/systemtenant_test.go index 83ae40c2e..4330d82c7 100644 --- a/internal/tenant/systemtenant/systemtenant_test.go +++ b/internal/tenant/systemtenant/systemtenant_test.go @@ -11,5 +11,5 @@ func TestSystemtenantRoundtrip(t *testing.T) { if Is(context.Background()) { t.Fatal() } - require.True(t, Is(Ctx)) + require.True(t, Is(UnsafeCtx)) } diff --git a/shards/shards.go b/shards/shards.go index 4d11baaf7..22c53c5cb 100644 --- a/shards/shards.go +++ b/shards/shards.go @@ -1065,7 +1065,7 @@ func (s *shardedSearcher) getLoaded() loaded { func mkRankedShard(s zoekt.Searcher) *rankedShard { q := query.Const{Value: true} - result, err := s.List(systemtenant.Ctx, &q, nil) + result, err := s.List(systemtenant.UnsafeCtx, &q, nil) if err != nil { return &rankedShard{Searcher: s} } From bd5c2ce6db68a74b371303594844059dfb7fdaf2 Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 15:59:11 +0100 Subject: [PATCH 4/4] add comment --- shards/shards.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/shards/shards.go b/shards/shards.go index 22c53c5cb..41df83347 100644 --- a/shards/shards.go +++ b/shards/shards.go @@ -1065,6 +1065,9 @@ 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 + // 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}