pyroscope(v2): query-backend tenant isolation check (#4184)

* pyroscope(v2): query-backend tenant isolation check

* asserting error type

* remove unresolved conflict

* fix tests

* skip invalid, log and metric

* just filter out invalid datasets

* improved metric, rename function, remove allocation, check tenant size

* metric rename

Author: alsoba13

pkg/experiment/query_backend/block_reader.go

@@ -3,16 +3,22 @@ package query_backend
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/grafana/dskit/multierror"
+	"github.com/grafana/dskit/tracing"
 	"github.com/opentracing/opentracing-go"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
+	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 	"github.com/grafana/pyroscope/pkg/experiment/block"
 	"github.com/grafana/pyroscope/pkg/objstore"
@@ -41,6 +47,8 @@ type BlockReader struct {
 	log     log.Logger
 	storage objstore.Bucket
 
+	metrics *metrics
+
 	// TODO:
 	//  - Use a worker pool instead of the errgroup.
 	//  - Reusable query context.
@@ -49,10 +57,11 @@ type BlockReader struct {
 	//    Instead, they should share the processing pipeline, if possible.
 }
 
-func NewBlockReader(logger log.Logger, storage objstore.Bucket) *BlockReader {
+func NewBlockReader(logger log.Logger, storage objstore.Bucket, reg prometheus.Registerer) *BlockReader {
 	return &BlockReader{
 		log:     logger,
 		storage: storage,
+		metrics: newMetrics(reg),
 	}
 }
 
@@ -71,7 +80,21 @@ func (b *BlockReader) Invoke(
 	g, ctx := errgroup.WithContext(ctx)
 	agg := newAggregator(req)
 
+	tenantMap := make(map[string]struct{})
+	for _, tenant := range req.Tenant {
+		tenantMap[tenant] = struct{}{}
+	}
+
 	for _, md := range req.QueryPlan.Root.Blocks {
+		md.Datasets, err = filterNotOwnedDatasets(md, tenantMap)
+		if err != nil {
+			b.metrics.datasetTenantIsolationFailure.Inc()
+			traceId, _ := tracing.ExtractTraceID(ctx)
+			level.Error(b.log).Log("msg", "trying to query datasets of other tenants", "valid-tenant", strings.Join(req.Tenant, ","), "block", md.Id, "err", err, "traceId", traceId)
+		}
+		if len(md.Datasets) == 0 {
+			continue
+		}
 		obj := block.NewObject(b.storage, md)
 		g.Go(util.RecoverPanic((&blockContext{
 			ctx: ctx,
@@ -120,6 +143,9 @@ func validateRequest(req *queryv1.InvokeRequest) (*request, error) {
 	if req.QueryPlan == nil || len(req.QueryPlan.Root.Blocks) == 0 {
 		return nil, fmt.Errorf("no blocks to query")
 	}
+	if len(req.Tenant) == 0 {
+		return nil, fmt.Errorf("no tenant provided")
+	}
 	matchers, err := parser.ParseMetricSelector(req.LabelSelector)
 	if err != nil {
 		return nil, fmt.Errorf("label selection is invalid: %w", err)
@@ -132,3 +158,20 @@ func validateRequest(req *queryv1.InvokeRequest) (*request, error) {
 	}
 	return &r, nil
 }
+
+// While the metastore is expected to already filter datasets of other tenants, we do an additional check to avoid
+// processing blocks or datasets belonging to the wrong tenant.
+func filterNotOwnedDatasets(b *metastorev1.BlockMeta, tenantMap map[string]struct{}) ([]*metastorev1.Dataset, error) {
+	errs := multierror.New()
+	datasets := make([]*metastorev1.Dataset, 0)
+	for _, dataset := range b.Datasets {
+		datasetTenant := b.StringTable[dataset.Tenant]
+		_, ok := tenantMap[datasetTenant]
+		if ok {
+			datasets = append(datasets, dataset)
+		} else {
+			errs.Add(fmt.Errorf(`dataset "%s" belongs to tenant "%s"`, b.StringTable[dataset.Name], datasetTenant))
+		}
+	}
+	return datasets, errs.Err()
+}

pkg/experiment/query_backend/block_reader_test.go

@@ -37,6 +37,7 @@ type testSuite struct {
 	reader *BlockReader
 	meta   []*metastorev1.BlockMeta
 	plan   *queryv1.QueryPlan
+	tenant []string
 }
 
 func (s *testSuite) SetupSuite() {
@@ -47,13 +48,19 @@ func (s *testSuite) SetupSuite() {
 func (s *testSuite) SetupTest() {
 	s.ctx = context.Background()
 	s.logger = test.NewTestingLogger(s.T())
-	s.reader = NewBlockReader(s.logger, &objstore.ReaderAtBucket{Bucket: s.bucket})
+	s.reader = NewBlockReader(s.logger, &objstore.ReaderAtBucket{Bucket: s.bucket}, nil)
 	s.meta = make([]*metastorev1.BlockMeta, len(s.blocks))
 	for i, b := range s.blocks {
 		s.meta[i] = b.CloneVT()
 	}
 	s.sanitizeMetadata()
 	s.plan = query_plan.Build(s.meta, 10, 10)
+	s.tenant = make([]string, 0)
+	for _, b := range s.plan.Root.Blocks {
+		for _, d := range b.Datasets {
+			s.tenant = append(s.tenant, b.StringTable[d.Tenant])
+		}
+	}
 }
 
 func (s *testSuite) loadFromDir(dir string) {
@@ -119,6 +126,7 @@ func (s *testSuite) Test_QueryTree_All() {
 			QueryType: queryv1.QueryType_QUERY_TREE,
 			Tree:      &queryv1.TreeQuery{MaxNodes: 16},
 		}},
+		Tenant: s.tenant,
 	})
 
 	s.Require().NoError(err)
@@ -142,6 +150,7 @@ func (s *testSuite) Test_QueryTree_Filter() {
 			QueryType: queryv1.QueryType_QUERY_TREE,
 			Tree:      &queryv1.TreeQuery{MaxNodes: 16},
 		}},
+		Tenant: s.tenant,
 	})
 
 	s.Require().NoError(err)
@@ -163,6 +172,7 @@ func (s *testSuite) Test_QueryPprof_Metadata() {
 			QueryType: queryv1.QueryType_QUERY_PPROF,
 			Pprof:     &queryv1.PprofQuery{},
 		}},
+		Tenant: s.tenant,
 	})
 
 	s.Require().NoError(err)
@@ -193,6 +203,7 @@ func (s *testSuite) Test_DatasetIndex_SeriesLabels_GroupBy() {
 				LabelNames: []string{"service_name", "__profile_type__"},
 			},
 		}},
+		Tenant: s.tenant,
 	})
 
 	s.Require().NoError(err)
@@ -215,6 +226,7 @@ func (s *testSuite) Test_SeriesLabels() {
 			QueryType:    queryv1.QueryType_QUERY_SERIES_LABELS,
 			SeriesLabels: &queryv1.SeriesLabelsQuery{},
 		}},
+		Tenant: s.tenant,
 	})
 
 	s.Require().NoError(err)
@@ -242,6 +254,7 @@ func (s *testSuite) Test_QueryTimeSeries() {
 		Query:         []*queryv1.Query{query},
 		QueryPlan:     s.plan,
 		LabelSelector: "{}",
+		Tenant:        s.tenant,
 	}
 
 	resp, err := s.reader.Invoke(s.ctx, req)
@@ -250,3 +263,24 @@ func (s *testSuite) Test_QueryTimeSeries() {
 	s.Require().Len(resp.Reports, 1)
 	s.Require().NotNil(resp.Reports[0].TimeSeries)
 }
+
+func (s *testSuite) Test_QueryTree_All_Tenant_Isolation() {
+	queryTenant := "some-tenant"
+
+	s.Require().NotContains(s.tenant, queryTenant)
+
+	resp, err := s.reader.Invoke(s.ctx, &queryv1.InvokeRequest{
+		EndTime:       time.Now().UnixMilli(),
+		LabelSelector: "{}",
+		QueryPlan:     s.plan,
+		Query: []*queryv1.Query{{
+			QueryType: queryv1.QueryType_QUERY_TREE,
+			Tree:      &queryv1.TreeQuery{MaxNodes: 16},
+		}},
+		Tenant: []string{queryTenant},
+	})
+
+	s.Require().NoError(err)
+	s.Require().NotNil(resp)
+	s.Require().Len(resp.Reports, 0)
+}

pkg/experiment/query_backend/metrics.go

@@ -0,0 +1,22 @@
+package query_backend
+
+import "github.com/prometheus/client_golang/prometheus"
+
+type metrics struct {
+	datasetTenantIsolationFailure prometheus.Counter
+}
+
+func newMetrics(reg prometheus.Registerer) *metrics {
+	m := &metrics{
+		datasetTenantIsolationFailure: prometheus.NewCounter(
+			prometheus.CounterOpts{
+				Namespace: "pyroscope",
+				Subsystem: "query_backend",
+				Name:      "dataset_tenant_isolation_failure_total",
+			}),
+	}
+	if reg != nil {
+		reg.MustRegister(m.datasetTenantIsolationFailure)
+	}
+	return m
+}

pkg/phlare/modules_experimental.go

@@ -329,7 +329,7 @@ func (f *Phlare) initQueryBackend() (services.Service, error) {
 		logger,
 		f.reg,
 		f.queryBackendClient,
-		querybackend.NewBlockReader(f.logger, f.storageBucket),
+		querybackend.NewBlockReader(f.logger, f.storageBucket, f.reg),
 	)
 	if err != nil {
 		return nil, err