fix: handle invalid selector errors

core/pkg/store/query.go

@@ -33,14 +33,23 @@ type Selector struct {
 	indexMap map[string]string
 }
 
+var validSelectorKeys = map[string]struct{}{
+	flagSetIdIndex: {},
+	sourceIndex:    {},
+}
+
 // NewSelector creates a new Selector from a selector expression string.
 // #1708 Until we decide on the Selector syntax, only a single key=value pair is supported
 // For example, to select flags from source "./mySource" or flagSetId "1234", use the expressions:
 // "source=./mySource" or "flagSetId=1234"
-func NewSelector(selectorExpression string) Selector {
-	return Selector{
-		indexMap: expressionToMap(selectorExpression),
+func NewSelector(selectorExpression string) (Selector, error) {
+	m := expressionToMap(selectorExpression)
+	for key := range m {
+		if _, ok := validSelectorKeys[key]; !ok {
+			return Selector{}, fmt.Errorf("invalid selector key %q, valid keys: %q, %q", key, flagSetIdIndex, sourceIndex)
+		}
 	}
+	return Selector{indexMap: m}, nil
 }
 
 func expressionToMap(sExp string) map[string]string {
@@ -70,13 +79,17 @@ func expressionToMap(sExp string) map[string]string {
 	return selectorMap
 }
 
-// WithIndex creates a new Selector from the current Selector and adds the given key-value-pair
-func (s Selector) WithIndex(key string, value string) Selector {
+func (s Selector) WithSource(source string) Selector    { return s.withIndex(sourceIndex, source) }
+func (s Selector) WithFlagSetId(id string) Selector     { return s.withIndex(flagSetIdIndex, id) }
+func (s Selector) withKey(key string) Selector          { return s.withIndex(keyIndex, key) }
+
+func (s Selector) withIndex(key, value string) Selector {
 	m := maps.Clone(s.indexMap)
-	m[key] = value
-	return Selector{
-		indexMap: m,
+	if m == nil {
+		m = make(map[string]string, 1)
 	}
+	m[key] = value
+	return Selector{indexMap: m}
 }
 
 func (s *Selector) IsEmpty() bool {

core/pkg/store/query_test.go

@@ -50,22 +50,27 @@ func TestSelector_IsEmpty(t *testing.T) {
 	}
 }
 
-func TestSelector_WithIndex(t *testing.T) {
-	oldS := Selector{indexMap: map[string]string{"source": "abc"}}
-	newS := oldS.WithIndex("flagSetId", "1234")
+func TestSelector_WithSourceAndFlagSetId(t *testing.T) {
+	s := Selector{}.WithSource("abc")
+	if s.indexMap[sourceIndex] != "abc" {
+		t.Errorf("WithSource did not set source")
+	}
 
-	if newS.indexMap["source"] != "abc" {
-		t.Errorf("WithIndex did not preserve existing keys")
+	s2 := s.WithFlagSetId("1234")
+	if s2.indexMap[sourceIndex] != "abc" {
+		t.Errorf("WithFlagSetId did not preserve source")
 	}
-	if newS.indexMap["flagSetId"] != "1234" {
-		t.Errorf("WithIndex did not add new key")
+	if s2.indexMap[flagSetIdIndex] != "1234" {
+		t.Errorf("WithFlagSetId did not set flagSetId")
 	}
+
 	// Ensure original is unchanged
-	if _, ok := oldS.indexMap["flagSetId"]; ok {
-		t.Errorf("WithIndex mutated original selector")
+	if _, ok := s.indexMap[flagSetIdIndex]; ok {
+		t.Errorf("WithFlagSetId mutated original selector")
 	}
 }
 
+
 func TestSelector_ToQuery(t *testing.T) {
 	tests := []struct {
 		name       string
@@ -176,6 +181,7 @@ func TestNewSelector(t *testing.T) {
 		name    string
 		input   string
 		wantMap map[string]string
+		wantErr bool
 	}{
 		// #1708 Until we decide on the Selector syntax, only a single key=value pair is supported
 		/*
@@ -205,12 +211,21 @@ func TestNewSelector(t *testing.T) {
 			input:   "",
 			wantMap: map[string]string{},
 		},
+		{
+			name:    "invalid key",
+			input:   "flagSetIds=abc",
+			wantErr: true,
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			s := NewSelector(tt.input)
-			if !reflect.DeepEqual(s.indexMap, tt.wantMap) {
+			s, err := NewSelector(tt.input)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("NewSelector(%q) error = %v, wantErr %v", tt.input, err, tt.wantErr)
+				return
+			}
+			if !tt.wantErr && !reflect.DeepEqual(s.indexMap, tt.wantMap) {
 				t.Errorf("NewSelector(%q) indexMap = %v, want %v", tt.input, s.indexMap, tt.wantMap)
 			}
 		})

core/pkg/store/store.go

@@ -153,7 +153,7 @@ func (s *Store) Get(_ context.Context, key string, selector *Selector) (model.Fl
 
 	// if present, use the selector to query the flags
 	if !selector.IsEmpty() {
-		selector := selector.WithIndex("key", key)
+		selector := selector.withKey(key)
 		indexId, constraints := selector.ToQuery()
 		s.logger.Debug(fmt.Sprintf("getting flag with query: %s, %v", indexId, constraints))
 		raw, err := txn.First(flagsTable, indexId, constraints...)
@@ -274,7 +274,7 @@ func (s *Store) Update(
 			seenFlagSetIds[id.flagSetId] = struct{}{}
 		}
 		for fsi := range seenFlagSetIds {
-			sel := NewSelector(flagSetIdIndex+"="+fsi).WithIndex(sourceIndex, source)
+			sel := Selector{}.WithFlagSetId(fsi).WithSource(source)
 			indexId, constraints := sel.ToQuery()
 			it, err := txn.Get(flagsTable, indexId, constraints...)
 			if err != nil {
@@ -284,7 +284,7 @@ func (s *Store) Update(
 			oldFlags = append(oldFlags, s.collect(it)...)
 		}
 	} else {
-		sel := NewSelector(sourceIndex + "=" + source)
+		sel := Selector{}.WithSource(source)
 		indexId, constraints := sel.ToQuery()
 		it, err := txn.Get(flagsTable, indexId, constraints...)
 		if err != nil {

core/pkg/store/store_test.go

@@ -238,8 +238,8 @@ func TestGet(t *testing.T) {
 
 	sources := []string{sourceA.Name, sourceB.Name, sourceC.Name}
 
-	sourceASelector := NewSelector("source=" + sourceA.Name)
-	flagSetIdCSelector := NewSelector("flagSetId=" + flagSetIdC)
+	sourceASelector := Selector{}.WithSource(sourceA.Name)
+	flagSetIdCSelector := Selector{}.WithFlagSetId(flagSetIdC)
 
 	t.Parallel()
 	tests := []struct {
@@ -377,10 +377,10 @@ func TestGetAllNoWatcher(t *testing.T) {
 
 	sources := []string{sourceA.Name, sourceB.Name, sourceC.Name}
 
-	sourceASelector := NewSelector("source=" + sourceA.Name)
-	flagSetIdCSelector := NewSelector("flagSetId=" + flagSetIdC)
+	sourceASelector := Selector{}.WithSource(sourceA.Name)
+	flagSetIdCSelector := Selector{}.WithFlagSetId(flagSetIdC)
 	// #1708 Until we decide on the Selector syntax, only a single key=value pair is supported
-	//flagSetIdAndCSelector := NewSelector("flagSetId=" + flagSetIdC + ",source=" + sourceC.Name)
+	//flagSetIdAndCSelector := Selector{}.WithFlagSetId(flagSetIdC).withIndex(sourceIndex, sourceC.Name)
 
 	t.Parallel()
 	tests := []struct {
@@ -506,10 +506,10 @@ func TestWatch(t *testing.T) {
 	pauseTime := 100 * time.Millisecond // time for updates to settle
 	timeout := 1000 * time.Millisecond  // time to make sure we get enough updates, and no extras
 
-	sourceASelector := NewSelector("source=" + sourceA)
-	flagSetIdCSelector := NewSelector("flagSetId=" + myFlagSetId)
-	emptySelector := NewSelector("")
-	sourceCSelector := NewSelector("source=" + sourceC)
+	sourceASelector := Selector{}.WithSource(sourceA)
+	flagSetIdCSelector := Selector{}.WithFlagSetId(myFlagSetId)
+	emptySelector := Selector{}
+	sourceCSelector := Selector{}.WithSource(sourceC)
 
 	tests := []struct {
 		name        string
@@ -786,15 +786,15 @@ func TestQueryMetadata(t *testing.T) {
 	// #1708 Until we decide on the Selector syntax, only a single key=value pair is supported
 	// 		 these tests should then also cover more complex selectors
 
-	selector := NewSelector("flagSetId=" + nonExistingFlagSetId)
+	selector := Selector{}.WithFlagSetId(nonExistingFlagSetId)
 	_, metadata, _ := store.GetAll(context.Background(), &selector)
 	assert.Equal(t, metadata, model.Metadata{"flagSetId": nonExistingFlagSetId}, "metadata did not match expected")
 
-	selector = NewSelector("flagSetId=" + nonExistingFlagSetId)
+	selector = Selector{}.WithFlagSetId(nonExistingFlagSetId)
 	_, metadata, _ = store.Get(context.Background(), "key", &selector)
 	assert.Equal(t, metadata, model.Metadata{"flagSetId": nonExistingFlagSetId}, "metadata did not match expected")
 
-	selector = NewSelector("source=" + otherSource)
+	selector = Selector{}.WithSource(otherSource)
 	_, metadata, _ = store.Get(context.Background(), "key", &selector)
 	assert.Equal(t, metadata, model.Metadata{"source": otherSource}, "metadata did not match expected")
 }

flagd/pkg/service/flag-evaluation/context_utils.go

@@ -2,8 +2,21 @@ package service
 
 import (
 	"net/http"
+
+	"connectrpc.com/connect"
+	flagdService "github.com/open-feature/flagd/flagd/pkg/service"
+	"github.com/open-feature/flagd/core/pkg/store"
 )
 
+func selectorFromHeader(header http.Header) (store.Selector, error) {
+	expr := header.Get(flagdService.FLAGD_SELECTOR_HEADER)
+	s, err := store.NewSelector(expr)
+	if err != nil {
+		return store.Selector{}, connect.NewError(connect.CodeInvalidArgument, err)
+	}
+	return s, nil
+}
+
 // MergeContextsAndHeaders merges evaluation contexts with static context values and header-based context.
 // highest priority > header-context-from-cli > static-context-from-cli > request-context > lowest priority
 // Header names are matched case-insensitively according to HTTP specification.

flagd/pkg/service/flag-evaluation/flag_evaluator.go

@@ -14,7 +14,6 @@ import (
 	"github.com/open-feature/flagd/core/pkg/service"
 	"github.com/open-feature/flagd/core/pkg/store"
 	"github.com/open-feature/flagd/core/pkg/telemetry"
-	flagdService "github.com/open-feature/flagd/flagd/pkg/service"
 	"github.com/rs/xid"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
@@ -75,8 +74,10 @@ func (s *OldFlagEvaluationService) ResolveAll(
 		Flags: make(map[string]*schemaV1.AnyFlag),
 	}
 
-	selectorExpression := req.Header().Get(flagdService.FLAGD_SELECTOR_HEADER)
-	selector := store.NewSelector(selectorExpression)
+	selector, err := selectorFromHeader(req.Header())
+	if err != nil {
+		return nil, err
+	}
 	ctx = context.WithValue(ctx, store.SelectorContextKey{}, selector)
 
 	values, _, err := s.eval.ResolveAllValues(ctx, reqID, mergeContexts(req.Msg.GetContext().AsMap(), s.contextValues, req.Header(), make(map[string]string)))
@@ -142,8 +143,10 @@ func (s *OldFlagEvaluationService) EventStream(
 	s.logger.Debug(fmt.Sprintf("starting event stream for request"))
 
 	requestNotificationChan := make(chan service.Notification, 1)
-	selectorExpression := req.Header().Get(flagdService.FLAGD_SELECTOR_HEADER)
-	selector := store.NewSelector(selectorExpression)
+	selector, err := selectorFromHeader(req.Header())
+	if err != nil {
+		return err
+	}
 	s.eventingConfiguration.Subscribe(ctx, req, &selector, requestNotificationChan)
 	defer s.eventingConfiguration.Unsubscribe(req)
 
@@ -185,11 +188,13 @@ func (s *OldFlagEvaluationService) ResolveBoolean(
 	ctx, span := s.flagEvalTracer.Start(ctx, "resolveBoolean", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 	res := connect.NewResponse(&schemaV1.ResolveBooleanResponse{})
-	selectorExpression := req.Header().Get(flagdService.FLAGD_SELECTOR_HEADER)
-	selector := store.NewSelector(selectorExpression)
+	selector, err := selectorFromHeader(req.Header())
+	if err != nil {
+		return nil, err
+	}
 	ctx = context.WithValue(ctx, store.SelectorContextKey{}, selector)
 
-	err := resolve[bool](
+	err = resolve[bool](
 		ctx,
 		s.logger,
 		s.eval.ResolveBooleanValue,
@@ -217,12 +222,14 @@ func (s *OldFlagEvaluationService) ResolveString(
 	ctx, span := s.flagEvalTracer.Start(ctx, "resolveString", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
-	selectorExpression := req.Header().Get(flagdService.FLAGD_SELECTOR_HEADER)
-	selector := store.NewSelector(selectorExpression)
+	selector, err := selectorFromHeader(req.Header())
+	if err != nil {
+		return nil, err
+	}
 	ctx = context.WithValue(ctx, store.SelectorContextKey{}, selector)
 
 	res := connect.NewResponse(&schemaV1.ResolveStringResponse{})
-	err := resolve[string](
+	err = resolve[string](
 		ctx,
 		s.logger,
 		s.eval.ResolveStringValue,
@@ -250,12 +257,14 @@ func (s *OldFlagEvaluationService) ResolveInt(
 	ctx, span := s.flagEvalTracer.Start(ctx, "resolveInt", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
-	selectorExpression := req.Header().Get(flagdService.FLAGD_SELECTOR_HEADER)
-	selector := store.NewSelector(selectorExpression)
+	selector, err := selectorFromHeader(req.Header())
+	if err != nil {
+		return nil, err
+	}
 	ctx = context.WithValue(ctx, store.SelectorContextKey{}, selector)
 
 	res := connect.NewResponse(&schemaV1.ResolveIntResponse{})
-	err := resolve[int64](
+	err = resolve[int64](
 		ctx,
 		s.logger,
 		s.eval.ResolveIntValue,
@@ -283,12 +292,14 @@ func (s *OldFlagEvaluationService) ResolveFloat(
 	ctx, span := s.flagEvalTracer.Start(ctx, "resolveFloat", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
-	selectorExpression := req.Header().Get(flagdService.FLAGD_SELECTOR_HEADER)
-	selector := store.NewSelector(selectorExpression)
+	selector, err := selectorFromHeader(req.Header())
+	if err != nil {
+		return nil, err
+	}
 	ctx = context.WithValue(ctx, store.SelectorContextKey{}, selector)
 
 	res := connect.NewResponse(&schemaV1.ResolveFloatResponse{})
-	err := resolve[float64](
+	err = resolve[float64](
 		ctx,
 		s.logger,
 		s.eval.ResolveFloatValue,
@@ -316,12 +327,14 @@ func (s *OldFlagEvaluationService) ResolveObject(
 	ctx, span := s.flagEvalTracer.Start(ctx, "resolveObject", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
-	selectorExpression := req.Header().Get(flagdService.FLAGD_SELECTOR_HEADER)
-	selector := store.NewSelector(selectorExpression)
+	selector, err := selectorFromHeader(req.Header())
+	if err != nil {
+		return nil, err
+	}
 	ctx = context.WithValue(ctx, store.SelectorContextKey{}, selector)
 
 	res := connect.NewResponse(&schemaV1.ResolveObjectResponse{})
-	err := resolve[map[string]any](
+	err = resolve[map[string]any](
 		ctx,
 		s.logger,
 		s.eval.ResolveObjectValue,