atomicutil: use atomicutil.Value wherever possible (#3517)

* atomicutil: use atomicutil.Value wherever possible

* fix test

* fix mux router
This commit is contained in:
Caleb Doxsey 2022-07-28 15:38:38 -06:00 committed by GitHub
parent 5c14d2c994
commit 0ac7e45a21
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
23 changed files with 121 additions and 215 deletions

View file

@ -3,11 +3,11 @@ package databroker
import (
"context"
"sync/atomic"
"google.golang.org/protobuf/types/known/emptypb"
"github.com/pomerium/pomerium/config"
"github.com/pomerium/pomerium/internal/atomicutil"
"github.com/pomerium/pomerium/internal/databroker"
databrokerpb "github.com/pomerium/pomerium/pkg/grpc/databroker"
registrypb "github.com/pomerium/pomerium/pkg/grpc/registry"
@ -17,12 +17,14 @@ import (
// A dataBrokerServer implements the data broker service interface.
type dataBrokerServer struct {
server *databroker.Server
sharedKey atomic.Value
sharedKey *atomicutil.Value[[]byte]
}
// newDataBrokerServer creates a new databroker service server.
func newDataBrokerServer(cfg *config.Config) *dataBrokerServer {
srv := &dataBrokerServer{}
srv := &dataBrokerServer{
sharedKey: atomicutil.NewValue([]byte{}),
}
srv.server = databroker.New(srv.getOptions(cfg)...)
srv.setKey(cfg)
return srv
@ -57,63 +59,63 @@ func (srv *dataBrokerServer) setKey(cfg *config.Config) {
// Databroker functions
func (srv *dataBrokerServer) AcquireLease(ctx context.Context, req *databrokerpb.AcquireLeaseRequest) (*databrokerpb.AcquireLeaseResponse, error) {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.AcquireLease(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 {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.Get(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 {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.Query(ctx, req)
}
func (srv *dataBrokerServer) Put(ctx context.Context, req *databrokerpb.PutRequest) (*databrokerpb.PutResponse, error) {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.Put(ctx, req)
}
func (srv *dataBrokerServer) ReleaseLease(ctx context.Context, req *databrokerpb.ReleaseLeaseRequest) (*emptypb.Empty, error) {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.ReleaseLease(ctx, req)
}
func (srv *dataBrokerServer) RenewLease(ctx context.Context, req *databrokerpb.RenewLeaseRequest) (*emptypb.Empty, error) {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.RenewLease(ctx, req)
}
func (srv *dataBrokerServer) SetOptions(ctx context.Context, req *databrokerpb.SetOptionsRequest) (*databrokerpb.SetOptionsResponse, error) {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.SetOptions(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 {
if err := grpcutil.RequireSignedJWT(stream.Context(), srv.sharedKey.Load()); err != nil {
return err
}
return srv.server.Sync(req, stream)
}
func (srv *dataBrokerServer) SyncLatest(req *databrokerpb.SyncLatestRequest, stream databrokerpb.DataBrokerService_SyncLatestServer) error {
if err := grpcutil.RequireSignedJWT(stream.Context(), srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(stream.Context(), srv.sharedKey.Load()); err != nil {
return err
}
return srv.server.SyncLatest(req, stream)
@ -122,21 +124,21 @@ func (srv *dataBrokerServer) SyncLatest(req *databrokerpb.SyncLatestRequest, str
// Registry functions
func (srv *dataBrokerServer) Report(ctx context.Context, req *registrypb.RegisterRequest) (*registrypb.RegisterResponse, error) {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.Report(ctx, req)
}
func (srv *dataBrokerServer) List(ctx context.Context, req *registrypb.ListRequest) (*registrypb.ServiceList, error) {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
return nil, err
}
return srv.server.List(ctx, req)
}
func (srv *dataBrokerServer) Watch(req *registrypb.ListRequest, stream registrypb.Registry_WatchServer) error {
if err := grpcutil.RequireSignedJWT(stream.Context(), srv.sharedKey.Load().([]byte)); err != nil {
if err := grpcutil.RequireSignedJWT(stream.Context(), srv.sharedKey.Load()); err != nil {
return err
}
return srv.server.Watch(req, stream)