databroker: rename cache service (#1790)

* rename cache folder * rename cache service everywhere * skip yaml in examples * Update docs/docs/topics/data-storage.md Co-authored-by: Travis Groth <travisgroth@users.noreply.github.com> Co-authored-by: Travis Groth <travisgroth@users.noreply.github.com>
2025-08-04 01:09:36 +02:00 · 2021-01-21 08:41:22 -07:00 · 2021-01-21 08:41:22 -07:00 · 70b4497595
commit 70b4497595
parent 0adb9e5dde
27 changed files with 115 additions and 108 deletions
--- a/databroker/cache.go
+++ b/databroker/cache.go
@ -0,0 +1,187 @@
+// Package databroker is a pomerium service that handles the storage of user
+// session state. It communicates over RPC with other pomerium services,
+// and can be configured to use a number of different backend databroker stores.
+package databroker
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+	"net"
+	"sync"
+
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+
+	"github.com/pomerium/pomerium/config"
+	"github.com/pomerium/pomerium/internal/directory"
+	"github.com/pomerium/pomerium/internal/identity"
+	"github.com/pomerium/pomerium/internal/identity/manager"
+	"github.com/pomerium/pomerium/internal/log"
+	"github.com/pomerium/pomerium/internal/telemetry"
+	"github.com/pomerium/pomerium/internal/urlutil"
+	"github.com/pomerium/pomerium/internal/version"
+	"github.com/pomerium/pomerium/pkg/cryptutil"
+	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/grpcutil"
+)
+
+// DataBroker represents the databroker service. The databroker service is a simple interface
+// for storing keyed blobs (bytes) of unstructured data.
+type DataBroker struct {
+	dataBrokerServer *dataBrokerServer
+	manager          *manager.Manager
+
+	localListener                net.Listener
+	localGRPCServer              *grpc.Server
+	localGRPCConnection          *grpc.ClientConn
+	dataBrokerStorageType        string // TODO remove in v0.11
+	deprecatedCacheClusterDomain string // TODO: remove in v0.11
+
+	mu                sync.Mutex
+	directoryProvider directory.Provider
+}
+
+// New creates a new databroker service.
+func New(cfg *config.Config) (*DataBroker, error) {
+	localListener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		return nil, err
+	}
+
+	sharedKey, _ := base64.StdEncoding.DecodeString(cfg.Options.SharedKey)
+
+	ui, si := grpcutil.AttachMetadataInterceptors(
+		metadata.Pairs(grpcutil.MetadataKeyPomeriumVersion, version.FullVersion()),
+	)
+
+	// No metrics handler because we have one in the control plane.  Add one
+	// if we no longer register with that grpc Server
+	localGRPCServer := grpc.NewServer(
+		grpc.StreamInterceptor(si),
+		grpc.UnaryInterceptor(ui),
+	)
+
+	clientStatsHandler := telemetry.NewGRPCClientStatsHandler(cfg.Options.Services)
+	clientDialOptions := []grpc.DialOption{
+		grpc.WithInsecure(),
+		grpc.WithChainUnaryInterceptor(clientStatsHandler.UnaryInterceptor, grpcutil.WithUnarySignedJWT(sharedKey)),
+		grpc.WithChainStreamInterceptor(grpcutil.WithStreamSignedJWT(sharedKey)),
+		grpc.WithStatsHandler(clientStatsHandler.Handler),
+	}
+
+	localGRPCConnection, err := grpc.DialContext(
+		context.Background(),
+		localListener.Addr().String(),
+		clientDialOptions...,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	dataBrokerServer := newDataBrokerServer(cfg)
+
+	c := &DataBroker{
+		dataBrokerServer:             dataBrokerServer,
+		localListener:                localListener,
+		localGRPCServer:              localGRPCServer,
+		localGRPCConnection:          localGRPCConnection,
+		deprecatedCacheClusterDomain: cfg.Options.GetDataBrokerURL().Hostname(),
+		dataBrokerStorageType:        cfg.Options.DataBrokerStorageType,
+	}
+	c.Register(c.localGRPCServer)
+
+	err = c.update(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return c, nil
+}
+
+// OnConfigChange is called whenever configuration is changed.
+func (c *DataBroker) OnConfigChange(cfg *config.Config) {
+	err := c.update(cfg)
+	if err != nil {
+		log.Error().Err(err).Msg("databroker: error updating configuration")
+	}
+
+	c.dataBrokerServer.OnConfigChange(cfg)
+}
+
+// Register registers all the gRPC services with the given server.
+func (c *DataBroker) Register(grpcServer *grpc.Server) {
+	databroker.RegisterDataBrokerServiceServer(grpcServer, c.dataBrokerServer)
+	directory.RegisterDirectoryServiceServer(grpcServer, c)
+}
+
+// Run runs the databroker components.
+func (c *DataBroker) Run(ctx context.Context) error {
+	eg, ctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		return c.localGRPCServer.Serve(c.localListener)
+	})
+	eg.Go(func() error {
+		<-ctx.Done()
+		c.localGRPCServer.Stop()
+		return nil
+	})
+	eg.Go(func() error {
+		return c.manager.Run(ctx)
+	})
+	return eg.Wait()
+}
+
+func (c *DataBroker) update(cfg *config.Config) error {
+	if err := validate(cfg.Options); err != nil {
+		return fmt.Errorf("databroker: bad option: %w", err)
+	}
+
+	authenticator, err := identity.NewAuthenticator(cfg.Options.GetOauthOptions())
+	if err != nil {
+		return fmt.Errorf("databroker: failed to create authenticator: %w", err)
+	}
+
+	directoryProvider := directory.GetProvider(directory.Options{
+		ServiceAccount: cfg.Options.ServiceAccount,
+		Provider:       cfg.Options.Provider,
+		ProviderURL:    cfg.Options.ProviderURL,
+		QPS:            cfg.Options.QPS,
+		ClientID:       cfg.Options.ClientID,
+		ClientSecret:   cfg.Options.ClientSecret,
+	})
+	c.mu.Lock()
+	c.directoryProvider = directoryProvider
+	c.mu.Unlock()
+
+	dataBrokerClient := databroker.NewDataBrokerServiceClient(c.localGRPCConnection)
+
+	options := []manager.Option{
+		manager.WithAuthenticator(authenticator),
+		manager.WithDirectoryProvider(directoryProvider),
+		manager.WithDataBrokerClient(dataBrokerClient),
+		manager.WithGroupRefreshInterval(cfg.Options.RefreshDirectoryInterval),
+		manager.WithGroupRefreshTimeout(cfg.Options.RefreshDirectoryTimeout),
+	}
+
+	if c.manager == nil {
+		c.manager = manager.New(options...)
+	} else {
+		c.manager.UpdateConfig(options...)
+	}
+
+	return nil
+}
+
+// validate checks that proper configuration settings are set to create
+// a databroker instance
+func validate(o *config.Options) error {
+	if _, err := cryptutil.NewAEADCipherFromBase64(o.SharedKey); err != nil {
+		return fmt.Errorf("invalid 'SHARED_SECRET': %w", err)
+	}
+	if err := urlutil.ValidateURL(o.DataBrokerURL); err != nil {
+		return fmt.Errorf("invalid 'DATA_BROKER_SERVICE_URL': %w", err)
+	}
+	return nil
+}
--- a/databroker/cache_test.go
+++ b/databroker/cache_test.go
@ -0,0 +1,40 @@
+package databroker
+
+import (
+	"io/ioutil"
+	"log"
+	"net/url"
+	"os"
+	"testing"
+
+	"github.com/pomerium/pomerium/config"
+	"github.com/pomerium/pomerium/pkg/cryptutil"
+)
+
+func TestNew(t *testing.T) {
+	dir, err := ioutil.TempDir("", "example")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer os.RemoveAll(dir)
+
+	tests := []struct {
+		name    string
+		opts    config.Options
+		wantErr bool
+	}{
+		{"good", config.Options{SharedKey: cryptutil.NewBase64Key(), DataBrokerURL: &url.URL{Scheme: "http", Host: "example"}}, false},
+		{"bad shared secret", config.Options{SharedKey: string([]byte(cryptutil.NewBase64Key())[:31]), DataBrokerURL: &url.URL{Scheme: "http", Host: "example"}}, true},
+		{"bad databroker url", config.Options{SharedKey: cryptutil.NewBase64Key(), DataBrokerURL: &url.URL{}}, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.opts.Provider = "google"
+			_, err := New(&config.Config{Options: &tt.opts})
+			if (err != nil) != tt.wantErr {
+				t.Errorf("New() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+		})
+	}
+}
--- a/databroker/databroker.go
+++ b/databroker/databroker.go
@ -0,0 +1,109 @@
+package databroker
+
+import (
+	"context"
+	"encoding/base64"
+	"sync/atomic"
+
+	"github.com/golang/protobuf/ptypes/empty"
+
+	"github.com/pomerium/pomerium/config"
+	"github.com/pomerium/pomerium/internal/databroker"
+	databrokerpb "github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/grpcutil"
+)
+
+// A dataBrokerServer implements the data broker service interface.
+type dataBrokerServer struct {
+	server    *databroker.Server
+	sharedKey atomic.Value
+}
+
+// newDataBrokerServer creates a new databroker service server.
+func newDataBrokerServer(cfg *config.Config) *dataBrokerServer {
+	srv := &dataBrokerServer{}
+	srv.server = databroker.New(srv.getOptions(cfg)...)
+	srv.setKey(cfg)
+	return srv
+}
+
+// OnConfigChange updates the underlying databroker server whenever configuration is changed.
+func (srv *dataBrokerServer) OnConfigChange(cfg *config.Config) {
+	srv.server.UpdateConfig(srv.getOptions(cfg)...)
+	srv.setKey(cfg)
+}
+
+func (srv *dataBrokerServer) getOptions(cfg *config.Config) []databroker.ServerOption {
+	return []databroker.ServerOption{
+		databroker.WithSharedKey(cfg.Options.SharedKey),
+		databroker.WithStorageType(cfg.Options.DataBrokerStorageType),
+		databroker.WithStorageConnectionString(cfg.Options.DataBrokerStorageConnectionString),
+		databroker.WithStorageCAFile(cfg.Options.DataBrokerStorageCAFile),
+		databroker.WithStorageCertificate(cfg.Options.DataBrokerCertificate),
+		databroker.WithStorageCertSkipVerify(cfg.Options.DataBrokerStorageCertSkipVerify),
+	}
+}
+
+func (srv *dataBrokerServer) setKey(cfg *config.Config) {
+	bs, _ := base64.StdEncoding.DecodeString(cfg.Options.SharedKey)
+	if bs == nil {
+		bs = make([]byte, 0)
+	}
+	srv.sharedKey.Store(bs)
+}
+
+func (srv *dataBrokerServer) Delete(ctx context.Context, req *databrokerpb.DeleteRequest) (*empty.Empty, error) {
+	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
+		return nil, err
+	}
+	return srv.server.Delete(ctx, req)
+}
+
+func (srv *dataBrokerServer) Get(ctx context.Context, req *databrokerpb.GetRequest) (*databrokerpb.GetResponse, error) {
+	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
+		return nil, err
+	}
+	return srv.server.Get(ctx, req)
+}
+
+func (srv *dataBrokerServer) GetAll(ctx context.Context, req *databrokerpb.GetAllRequest) (*databrokerpb.GetAllResponse, error) {
+	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
+		return nil, err
+	}
+	return srv.server.GetAll(ctx, req)
+}
+
+func (srv *dataBrokerServer) Query(ctx context.Context, req *databrokerpb.QueryRequest) (*databrokerpb.QueryResponse, error) {
+	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
+		return nil, err
+	}
+	return srv.server.Query(ctx, req)
+}
+
+func (srv *dataBrokerServer) Set(ctx context.Context, req *databrokerpb.SetRequest) (*databrokerpb.SetResponse, error) {
+	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
+		return nil, err
+	}
+	return srv.server.Set(ctx, req)
+}
+
+func (srv *dataBrokerServer) Sync(req *databrokerpb.SyncRequest, stream databrokerpb.DataBrokerService_SyncServer) error {
+	if err := grpcutil.RequireSignedJWT(stream.Context(), srv.sharedKey.Load().([]byte)); err != nil {
+		return err
+	}
+	return srv.server.Sync(req, stream)
+}
+
+func (srv *dataBrokerServer) GetTypes(ctx context.Context, req *empty.Empty) (*databrokerpb.GetTypesResponse, error) {
+	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
+		return nil, err
+	}
+	return srv.server.GetTypes(ctx, req)
+}
+
+func (srv *dataBrokerServer) SyncTypes(req *empty.Empty, stream databrokerpb.DataBrokerService_SyncTypesServer) error {
+	if err := grpcutil.RequireSignedJWT(stream.Context(), srv.sharedKey.Load().([]byte)); err != nil {
+		return err
+	}
+	return srv.server.SyncTypes(req, stream)
+}
--- a/databroker/databroker_test.go
+++ b/databroker/databroker_test.go
@ -0,0 +1,125 @@
+package databroker
+
+import (
+	"context"
+	"net"
+	"strconv"
+	"testing"
+
+	"github.com/golang/protobuf/ptypes"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/test/bufconn"
+
+	internal_databroker "github.com/pomerium/pomerium/internal/databroker"
+	"github.com/pomerium/pomerium/internal/log"
+	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/grpc/session"
+	"github.com/pomerium/pomerium/pkg/grpc/user"
+)
+
+const bufSize = 1024 * 1024
+
+var lis *bufconn.Listener
+
+func init() {
+	lis = bufconn.Listen(bufSize)
+	s := grpc.NewServer()
+	internalSrv := internal_databroker.New()
+	srv := &dataBrokerServer{server: internalSrv}
+	srv.sharedKey.Store([]byte{})
+	databroker.RegisterDataBrokerServiceServer(s, srv)
+
+	go func() {
+		if err := s.Serve(lis); err != nil {
+			log.Fatal().Err(err).Msg("Server exited with error")
+		}
+	}()
+}
+
+func bufDialer(context.Context, string) (net.Conn, error) {
+	return lis.Dial()
+}
+
+func TestServerSync(t *testing.T) {
+	ctx := context.Background()
+	conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
+	require.NoError(t, err)
+	defer conn.Close()
+	c := databroker.NewDataBrokerServiceClient(conn)
+	any, _ := ptypes.MarshalAny(new(user.User))
+	numRecords := 200
+
+	for i := 0; i < numRecords; i++ {
+		c.Set(ctx, &databroker.SetRequest{Type: any.TypeUrl, Id: strconv.Itoa(i), Data: any})
+	}
+
+	t.Run("Sync ok", func(t *testing.T) {
+		client, _ := c.Sync(ctx, &databroker.SyncRequest{Type: any.GetTypeUrl()})
+		count := 0
+		for {
+			res, err := client.Recv()
+			if err != nil {
+				break
+			}
+			count += len(res.Records)
+			if count == numRecords {
+				break
+			}
+		}
+	})
+	t.Run("Error occurred while syncing", func(t *testing.T) {
+		ctx, cancelFunc := context.WithCancel(ctx)
+		defer cancelFunc()
+
+		client, _ := c.Sync(ctx, &databroker.SyncRequest{Type: any.GetTypeUrl()})
+		count := 0
+		numRecordsWanted := 100
+		cancelFuncCalled := false
+		for {
+			res, err := client.Recv()
+			if err != nil {
+				assert.True(t, cancelFuncCalled)
+				break
+			}
+			count += len(res.Records)
+			if count == numRecordsWanted {
+				cancelFunc()
+				cancelFuncCalled = true
+			}
+		}
+	})
+}
+
+func BenchmarkSync(b *testing.B) {
+	ctx := context.Background()
+	conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
+	if err != nil {
+		b.Fatalf("Failed to dial bufnet: %v", err)
+	}
+	defer conn.Close()
+	c := databroker.NewDataBrokerServiceClient(conn)
+	any, _ := ptypes.MarshalAny(new(session.Session))
+	numRecords := 10000
+
+	for i := 0; i < numRecords; i++ {
+		c.Set(ctx, &databroker.SetRequest{Type: any.TypeUrl, Id: strconv.Itoa(i), Data: any})
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		client, _ := c.Sync(ctx, &databroker.SyncRequest{Type: any.GetTypeUrl()})
+		count := 0
+		for {
+			res, err := client.Recv()
+			if err != nil {
+				break
+			}
+			count += len(res.Records)
+			if count == numRecords {
+				break
+			}
+		}
+	}
+}
--- a/databroker/directory.go
+++ b/databroker/directory.go
@ -0,0 +1,44 @@
+package databroker
+
+import (
+	"context"
+	"errors"
+
+	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/emptypb"
+
+	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/grpc/directory"
+)
+
+// RefreshUser refreshes a user's directory information.
+func (c *DataBroker) RefreshUser(ctx context.Context, req *directory.RefreshUserRequest) (*emptypb.Empty, error) {
+	c.mu.Lock()
+	dp := c.directoryProvider
+	c.mu.Unlock()
+
+	if dp == nil {
+		return nil, errors.New("no directory provider is available for refresh")
+	}
+
+	u, err := dp.User(ctx, req.GetUserId(), req.GetAccessToken())
+	if err != nil {
+		return nil, err
+	}
+
+	any, err := anypb.New(u)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = c.dataBrokerServer.Set(ctx, &databroker.SetRequest{
+		Type: any.GetTypeUrl(),
+		Id:   u.GetId(),
+		Data: any,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return new(emptypb.Empty), nil
+}