3pp-mirror/pomerium
1
0
Fork 0
mirror of https://github.com/pomerium/pomerium.git synced 2025-04-28 09:56:31 +02:00
Code Issues Projects Releases Packages Wiki Activity

storage: add sync querier (#5570)

Browse source 
* storage: add fallback querier

* storage: add sync querier

* storage: add typed querier

* use synced querier
This commit is contained in:
Caleb Doxsey 2025-04-23 10:15:48 -06:00 • committed by GitHub
parent e1d84a1dde
commit 8738066ce4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 569 additions and 214 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
authorize/authorize.go
View file

@ -11,7 +11,6 @@ import (
oteltrace "go.opentelemetry.io/otel/trace"
"golang.org/x/sync/errgroup"
"github.com/pomerium/datasource/pkg/directory"
"github.com/pomerium/pomerium/authorize/evaluator"
"github.com/pomerium/pomerium/authorize/internal/store"
"github.com/pomerium/pomerium/config"
@ -20,17 +19,15 @@ import (
"github.com/pomerium/pomerium/internal/telemetry/metrics"
"github.com/pomerium/pomerium/pkg/cryptutil"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/storage"
"github.com/pomerium/pomerium/pkg/telemetry/trace"
)
// Authorize struct holds
type Authorize struct {
state *atomicutil.Value[*authorizeState]
store *store.Store
currentConfig *atomicutil.Value[*config.Config]
accessTracker *AccessTracker
groupsCacheWarmer *cacheWarmer
state *atomicutil.Value[*authorizeState]
store *store.Store
currentConfig *atomicutil.Value[*config.Config]
accessTracker *AccessTracker
tracerProvider oteltrace.TracerProvider
tracer oteltrace.Tracer
@ -48,13 +45,12 @@ func New(ctx context.Context, cfg *config.Config) (*Authorize, error) {
}
a.accessTracker = NewAccessTracker(a, accessTrackerMaxSize, accessTrackerDebouncePeriod)
state, err := newAuthorizeStateFromConfig(ctx, tracerProvider, cfg, a.store, nil)
state, err := newAuthorizeStateFromConfig(ctx, nil, tracerProvider, cfg, a.store)
if err != nil {
return nil, err
}
a.state = atomicutil.NewValue(state)
a.groupsCacheWarmer = newCacheWarmer(state.dataBrokerClientConnection, storage.GlobalCache, directory.GroupRecordType)
return a, nil
}
@ -70,10 +66,6 @@ func (a *Authorize) Run(ctx context.Context) error {
a.accessTracker.Run(ctx)
return nil
})
eg.Go(func() error {
a.groupsCacheWarmer.Run(ctx)
return nil
})
return eg.Wait()
}
@ -154,13 +146,9 @@ func newPolicyEvaluator(
func (a *Authorize) OnConfigChange(ctx context.Context, cfg *config.Config) {
currentState := a.state.Load()
a.currentConfig.Store(cfg)
if newState, err := newAuthorizeStateFromConfig(ctx, a.tracerProvider, cfg, a.store, currentState.evaluator); err != nil {
if newState, err := newAuthorizeStateFromConfig(ctx, currentState, a.tracerProvider, cfg, a.store); err != nil {
log.Ctx(ctx).Error().Err(err).Msg("authorize: error updating state")
} else {
a.state.Store(newState)
if currentState.dataBrokerClientConnection != newState.dataBrokerClientConnection {
a.groupsCacheWarmer.UpdateConn(newState.dataBrokerClientConnection)
}
}
}

122
authorize/cache_warmer.go
View file

@ -1,122 +0,0 @@
package authorize
import (
"context"
"time"
"google.golang.org/grpc"
"github.com/pomerium/pomerium/internal/log"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/storage"
)
type cacheWarmer struct {
cc *grpc.ClientConn
cache storage.Cache
typeURL string
updatedCC chan *grpc.ClientConn
}
func newCacheWarmer(
cc *grpc.ClientConn,
cache storage.Cache,
typeURL string,
) *cacheWarmer {
return &cacheWarmer{
cc: cc,
cache: cache,
typeURL: typeURL,
updatedCC: make(chan *grpc.ClientConn, 1),
}
}
func (cw *cacheWarmer) UpdateConn(cc *grpc.ClientConn) {
for {
select {
case cw.updatedCC <- cc:
return
default:
}
select {
case <-cw.updatedCC:
default:
}
}
}
func (cw *cacheWarmer) Run(ctx context.Context) {
// Run a syncer for the cache warmer until the underlying databroker connection is changed.
// When that happens cancel the currently running syncer and start a new one.
runCtx, runCancel := context.WithCancel(ctx)
go cw.run(runCtx, cw.cc)
for {
select {
case <-ctx.Done():
runCancel()
return
case cc := <-cw.updatedCC:
log.Ctx(ctx).Info().Msg("cache-warmer: received updated databroker client connection, restarting syncer")
cw.cc = cc
runCancel()
runCtx, runCancel = context.WithCancel(ctx)
go cw.run(runCtx, cw.cc)
}
}
}
func (cw *cacheWarmer) run(ctx context.Context, cc *grpc.ClientConn) {
log.Ctx(ctx).Debug().Str("type-url", cw.typeURL).Msg("cache-warmer: running databroker syncer to warm cache")
_ = databroker.NewSyncer(ctx, "cache-warmer", cacheWarmerSyncerHandler{
client: databroker.NewDataBrokerServiceClient(cc),
cache: cw.cache,
}, databroker.WithTypeURL(cw.typeURL)).Run(ctx)
}
type cacheWarmerSyncerHandler struct {
client databroker.DataBrokerServiceClient
cache storage.Cache
}
func (h cacheWarmerSyncerHandler) GetDataBrokerServiceClient() databroker.DataBrokerServiceClient {
return h.client
}
func (h cacheWarmerSyncerHandler) ClearRecords(_ context.Context) {
h.cache.InvalidateAll()
}
func (h cacheWarmerSyncerHandler) UpdateRecords(ctx context.Context, serverVersion uint64, records []*databroker.Record) {
for _, record := range records {
req := &databroker.QueryRequest{
Type: record.Type,
Limit: 1,
}
req.SetFilterByIDOrIndex(record.Id)
res := &databroker.QueryResponse{
Records: []*databroker.Record{record},
TotalCount: 1,
ServerVersion: serverVersion,
RecordVersion: record.Version,
}
expiry := time.Now().Add(time.Hour * 24 * 365)
key, err := storage.MarshalQueryRequest(req)
if err != nil {
log.Ctx(ctx).Error().Err(err).Msg("cache-warmer: failed to marshal query request")
continue
}
value, err := storage.MarshalQueryResponse(res)
if err != nil {
log.Ctx(ctx).Error().Err(err).Msg("cache-warmer: failed to marshal query response")
continue
}
h.cache.Set(expiry, key, value)
}
}

52
authorize/cache_warmer_test.go
View file

@ -1,52 +0,0 @@
package authorize
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.opentelemetry.io/otel/trace/noop"
"google.golang.org/grpc"
"github.com/pomerium/pomerium/internal/databroker"
"github.com/pomerium/pomerium/internal/testutil"
databrokerpb "github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/protoutil"
"github.com/pomerium/pomerium/pkg/storage"
)
func TestCacheWarmer(t *testing.T) {
t.Parallel()
ctx := testutil.GetContext(t, 10*time.Minute)
cc := testutil.NewGRPCServer(t, func(srv *grpc.Server) {
databrokerpb.RegisterDataBrokerServiceServer(srv, databroker.New(ctx, noop.NewTracerProvider()))
})
t.Cleanup(func() { cc.Close() })
client := databrokerpb.NewDataBrokerServiceClient(cc)
_, err := client.Put(ctx, &databrokerpb.PutRequest{
Records: []*databrokerpb.Record{
{Type: "example.com/record", Id: "e1", Data: protoutil.NewAnyBool(true)},
{Type: "example.com/record", Id: "e2", Data: protoutil.NewAnyBool(true)},
},
})
require.NoError(t, err)
cache := storage.NewGlobalCache(time.Minute)
cw := newCacheWarmer(cc, cache, "example.com/record")
go cw.Run(ctx)
assert.Eventually(t, func() bool {
req := &databrokerpb.QueryRequest{
Type: "example.com/record",
Limit: 1,
}
req.SetFilterByIDOrIndex("e1")
res, err := storage.NewCachingQuerier(storage.NewStaticQuerier(), cache).Query(ctx, req)
require.NoError(t, err)
return len(res.GetRecords()) == 1
}, 10*time.Second, time.Millisecond*100)
}

21
authorize/grpc.go
View file

@ -34,11 +34,7 @@ func (a *Authorize) Check(ctx context.Context, in *envoy_service_auth_v3.CheckRe
ctx, span := a.tracer.Start(ctx, "authorize.grpc.Check")
defer span.End()
querier := storage.NewCachingQuerier(
storage.NewQuerier(a.state.Load().dataBrokerClient),
storage.GlobalCache,
)
ctx = storage.WithQuerier(ctx, querier)
ctx = a.withQuerierForCheckRequest(ctx)
state := a.state.Load()
@ -172,6 +168,21 @@ func (a *Authorize) getMatchingPolicy(routeID uint64) *config.Policy {
return nil
}
func (a *Authorize) withQuerierForCheckRequest(ctx context.Context) context.Context {
state := a.state.Load()
q := storage.NewQuerier(state.dataBrokerClient)
// if sync queriers are enabled, use those
if len(state.syncQueriers) > 0 {
m := map[string]storage.Querier{}
for recordType, sq := range state.syncQueriers {
m[recordType] = storage.NewFallbackQuerier(sq, q)
}
q = storage.NewTypedQuerier(q, m)
}
q = storage.NewCachingQuerier(q, storage.GlobalCache)
return storage.WithQuerier(ctx, q)
}
func getHTTPRequestFromCheckRequest(req *envoy_service_auth_v3.CheckRequest) *http.Request {
hattrs := req.GetAttributes().GetRequest().GetHttp()
u := getCheckRequestURL(req)

38
authorize/state.go
View file

@ -9,12 +9,17 @@ import (
oteltrace "go.opentelemetry.io/otel/trace"
googlegrpc "google.golang.org/grpc"
"github.com/pomerium/datasource/pkg/directory"
"github.com/pomerium/pomerium/authorize/evaluator"
"github.com/pomerium/pomerium/authorize/internal/store"
"github.com/pomerium/pomerium/config"
"github.com/pomerium/pomerium/internal/authenticateflow"
"github.com/pomerium/pomerium/pkg/grpc"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/grpc/session"
"github.com/pomerium/pomerium/pkg/grpc/user"
"github.com/pomerium/pomerium/pkg/grpcutil"
"github.com/pomerium/pomerium/pkg/storage"
)
var outboundGRPCConnection = new(grpc.CachedOutboundGRPClientConn)
@ -30,14 +35,15 @@ type authorizeState struct {
dataBrokerClient databroker.DataBrokerServiceClient
sessionStore *config.SessionStore
authenticateFlow authenticateFlow
syncQueriers map[string]storage.Querier
}
func newAuthorizeStateFromConfig(
ctx context.Context,
previousState *authorizeState,
tracerProvider oteltrace.TracerProvider,
cfg *config.Config,
store *store.Store,
previousPolicyEvaluator *evaluator.Evaluator,
) (*authorizeState, error) {
if err := validateOptions(cfg.Options); err != nil {
return nil, fmt.Errorf("authorize: bad options: %w", err)
@ -46,8 +52,12 @@ func newAuthorizeStateFromConfig(
state := new(authorizeState)
var err error
var previousEvaluator *evaluator.Evaluator
if previousState != nil {
previousEvaluator = previousState.evaluator
}
state.evaluator, err = newPolicyEvaluator(ctx, cfg.Options, store, previousPolicyEvaluator)
state.evaluator, err = newPolicyEvaluator(ctx, cfg.Options, store, previousEvaluator)
if err != nil {
return nil, fmt.Errorf("authorize: failed to update policy with options: %w", err)
}
@ -88,5 +98,29 @@ func newAuthorizeStateFromConfig(
return nil, err
}
state.syncQueriers = make(map[string]storage.Querier)
if previousState != nil {
if previousState.dataBrokerClientConnection == state.dataBrokerClientConnection {
state.syncQueriers = previousState.syncQueriers
} else {
for _, v := range previousState.syncQueriers {
v.Stop()
}
}
}
if cfg.Options.IsRuntimeFlagSet(config.RuntimeFlagAuthorizeUseSyncedData) {
for _, recordType := range []string{
grpcutil.GetTypeURL(new(session.Session)),
grpcutil.GetTypeURL(new(user.User)),
grpcutil.GetTypeURL(new(user.ServiceAccount)),
directory.GroupRecordType,
directory.UserRecordType,
} {
if _, ok := state.syncQueriers[recordType]; !ok {
state.syncQueriers[recordType] = storage.NewSyncQuerier(state.dataBrokerClient, recordType)
}
}
}
return state, nil
}

4
config/runtime_flags.go
View file

@ -25,6 +25,10 @@ var (
// RuntimeFlagAddExtraMetricsLabels enables adding extra labels to metrics (host and installation id)
RuntimeFlagAddExtraMetricsLabels = runtimeFlag("add_extra_metrics_labels", true)
// RuntimeFlagAuthorizeUseSyncedData enables synced data for querying the databroker for
// certain types of data.
RuntimeFlagAuthorizeUseSyncedData = runtimeFlag("authorize_use_synced_data", true)
)
// RuntimeFlag is a runtime flag that can flip on/off certain features

23
internal/testenv/selftests/tracing_test.go
View file

@ -81,14 +81,16 @@ func TestOTLPTracing(t *testing.T) {
results := NewTraceResults(srv.FlushResourceSpans())
var (
testEnvironmentLocalTest = fmt.Sprintf("Test Environment: %s", t.Name())
testEnvironmentAuthenticate = "Test Environment: Authenticate"
authenticateOAuth2Client = "Authenticate: OAuth2 Client: GET /.well-known/jwks.json"
idpServerGetUserinfo = "IDP: Server: GET /oidc/userinfo"
idpServerPostToken = "IDP: Server: POST /oidc/token"
controlPlaneEnvoyAccessLogs = "Control Plane: envoy.service.accesslog.v3.AccessLogService/StreamAccessLogs"
controlPlaneEnvoyDiscovery = "Control Plane: envoy.service.discovery.v3.AggregatedDiscoveryService/DeltaAggregatedResources"
controlPlaneExport = "Control Plane: opentelemetry.proto.collector.trace.v1.TraceService/Export"
testEnvironmentLocalTest = fmt.Sprintf("Test Environment: %s", t.Name())
testEnvironmentAuthenticate = "Test Environment: Authenticate"
authenticateOAuth2Client = "Authenticate: OAuth2 Client: GET /.well-known/jwks.json"
authorizeDatabrokerSync = "Authorize: databroker.DataBrokerService/Sync"
authorizeDatabrokerSyncLatest = "Authorize: databroker.DataBrokerService/SyncLatest"
idpServerGetUserinfo = "IDP: Server: GET /oidc/userinfo"
idpServerPostToken = "IDP: Server: POST /oidc/token"
controlPlaneEnvoyAccessLogs = "Control Plane: envoy.service.accesslog.v3.AccessLogService/StreamAccessLogs"
controlPlaneEnvoyDiscovery = "Control Plane: envoy.service.discovery.v3.AggregatedDiscoveryService/DeltaAggregatedResources"
controlPlaneExport = "Control Plane: opentelemetry.proto.collector.trace.v1.TraceService/Export"
)
results.MatchTraces(t,
@ -96,11 +98,13 @@ func TestOTLPTracing(t *testing.T) {
Exact: true,
CheckDetachedSpans: true,
},
Match{Name: testEnvironmentLocalTest, TraceCount: 1, Services: []string{"Authorize", "Test Environment", "Control Plane", "Data Broker"}},
Match{Name: testEnvironmentLocalTest, TraceCount: 1, Services: []string{"Test Environment", "Control Plane", "Data Broker"}},
Match{Name: testEnvironmentAuthenticate, TraceCount: 1, Services: allServices},
Match{Name: authenticateOAuth2Client, TraceCount: Greater(0)},
Match{Name: idpServerGetUserinfo, TraceCount: EqualToMatch(authenticateOAuth2Client)},
Match{Name: idpServerPostToken, TraceCount: EqualToMatch(authenticateOAuth2Client)},
Match{Name: authorizeDatabrokerSync, TraceCount: Greater(0)},
Match{Name: authorizeDatabrokerSyncLatest, TraceCount: Greater(0)},
Match{Name: controlPlaneEnvoyDiscovery, TraceCount: 1},
Match{Name: controlPlaneExport, TraceCount: Greater(0)},
Match{Name: controlPlaneEnvoyAccessLogs, TraceCount: Any{}},
@ -283,6 +287,7 @@ func (s *SamplingTestSuite) TestExternalTraceparentNeverSample() {
"IDP: Server: POST /oidc/token": {},
"IDP: Server: GET /oidc/userinfo": {},
"Authenticate: OAuth2 Client: GET /.well-known/jwks.json": {},
"Authorize: databroker.DataBrokerService/SyncLatest": {},
}
actual := slices.Collect(maps.Keys(traces.ByName))
for _, name := range actual {

11
pkg/storage/postgres/tracing_test.go
View file

@ -58,12 +58,13 @@ func TestQueryTracing(t *testing.T) {
results := tracetest.NewTraceResults(receiver.FlushResourceSpans())
traces, exists := results.GetTraces().ByParticipant["Data Broker"]
require.True(t, exists)
require.Len(t, traces, 1)
var found bool
for _, span := range traces[0].Spans {
if span.Scope.GetName() == "github.com/exaring/otelpgx" {
found = true
break
for _, trace := range traces {
for _, span := range trace.Spans {
if span.Scope.GetName() == "github.com/exaring/otelpgx" {
found = true
break
}
}
}
assert.True(t, found, "no spans with otelpgx scope found")

9
pkg/storage/querier.go
View file

@ -3,6 +3,7 @@ package storage
import (
"context"
"encoding/json"
"errors"
grpc "google.golang.org/grpc"
"google.golang.org/grpc/codes"
@ -14,10 +15,14 @@ import (
"github.com/pomerium/pomerium/pkg/grpcutil"
)
// ErrUnavailable indicates that a querier is not available.
var ErrUnavailable = errors.New("unavailable")
// A Querier is a read-only subset of the client methods
type Querier interface {
InvalidateCache(ctx context.Context, in *databroker.QueryRequest)
Query(ctx context.Context, in *databroker.QueryRequest, opts ...grpc.CallOption) (*databroker.QueryResponse, error)
Stop()
}
// nilQuerier always returns NotFound.
@ -26,9 +31,11 @@ type nilQuerier struct{}
func (nilQuerier) InvalidateCache(_ context.Context, _ *databroker.QueryRequest) {}
func (nilQuerier) Query(_ context.Context, _ *databroker.QueryRequest, _ ...grpc.CallOption) (*databroker.QueryResponse, error) {
return nil, status.Error(codes.NotFound, "not found")
return nil, errors.Join(ErrUnavailable, status.Error(codes.NotFound, "not found"))
}
func (nilQuerier) Stop() {}
type querierKey struct{}
// GetQuerier gets the databroker Querier from the context.

2
pkg/storage/querier_caching.go
View file

@ -50,6 +50,8 @@ func (q *cachingQuerier) Query(ctx context.Context, in *databroker.QueryRequest,
return res, nil
}
func (*cachingQuerier) Stop() {}
func (q *cachingQuerier) getCacheKey(in *databroker.QueryRequest) ([]byte, error) {
in = proto.Clone(in).(*databroker.QueryRequest)
in.MinimumRecordVersionHint = nil

2
pkg/storage/querier_client.go
View file

@ -23,3 +23,5 @@ func (q *clientQuerier) InvalidateCache(_ context.Context, _ *databroker.QueryRe
func (q *clientQuerier) Query(ctx context.Context, in *databroker.QueryRequest, opts ...grpc.CallOption) (*databroker.QueryResponse, error) {
return q.client.Query(ctx, in, opts...)
}
func (*clientQuerier) Stop() {}

49
pkg/storage/querier_fallback.go Normal file
View file

@ -0,0 +1,49 @@
package storage
import (
"context"
"errors"
grpc "google.golang.org/grpc"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
)
type fallbackQuerier []Querier
// NewFallbackQuerier creates a new fallback-querier. The first call to Query that
// does not return an error will be used.
func NewFallbackQuerier(queriers ...Querier) Querier {
return fallbackQuerier(queriers)
}
// InvalidateCache invalidates the cache of all the queriers.
func (q fallbackQuerier) InvalidateCache(ctx context.Context, req *databroker.QueryRequest) {
for _, qq := range q {
qq.InvalidateCache(ctx, req)
}
}
// Query returns the first querier's results that doesn't result in an error.
func (q fallbackQuerier) Query(ctx context.Context, req *databroker.QueryRequest, opts ...grpc.CallOption) (*databroker.QueryResponse, error) {
if len(q) == 0 {
return nil, ErrUnavailable
}
var merr error
for _, qq := range q {
res, err := qq.Query(ctx, req, opts...)
if err == nil {
return res, nil
}
merr = errors.Join(merr, err)
}
return nil, merr
}
// Stop stops all the queriers.
func (q fallbackQuerier) Stop() {
for _, qq := range q {
qq.Stop()
}
}

36
pkg/storage/querier_fallback_test.go Normal file
View file

@ -0,0 +1,36 @@
package storage_test
import (
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/stretchr/testify/assert"
"google.golang.org/protobuf/testing/protocmp"
"github.com/pomerium/pomerium/internal/testutil"
databrokerpb "github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/storage"
)
func TestFallbackQuerier(t *testing.T) {
t.Parallel()
ctx := testutil.GetContext(t, time.Minute)
q1 := storage.GetQuerier(ctx) // nil querier
q2 := storage.NewStaticQuerier(&databrokerpb.Record{
Type: "t1",
Id: "r1",
Version: 1,
})
res, err := storage.NewFallbackQuerier(q1, q2).Query(ctx, &databrokerpb.QueryRequest{
Type: "t1",
Limit: 1,
})
assert.NoError(t, err, "should fallback")
assert.Empty(t, cmp.Diff(&databrokerpb.QueryResponse{
Records: []*databrokerpb.Record{{Type: "t1", Id: "r1", Version: 1}},
TotalCount: 1,
RecordVersion: 1,
}, res, protocmp.Transform()))
}

2
pkg/storage/querier_static.go
View file

@ -81,3 +81,5 @@ func (q *staticQuerier) InvalidateCache(_ context.Context, _ *databroker.QueryRe
func (q *staticQuerier) Query(_ context.Context, req *databroker.QueryRequest, _ ...grpc.CallOption) (*databroker.QueryResponse, error) {
return QueryRecordCollections(q.records, req)
}
func (*staticQuerier) Stop() {}

184
pkg/storage/querier_sync.go Normal file
View file

@ -0,0 +1,184 @@
package storage
import (
"context"
"errors"
"fmt"
"io"
"sync"
"time"
"github.com/cenkalti/backoff/v4"
grpc "google.golang.org/grpc"
"google.golang.org/grpc/codes"
status "google.golang.org/grpc/status"
"github.com/pomerium/pomerium/internal/log"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
)
type syncQuerier struct {
client databroker.DataBrokerServiceClient
recordType string
cancel context.CancelFunc
serverVersion uint64
latestRecordVersion uint64
mu sync.RWMutex
ready bool
records RecordCollection
}
// NewSyncQuerier creates a new Querier backed by an in-memory record collection
// filled via sync calls to the databroker.
func NewSyncQuerier(
client databroker.DataBrokerServiceClient,
recordType string,
) Querier {
q := &syncQuerier{
client: client,
recordType: recordType,
records: NewRecordCollection(),
}
ctx, cancel := context.WithCancel(context.Background())
q.cancel = cancel
go q.run(ctx)
return q
}
func (q *syncQuerier) InvalidateCache(_ context.Context, _ *databroker.QueryRequest) {
// do nothing
}
func (q *syncQuerier) Query(_ context.Context, req *databroker.QueryRequest, _ ...grpc.CallOption) (*databroker.QueryResponse, error) {
q.mu.RLock()
if !q.canHandleQueryLocked(req) {
q.mu.RUnlock()
return nil, ErrUnavailable
}
defer q.mu.RUnlock()
return QueryRecordCollections(map[string]RecordCollection{
q.recordType: q.records,
}, req)
}
func (q *syncQuerier) Stop() {
q.cancel()
}
func (q *syncQuerier) canHandleQueryLocked(req *databroker.QueryRequest) bool {
if !q.ready {
return false
}
if req.GetType() != q.recordType {
return false
}
if req.MinimumRecordVersionHint != nil && q.latestRecordVersion < *req.MinimumRecordVersionHint {
return false
}
return true
}
func (q *syncQuerier) run(ctx context.Context) {
bo := backoff.WithContext(backoff.NewExponentialBackOff(backoff.WithMaxElapsedTime(0)), ctx)
_ = backoff.RetryNotify(func() error {
if q.serverVersion == 0 {
err := q.syncLatest(ctx)
if err != nil {
return err
}
}
return q.sync(ctx)
}, bo, func(err error, d time.Duration) {
log.Ctx(ctx).Error().
Err(err).
Dur("delay", d).
Msg("storage/sync-querier: error syncing records")
})
}
func (q *syncQuerier) syncLatest(ctx context.Context) error {
stream, err := q.client.SyncLatest(ctx, &databroker.SyncLatestRequest{
Type: q.recordType,
})
if err != nil {
return fmt.Errorf("error starting sync latest stream: %w", err)
}
q.mu.Lock()
q.ready = false
q.records.Clear()
q.mu.Unlock()
for {
res, err := stream.Recv()
if errors.Is(err, io.EOF) {
break
} else if err != nil {
return fmt.Errorf("error receiving sync latest message: %w", err)
}
switch res := res.Response.(type) {
case *databroker.SyncLatestResponse_Record:
q.mu.Lock()
q.records.Put(res.Record)
q.mu.Unlock()
case *databroker.SyncLatestResponse_Versions:
q.mu.Lock()
q.serverVersion = res.Versions.ServerVersion
q.latestRecordVersion = res.Versions.LatestRecordVersion
q.mu.Unlock()
default:
return fmt.Errorf("unknown message type from sync latest: %T", res)
}
}
q.mu.Lock()
log.Ctx(ctx).Info().
Str("record-type", q.recordType).
Int("record-count", q.records.Len()).
Uint64("latest-record-version", q.latestRecordVersion).
Msg("storage/sync-querier: synced latest records")
q.ready = true
q.mu.Unlock()
return nil
}
func (q *syncQuerier) sync(ctx context.Context) error {
q.mu.RLock()
req := &databroker.SyncRequest{
ServerVersion: q.serverVersion,
RecordVersion: q.latestRecordVersion,
Type: q.recordType,
}
q.mu.RUnlock()
stream, err := q.client.Sync(ctx, req)
if err != nil {
return fmt.Errorf("error starting sync stream: %w", err)
}
for {
res, err := stream.Recv()
if status.Code(err) == codes.Aborted {
// this indicates the server version changed, so we need to reset
q.mu.Lock()
q.serverVersion = 0
q.latestRecordVersion = 0
q.mu.Unlock()
return fmt.Errorf("stream was aborted due to mismatched server versions: %w", err)
} else if err != nil {
return fmt.Errorf("error receiving sync message: %w", err)
}
q.mu.Lock()
q.latestRecordVersion = max(q.latestRecordVersion, res.Record.Version)
q.records.Put(res.Record)
q.mu.Unlock()
}
}

89
pkg/storage/querier_sync_test.go Normal file
View file

@ -0,0 +1,89 @@
package storage_test
import (
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.opentelemetry.io/otel/trace/noop"
grpc "google.golang.org/grpc"
"google.golang.org/protobuf/testing/protocmp"
"google.golang.org/protobuf/types/known/structpb"
"github.com/pomerium/pomerium/internal/databroker"
"github.com/pomerium/pomerium/internal/testutil"
databrokerpb "github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/protoutil"
"github.com/pomerium/pomerium/pkg/storage"
)
func TestSyncQuerier(t *testing.T) {
t.Parallel()
ctx := testutil.GetContext(t, 10*time.Minute)
cc := testutil.NewGRPCServer(t, func(srv *grpc.Server) {
databrokerpb.RegisterDataBrokerServiceServer(srv, databroker.New(ctx, noop.NewTracerProvider()))
})
t.Cleanup(func() { cc.Close() })
client := databrokerpb.NewDataBrokerServiceClient(cc)
r1 := &databrokerpb.Record{
Type: "t1",
Id: "r1",
Data: protoutil.ToAny("q2"),
}
_, err := client.Put(ctx, &databrokerpb.PutRequest{
Records: []*databrokerpb.Record{r1},
})
require.NoError(t, err)
r2 := &databrokerpb.Record{
Type: "t1",
Id: "r2",
Data: protoutil.ToAny("q2"),
}
q := storage.NewSyncQuerier(client, "t1")
t.Cleanup(q.Stop)
assert.EventuallyWithT(t, func(c *assert.CollectT) {
res, err := q.Query(ctx, &databrokerpb.QueryRequest{
Type: "t1",
Filter: newStruct(t, map[string]any{
"id": "r1",
}),
Limit: 1,
})
if assert.NoError(c, err) && assert.Len(c, res.Records, 1) {
assert.Empty(c, cmp.Diff(r1.Data, res.Records[0].Data, protocmp.Transform()))
}
}, time.Second*10, time.Millisecond*50, "should sync records")
_, err = client.Put(ctx, &databrokerpb.PutRequest{
Records: []*databrokerpb.Record{r2},
})
require.NoError(t, err)
assert.EventuallyWithT(t, func(c *assert.CollectT) {
res, err := q.Query(ctx, &databrokerpb.QueryRequest{
Type: "t1",
Filter: newStruct(t, map[string]any{
"id": "r2",
}),
Limit: 1,
})
if assert.NoError(c, err) && assert.Len(c, res.Records, 1) {
assert.Empty(c, cmp.Diff(r2.Data, res.Records[0].Data, protocmp.Transform()))
}
}, time.Second*10, time.Millisecond*50, "should pick up changes")
}
func newStruct(t *testing.T, m map[string]any) *structpb.Struct {
t.Helper()
s, err := structpb.NewStruct(m)
require.NoError(t, err)
return s
}

45
pkg/storage/querier_typed.go Normal file
View file

@ -0,0 +1,45 @@
package storage
import (
"context"
grpc "google.golang.org/grpc"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
)
type typedQuerier struct {
defaultQuerier Querier
queriersByType map[string]Querier
}
// NewTypedQuerier creates a new Querier that dispatches to other queries based on the type.
func NewTypedQuerier(defaultQuerier Querier, queriersByType map[string]Querier) Querier {
return &typedQuerier{
defaultQuerier: defaultQuerier,
queriersByType: queriersByType,
}
}
func (q *typedQuerier) InvalidateCache(ctx context.Context, req *databroker.QueryRequest) {
qq, ok := q.queriersByType[req.Type]
if !ok {
qq = q.defaultQuerier
}
qq.InvalidateCache(ctx, req)
}
func (q *typedQuerier) Query(ctx context.Context, req *databroker.QueryRequest, opts ...grpc.CallOption) (*databroker.QueryResponse, error) {
qq, ok := q.queriersByType[req.Type]
if !ok {
qq = q.defaultQuerier
}
return qq.Query(ctx, req, opts...)
}
func (q *typedQuerier) Stop() {
q.defaultQuerier.Stop()
for _, qq := range q.queriersByType {
qq.Stop()
}
}

68
pkg/storage/querier_typed_test.go Normal file
View file

@ -0,0 +1,68 @@
package storage_test
import (
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/stretchr/testify/assert"
"google.golang.org/protobuf/testing/protocmp"
"github.com/pomerium/pomerium/internal/testutil"
databrokerpb "github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/storage"
)
func TestTypedQuerier(t *testing.T) {
t.Parallel()
ctx := testutil.GetContext(t, time.Minute)
q1 := storage.NewStaticQuerier(&databrokerpb.Record{
Type: "t1",
Id: "r1",
})
q2 := storage.NewStaticQuerier(&databrokerpb.Record{
Type: "t2",
Id: "r2",
})
q3 := storage.NewStaticQuerier(&databrokerpb.Record{
Type: "t3",
Id: "r3",
})
q := storage.NewTypedQuerier(q1, map[string]storage.Querier{
"t2": q2,
"t3": q3,
})
res, err := q.Query(ctx, &databrokerpb.QueryRequest{
Type: "t1",
Limit: 1,
})
assert.NoError(t, err)
assert.Empty(t, cmp.Diff(&databrokerpb.QueryResponse{
Records: []*databrokerpb.Record{{Type: "t1", Id: "r1"}},
TotalCount: 1,
}, res, protocmp.Transform()))
res, err = q.Query(ctx, &databrokerpb.QueryRequest{
Type: "t2",
Limit: 1,
})
assert.NoError(t, err)
assert.Empty(t, cmp.Diff(&databrokerpb.QueryResponse{
Records: []*databrokerpb.Record{{Type: "t2", Id: "r2"}},
TotalCount: 1,
}, res, protocmp.Transform()))
res, err = q.Query(ctx, &databrokerpb.QueryRequest{
Type: "t3",
Limit: 1,
})
assert.NoError(t, err)
assert.Empty(t, cmp.Diff(&databrokerpb.QueryResponse{
Records: []*databrokerpb.Record{{Type: "t3", Id: "r3"}},
TotalCount: 1,
}, res, protocmp.Transform()))
}

2
pkg/telemetry/trace/client_test.go
View file

@ -297,6 +297,8 @@ func (h *errHandler) Handle(err error) {
}
func TestNewTraceClientFromConfig(t *testing.T) {
t.Skip("failing because authorize uses databroker sync now")
env := testenv.New(t, testenv.WithTraceDebugFlags(testenv.StandardTraceDebugFlags))
receiver := scenarios.NewOTLPTraceReceiver()