Skip to content

Commit d24d2a1

Browse files
nicklaslclaude
andcommitted
feat(go): support encrypted CDN resolver state
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 35d700e commit d24d2a1

12 files changed

Lines changed: 251 additions & 54 deletions

File tree

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package local_resolver
2+
3+
import (
4+
"testing"
5+
6+
"github.com/spotify/confidence-resolver/openfeature-provider/go/confidence/internal/proto/resolver"
7+
"github.com/spotify/confidence-resolver/openfeature-provider/go/confidence/internal/proto/wasm"
8+
tu "github.com/spotify/confidence-resolver/openfeature-provider/go/confidence/internal/testutil"
9+
"google.golang.org/protobuf/types/known/structpb"
10+
)
11+
12+
func TestSetEncryptedResolverState_ResolvesFlags(t *testing.T) {
13+
encryptedState := tu.LoadTestEncryptedState(t)
14+
encryptionKey := tu.LoadTestEncryptionKey(t)
15+
16+
r := NewWasmResolverFactory(NoOpLogSink).New()
17+
defer r.Close(nil)
18+
19+
err := r.SetEncryptedResolverState(&wasm.SetEncryptedResolverStateRequest{
20+
EncryptedState: encryptedState,
21+
EncryptionKey: encryptionKey,
22+
Sdk: &resolver.Sdk{
23+
Sdk: &resolver.Sdk_Id{Id: resolver.SdkId_SDK_ID_GO_LOCAL_PROVIDER},
24+
Version: "0.0.0-test",
25+
},
26+
})
27+
if err != nil {
28+
t.Fatalf("SetEncryptedResolverState failed: %v", err)
29+
}
30+
31+
ctx, _ := structpb.NewStruct(map[string]interface{}{
32+
"targeting_key": "tutorial_visitor",
33+
"visitor_id": "tutorial_visitor",
34+
})
35+
resp, err := r.ResolveProcess(&wasm.ResolveProcessRequest{
36+
Resolve: &wasm.ResolveProcessRequest_DeferredMaterializations{
37+
DeferredMaterializations: &resolver.ResolveFlagsRequest{
38+
Flags: []string{"flags/tutorial-feature"},
39+
ClientSecret: "mkjJruAATQWjeY7foFIWfVAcBWnci2YF",
40+
Apply: true,
41+
EvaluationContext: ctx,
42+
},
43+
},
44+
})
45+
if err != nil {
46+
t.Fatalf("ResolveProcess failed: %v", err)
47+
}
48+
49+
resolved := resp.GetResolved()
50+
if resolved == nil {
51+
t.Fatal("Expected resolved response")
52+
}
53+
flags := resolved.GetResponse().GetResolvedFlags()
54+
if len(flags) != 1 {
55+
t.Fatalf("Expected 1 resolved flag, got %d", len(flags))
56+
}
57+
if flags[0].Reason != resolver.ResolveReason_RESOLVE_REASON_MATCH {
58+
t.Errorf("Expected MATCH, got %v", flags[0].Reason)
59+
}
60+
}
61+
62+
func TestSetEncryptedResolverState_RejectsWrongKey(t *testing.T) {
63+
encryptedState := tu.LoadTestEncryptedState(t)
64+
65+
r := NewWasmResolverFactory(NoOpLogSink).New()
66+
defer r.Close(nil)
67+
68+
err := r.SetEncryptedResolverState(&wasm.SetEncryptedResolverStateRequest{
69+
EncryptedState: encryptedState,
70+
EncryptionKey: make([]byte, 32),
71+
})
72+
if err == nil {
73+
t.Fatal("Expected error with wrong key")
74+
}
75+
}

openfeature-provider/go/confidence/internal/local_resolver/local_resolver.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ type LocalResolverFactory interface {
1919

2020
type LocalResolver interface {
2121
SetResolverState(*wasm.SetResolverStateRequest) error
22+
SetEncryptedResolverState(*wasm.SetEncryptedResolverStateRequest) error
2223
ResolveProcess(*wasm.ResolveProcessRequest) (*wasm.ResolveProcessResponse, error)
2324
RegisterResolve(*wasm.RegisterResolveRequest)
2425
ApplyFlags(*resolver.ApplyFlagsRequest) error

openfeature-provider/go/confidence/internal/local_resolver/pool.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,13 @@ func (s *PooledResolver) SetResolverState(request *wasm.SetResolverStateRequest)
108108
})
109109
}
110110

111+
// SetEncryptedResolverState implements LocalResolver.
112+
func (s *PooledResolver) SetEncryptedResolverState(request *wasm.SetEncryptedResolverStateRequest) error {
113+
return s.maintenance(func(lr LocalResolver) error {
114+
return lr.SetEncryptedResolverState(request)
115+
})
116+
}
117+
111118
// FlushAllLogs implements LocalResolver.
112119
func (s *PooledResolver) FlushAllLogs() error {
113120
return s.maintenance(func(lr LocalResolver) error {

openfeature-provider/go/confidence/internal/local_resolver/recover.go

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,9 @@ type RecoveringResolver struct {
3939
current atomic.Value // holds LocalResolver
4040
broken atomic.Bool // indicates an instance has panicked
4141

42-
lastState atomic.Value // holds *wasm.SetResolverStateRequest
42+
lastState atomic.Value // holds *wasm.SetResolverStateRequest
43+
lastEncryptedState atomic.Value // holds *wasm.SetEncryptedResolverStateRequest
44+
useEncrypted atomic.Bool
4345
}
4446

4547
func (r *RecoveringResolver) get() LocalResolver {
@@ -60,9 +62,12 @@ func (r *RecoveringResolver) startRecreate() {
6062
}()
6163
old := r.get()
6264
newLR := r.factory.New()
63-
if v := r.lastState.Load(); v != nil {
64-
state := v.(*wasm.SetResolverStateRequest)
65-
_ = newLR.SetResolverState(state)
65+
if r.useEncrypted.Load() {
66+
if v := r.lastEncryptedState.Load(); v != nil {
67+
_ = newLR.SetEncryptedResolverState(v.(*wasm.SetEncryptedResolverStateRequest))
68+
}
69+
} else if v := r.lastState.Load(); v != nil {
70+
_ = newLR.SetResolverState(v.(*wasm.SetResolverStateRequest))
6671
}
6772
r.current.Store(newLR)
6873
if old != nil {
@@ -95,6 +100,18 @@ func (r *RecoveringResolver) SetResolverState(request *wasm.SetResolverStateRequ
95100
err = lr.SetResolverState(request)
96101
if err == nil {
97102
r.lastState.Store(request)
103+
r.useEncrypted.Store(false)
104+
}
105+
})
106+
return
107+
}
108+
109+
func (r *RecoveringResolver) SetEncryptedResolverState(request *wasm.SetEncryptedResolverStateRequest) (err error) {
110+
r.withRecover("SetEncryptedResolverState", &err, func(lr LocalResolver) {
111+
err = lr.SetEncryptedResolverState(request)
112+
if err == nil {
113+
r.lastEncryptedState.Store(request)
114+
r.useEncrypted.Store(true)
98115
}
99116
})
100117
return

openfeature-provider/go/confidence/internal/local_resolver/recover_test.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ type mockResolver struct {
1717
}
1818

1919
func (m *mockResolver) SetResolverState(*wasm.SetResolverStateRequest) error { return nil }
20+
func (m *mockResolver) SetEncryptedResolverState(*wasm.SetEncryptedResolverStateRequest) error {
21+
return nil
22+
}
2023
func (m *mockResolver) ResolveProcess(*wasm.ResolveProcessRequest) (*wasm.ResolveProcessResponse, error) {
2124
if m.shouldPanic {
2225
panic("simulated WASM crash")

openfeature-provider/go/confidence/internal/local_resolver/wasm.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,10 @@ func (r *WasmResolver) SetResolverState(request *wasm.SetResolverStateRequest) e
6262
return r.call("wasm_msg_guest_set_resolver_state", request, nil)
6363
}
6464

65+
func (r *WasmResolver) SetEncryptedResolverState(request *wasm.SetEncryptedResolverStateRequest) error {
66+
return r.call("wasm_msg_guest_set_encrypted_resolver_state", request, nil)
67+
}
68+
6569
func (r *WasmResolver) ResolveProcess(request *wasm.ResolveProcessRequest) (*wasm.ResolveProcessResponse, error) {
6670
resp := &wasm.ResolveProcessResponse{}
6771
err := r.call("wasm_msg_guest_resolve_flags", request, resp)

openfeature-provider/go/confidence/internal/testutil/helpers.go

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package testutil
22

33
import (
44
"context"
5+
"encoding/hex"
56
"encoding/json"
67
"os"
78
"path/filepath"
@@ -74,6 +75,28 @@ func LoadTestAccountID(t *testing.T) string {
7475
return strings.TrimSpace(string(data))
7576
}
7677

78+
func LoadTestEncryptedState(t *testing.T) []byte {
79+
dataPath := filepath.Join(repoRoot, "data", "resolver_state_encrypted.pb")
80+
data, err := os.ReadFile(dataPath)
81+
if err != nil {
82+
t.Skipf("Skipping test - could not load encrypted test state: %v", err)
83+
}
84+
return data
85+
}
86+
87+
func LoadTestEncryptionKey(t *testing.T) []byte {
88+
dataPath := filepath.Join(repoRoot, "data", "encryption_key_test.hex")
89+
hexKey, err := os.ReadFile(dataPath)
90+
if err != nil {
91+
t.Skipf("Skipping test - could not load test encryption key: %v", err)
92+
}
93+
key, err := hex.DecodeString(strings.TrimSpace(string(hexKey)))
94+
if err != nil {
95+
t.Fatalf("Invalid hex encryption key: %v", err)
96+
}
97+
return key
98+
}
99+
77100
// CreateStateWithMaterializedSegment creates a test state with a flag using MaterializedSegmentCriterion
78101
func CreateStateWithMaterializedSegment() []byte {
79102
clientName := "clients/test-client"
@@ -428,9 +451,12 @@ func (m *MockedLocalResolver) ResolveProcess(*wasm.ResolveProcessRequest) (*wasm
428451
return m.Response, m.Err
429452
}
430453
func (m MockedLocalResolver) SetResolverState(*wasm.SetResolverStateRequest) error { return nil }
431-
func (m MockedLocalResolver) PrometheusSnapshot(_ uint32, _ bool) string { return "" }
432-
func (m MockedLocalResolver) RegisterResolve(*wasm.RegisterResolveRequest) {}
433-
func (m MockedLocalResolver) ApplyFlags(*resolver.ApplyFlagsRequest) error { return nil }
454+
func (m MockedLocalResolver) SetEncryptedResolverState(*wasm.SetEncryptedResolverStateRequest) error {
455+
return nil
456+
}
457+
func (m MockedLocalResolver) PrometheusSnapshot(_ uint32, _ bool) string { return "" }
458+
func (m MockedLocalResolver) RegisterResolve(*wasm.RegisterResolveRequest) {}
459+
func (m MockedLocalResolver) ApplyFlags(*resolver.ApplyFlagsRequest) error { return nil }
434460

435461
func MustJSONToProto(jsonString string) *structpb.Value {
436462
var v structpb.Value

openfeature-provider/go/confidence/materialization.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,10 @@ func (m *materializationSupportedResolver) SetResolverState(request *wasm.SetRes
189189
return m.current.SetResolverState(request)
190190
}
191191

192+
func (m *materializationSupportedResolver) SetEncryptedResolverState(request *wasm.SetEncryptedResolverStateRequest) (err error) {
193+
return m.current.SetEncryptedResolverState(request)
194+
}
195+
192196
func (m *materializationSupportedResolver) PrometheusSnapshot(bucketsPerDecade uint32, openmetrics bool) string {
193197
return m.current.PrometheusSnapshot(bucketsPerDecade, openmetrics)
194198
}

openfeature-provider/go/confidence/provider.go

Lines changed: 46 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package confidence
22

33
import (
44
"context"
5+
"encoding/hex"
56
"errors"
67
"fmt"
78
"log/slog"
@@ -32,6 +33,7 @@ type Option func(*providerOptions)
3233
type providerOptions struct {
3334
statePollInterval time.Duration
3435
logPollInterval time.Duration
36+
encryptionKey string
3537
}
3638

3739
// WithStatePollInterval sets the interval for polling state updates
@@ -48,6 +50,13 @@ func WithLogPollInterval(d time.Duration) Option {
4850
}
4951
}
5052

53+
// WithEncryptionKey sets the hex-encoded AES-256 encryption key for decrypting CDN state.
54+
func WithEncryptionKey(key string) Option {
55+
return func(o *providerOptions) {
56+
o.encryptionKey = key
57+
}
58+
}
59+
5160
// LocalResolverProvider implements the OpenFeature FeatureProvider interface
5261
// for local flag resolution using the Confidence WASM resolver
5362
type LocalResolverProvider struct {
@@ -56,6 +65,7 @@ type LocalResolverProvider struct {
5665
stateProvider StateProvider
5766
flagLogger FlagLogger
5867
clientSecret string
68+
encryptionKey string
5969
logger *slog.Logger
6070
cancelFunc context.CancelFunc
6171
wg sync.WaitGroup
@@ -107,6 +117,7 @@ func NewLocalResolverProvider(
107117
stateProvider: stateProvider,
108118
flagLogger: flagLogger,
109119
clientSecret: clientSecret,
120+
encryptionKey: options.encryptionKey,
110121
logger: logger,
111122
statePollInterval: statePollInterval,
112123
logPollInterval: logPollInterval,
@@ -478,21 +489,12 @@ func (p *LocalResolverProvider) Init(evaluationContext openfeature.EvaluationCon
478489
return fmt.Errorf("failed to fetch initial state: %w", err)
479490
}
480491

481-
if accountId == "" {
482-
p.logger.Error("AccountID is empty in the fetched state, this should not happen")
483-
return fmt.Errorf("AccountID is empty in the initial state")
492+
sdk := &resolvertypes.Sdk{
493+
Sdk: &resolvertypes.Sdk_Id{Id: resolvertypes.SdkId_SDK_ID_GO_LOCAL_PROVIDER},
494+
Version: Version,
484495
}
485496

486-
// Update resolver with initial state (triggers WASM compilation and initialization)
487-
setResolverStateRequest := &wasm.SetResolverStateRequest{
488-
State: initialState,
489-
AccountId: accountId,
490-
Sdk: &resolvertypes.Sdk{
491-
Sdk: &resolvertypes.Sdk_Id{Id: resolvertypes.SdkId_SDK_ID_GO_LOCAL_PROVIDER},
492-
Version: Version,
493-
},
494-
}
495-
if err := p.resolver.SetResolverState(setResolverStateRequest); err != nil {
497+
if err := p.setResolverState(initialState, accountId, sdk); err != nil {
496498
p.logger.Error("Failed to initialize resolver with initial state", "error", err)
497499
return fmt.Errorf("failed to initialize resolver: %w", err)
498500
}
@@ -504,6 +506,33 @@ func (p *LocalResolverProvider) Init(evaluationContext openfeature.EvaluationCon
504506
return nil
505507
}
506508

509+
func (p *LocalResolverProvider) setResolverState(state []byte, accountId string, sdk *resolvertypes.Sdk) error {
510+
if p.encryptionKey != "" {
511+
if fetcher, ok := p.stateProvider.(*FlagsAdminStateFetcher); ok {
512+
if cdnBytes := fetcher.RawCdnBytes(); cdnBytes != nil {
513+
keyBytes, err := hex.DecodeString(p.encryptionKey)
514+
if err != nil {
515+
return fmt.Errorf("invalid encryption key: %w", err)
516+
}
517+
return p.resolver.SetEncryptedResolverState(&wasm.SetEncryptedResolverStateRequest{
518+
EncryptedState: cdnBytes,
519+
EncryptionKey: keyBytes,
520+
Sdk: sdk,
521+
})
522+
}
523+
}
524+
}
525+
526+
if accountId == "" {
527+
return fmt.Errorf("AccountID is empty in the state")
528+
}
529+
return p.resolver.SetResolverState(&wasm.SetResolverStateRequest{
530+
State: state,
531+
AccountId: accountId,
532+
Sdk: sdk,
533+
})
534+
}
535+
507536
// Shutdown closes the provider and cleans up resources (part of StateHandler interface)
508537
func (p *LocalResolverProvider) Shutdown() {
509538
ctx := context.Background()
@@ -566,28 +595,17 @@ func (p *LocalResolverProvider) startScheduledTasks(parentCtx context.Context) {
566595
for {
567596
select {
568597
case <-stateTicker.C:
569-
// Fetch latest state and accountID
570598
state, accountId, err := p.stateProvider.Provide(ctx)
571599
if err != nil {
572600
p.logger.Error("State fetch failed", "error", err)
573601
continue
574602
}
575603

576-
if accountId == "" {
577-
p.logger.Error("AccountID inside fetched state is empty, skipping this state update attempt")
578-
continue
579-
}
580-
581-
// Update state
582-
setResolverStateRequest := &wasm.SetResolverStateRequest{
583-
State: state,
584-
AccountId: accountId,
585-
Sdk: &resolvertypes.Sdk{
586-
Sdk: &resolvertypes.Sdk_Id{Id: resolvertypes.SdkId_SDK_ID_GO_LOCAL_PROVIDER},
587-
Version: Version,
588-
},
604+
sdk := &resolvertypes.Sdk{
605+
Sdk: &resolvertypes.Sdk_Id{Id: resolvertypes.SdkId_SDK_ID_GO_LOCAL_PROVIDER},
606+
Version: Version,
589607
}
590-
if err := p.resolver.SetResolverState(setResolverStateRequest); err != nil {
608+
if err := p.setResolverState(state, accountId, sdk); err != nil {
591609
p.logger.Error("Failed to update state", "error", err)
592610
}
593611
case <-ctx.Done():

openfeature-provider/go/confidence/provider_builder.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ const confidenceDomain = "edge-grpc.spotify.com"
1919

2020
type ProviderConfig struct {
2121
ClientSecret string
22+
EncryptionKey string // Optional: hex-encoded AES-256 key for decrypting CDN state
2223
Logger *slog.Logger
2324
TransportHooks TransportHooks // Optional: defaults to DefaultTransportHooks
2425
MaterializationStore MaterializationStore // Optional
@@ -72,7 +73,7 @@ func NewProvider(ctx context.Context, config ProviderConfig) (*LocalResolverProv
7273
flagLoggerService := resolverv1.NewInternalFlagLoggerServiceClient(conn)
7374
// Build HTTP transport using hooks and pass into state fetcher
7475
transport := hooks.WrapHTTP(http.DefaultTransport)
75-
stateProvider := NewFlagsAdminStateFetcherWithTransport(config.ClientSecret, logger, transport)
76+
stateProvider := NewFlagsAdminStateFetcherWithEncryption(config.ClientSecret, config.EncryptionKey, logger, transport)
7677
flagLogger := fl.NewGrpcWasmFlagLogger(flagLoggerService, config.ClientSecret, logger)
7778
materializationStore := config.MaterializationStore
7879
if materializationStore == nil {
@@ -87,6 +88,9 @@ func NewProvider(ctx context.Context, config ProviderConfig) (*LocalResolverProv
8788
}
8889
resolverSupplierWithMaterialization := wrapResolverSupplierWithMaterializations(resolverSupplier, materializationStore)
8990
providerOpts := buildProviderOptions(config.StatePollInterval, config.LogPollInterval)
91+
if config.EncryptionKey != "" {
92+
providerOpts = append(providerOpts, WithEncryptionKey(config.EncryptionKey))
93+
}
9094
provider := NewLocalResolverProvider(resolverSupplierWithMaterialization, stateProvider, flagLogger, config.ClientSecret, logger, providerOpts...)
9195
return provider, nil
9296
}

0 commit comments

Comments
 (0)