databroker: add patch method (#4704)

Add a Patch() method to the databroker gRPC service.

Update the storage.Backend interface to include the Patch() method now
that all the storage.Backend implementations include it.

Add a test to exercise the patch method under concurrent usage.

databroker/databroker.go

@@ -93,6 +93,13 @@ func (srv *dataBrokerServer) Put(ctx context.Context, req *databrokerpb.PutReque
 	return srv.server.Put(ctx, req)
 }
 
+func (srv *dataBrokerServer) Patch(ctx context.Context, req *databrokerpb.PatchRequest) (*databrokerpb.PatchResponse, error) {
+	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
+		return nil, err
+	}
+	return srv.server.Patch(ctx, req)
+}
+
 func (srv *dataBrokerServer) ReleaseLease(ctx context.Context, req *databrokerpb.ReleaseLeaseRequest) (*emptypb.Empty, error) {
 	if err := grpcutil.RequireSignedJWT(ctx, srv.sharedKey.Load()); err != nil {
 		return nil, err

internal/databroker/server.go

@@ -234,6 +234,45 @@ func (srv *Server) Put(ctx context.Context, req *databroker.PutRequest) (*databr
 	return res, nil
 }
 
+// Patch updates specific fields of an existing record.
+func (srv *Server) Patch(ctx context.Context, req *databroker.PatchRequest) (*databroker.PatchResponse, error) {
+	ctx, span := trace.StartSpan(ctx, "databroker.grpc.Patch")
+	defer span.End()
+
+	records := req.GetRecords()
+	if len(records) == 1 {
+		log.Info(ctx).
+			Str("record-type", records[0].GetType()).
+			Str("record-id", records[0].GetId()).
+			Msg("patch")
+	} else {
+		var recordType string
+		for _, record := range records {
+			recordType = record.GetType()
+		}
+		log.Info(ctx).
+			Int("record-count", len(records)).
+			Str("record-type", recordType).
+			Msg("patch")
+	}
+
+	db, err := srv.getBackend()
+	if err != nil {
+		return nil, err
+	}
+
+	serverVersion, patchedRecords, err := db.Patch(ctx, records, req.GetFieldMask())
+	if err != nil {
+		return nil, err
+	}
+	res := &databroker.PatchResponse{
+		ServerVersion: serverVersion,
+		Records:       patchedRecords,
+	}
+
+	return res, nil
+}
+
 // ReleaseLease releases a lease.
 func (srv *Server) ReleaseLease(ctx context.Context, req *databroker.ReleaseLeaseRequest) (*emptypb.Empty, error) {
 	ctx, span := trace.StartSpan(ctx, "databroker.grpc.ReleaseLease")

internal/databroker/server_test.go

@@ -18,6 +18,7 @@ import (
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/durationpb"
+	"google.golang.org/protobuf/types/known/fieldmaskpb"
 	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
@@ -84,6 +85,58 @@ func TestServer_Get(t *testing.T) {
 	})
 }
 
+func TestServer_Patch(t *testing.T) {
+	cfg := newServerConfig()
+	srv := newServer(cfg)
+
+	s := &session.Session{
+		Id:         "1",
+		OauthToken: &session.OAuthToken{AccessToken: "access-token"},
+	}
+	data := protoutil.NewAny(s)
+	_, err := srv.Put(context.Background(), &databroker.PutRequest{
+		Records: []*databroker.Record{{
+			Type: data.TypeUrl,
+			Id:   s.Id,
+			Data: data,
+		}},
+	})
+	require.NoError(t, err)
+
+	fm, err := fieldmaskpb.New(s, "accessed_at")
+	require.NoError(t, err)
+
+	now := timestamppb.Now()
+	s.AccessedAt = now
+	s.OauthToken.AccessToken = "access-token-field-ignored"
+	data = protoutil.NewAny(s)
+	patchResponse, err := srv.Patch(context.Background(), &databroker.PatchRequest{
+		Records: []*databroker.Record{{
+			Type: data.TypeUrl,
+			Id:   s.Id,
+			Data: data,
+		}},
+		FieldMask: fm,
+	})
+	require.NoError(t, err)
+	testutil.AssertProtoEqual(t, protoutil.NewAny(&session.Session{
+		Id:         "1",
+		AccessedAt: now,
+		OauthToken: &session.OAuthToken{AccessToken: "access-token"},
+	}), patchResponse.GetRecord().GetData())
+
+	getResponse, err := srv.Get(context.Background(), &databroker.GetRequest{
+		Type: data.TypeUrl,
+		Id:   s.Id,
+	})
+	require.NoError(t, err)
+	testutil.AssertProtoEqual(t, protoutil.NewAny(&session.Session{
+		Id:         "1",
+		AccessedAt: now,
+		OauthToken: &session.OAuthToken{AccessToken: "access-token"},
+	}), getResponse.GetRecord().GetData())
+}
+
 func TestServer_Options(t *testing.T) {
 	cfg := newServerConfig()
 	srv := newServer(cfg)

pkg/grpc/databroker/databroker.go

@@ -150,6 +150,15 @@ func (x *PutResponse) GetRecord() *Record {
 	return records[0]
 }
 
+// GetRecord gets the first record, or nil if there are none.
+func (x *PatchResponse) GetRecord() *Record {
+	records := x.GetRecords()
+	if len(records) == 0 {
+		return nil
+	}
+	return records[0]
+}
+
 // SetFilterByID sets the filter to an id.
 func (x *QueryRequest) SetFilterByID(id string) {
 	x.Filter = &structpb.Struct{Fields: map[string]*structpb.Value{

pkg/grpc/databroker/databroker.proto

@@ -6,6 +6,7 @@ option go_package = "github.com/pomerium/pomerium/pkg/grpc/databroker";
 import "google/protobuf/any.proto";
 import "google/protobuf/duration.proto";
 import "google/protobuf/empty.proto";
+import "google/protobuf/field_mask.proto";
 import "google/protobuf/struct.proto";
 import "google/protobuf/timestamp.proto";
 
@@ -60,6 +61,15 @@ message PutResponse {
   repeated Record records = 2;
 }
 
+message PatchRequest {
+  repeated Record records = 1;
+  google.protobuf.FieldMask field_mask = 2;
+}
+message PatchResponse {
+  uint64 server_version = 1;
+  repeated Record records = 2;
+}
+
 message SetOptionsRequest {
   string type = 1;
   Options options = 2;
@@ -114,6 +124,8 @@ service DataBrokerService {
   rpc ListTypes(google.protobuf.Empty) returns (ListTypesResponse);
   // Put saves a record.
   rpc Put(PutRequest) returns (PutResponse);
+  // Patch updates specific fields of an existing record.
+  rpc Patch(PatchRequest) returns (PatchResponse);
   // Query queries for records.
   rpc Query(QueryRequest) returns (QueryResponse);
   // ReleaseLease releases a distributed mutex lease.

pkg/grpc/databroker/mock_databroker/databroker.pb.go

@@ -133,6 +133,26 @@ func (mr *MockDataBrokerServiceClientMockRecorder) ListTypes(ctx, in interface{}
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListTypes", reflect.TypeOf((*MockDataBrokerServiceClient)(nil).ListTypes), varargs...)
 }
 
+// Patch mocks base method.
+func (m *MockDataBrokerServiceClient) Patch(ctx context.Context, in *databroker.PatchRequest, opts ...grpc.CallOption) (*databroker.PatchResponse, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{ctx, in}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Patch", varargs...)
+	ret0, _ := ret[0].(*databroker.PatchResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Patch indicates an expected call of Patch.
+func (mr *MockDataBrokerServiceClientMockRecorder) Patch(ctx, in interface{}, opts ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{ctx, in}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Patch", reflect.TypeOf((*MockDataBrokerServiceClient)(nil).Patch), varargs...)
+}
+
 // Put mocks base method.
 func (m *MockDataBrokerServiceClient) Put(ctx context.Context, in *databroker.PutRequest, opts ...grpc.CallOption) (*databroker.PutResponse, error) {
 	m.ctrl.T.Helper()
@@ -587,6 +607,21 @@ func (mr *MockDataBrokerServiceServerMockRecorder) ListTypes(arg0, arg1 interfac
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListTypes", reflect.TypeOf((*MockDataBrokerServiceServer)(nil).ListTypes), arg0, arg1)
 }
 
+// Patch mocks base method.
+func (m *MockDataBrokerServiceServer) Patch(arg0 context.Context, arg1 *databroker.PatchRequest) (*databroker.PatchResponse, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Patch", arg0, arg1)
+	ret0, _ := ret[0].(*databroker.PatchResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Patch indicates an expected call of Patch.
+func (mr *MockDataBrokerServiceServerMockRecorder) Patch(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Patch", reflect.TypeOf((*MockDataBrokerServiceServer)(nil).Patch), arg0, arg1)
+}
+
 // Put mocks base method.
 func (m *MockDataBrokerServiceServer) Put(arg0 context.Context, arg1 *databroker.PutRequest) (*databroker.PutResponse, error) {
 	m.ctrl.T.Helper()

pkg/storage/storage.go

@@ -11,6 +11,7 @@ import (
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/fieldmaskpb"
 
 	"github.com/pomerium/pomerium/internal/log"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
@@ -37,6 +38,8 @@ type Backend interface {
 	ListTypes(ctx context.Context) ([]string, error)
 	// Put is used to insert or update records.
 	Put(ctx context.Context, records []*databroker.Record) (serverVersion uint64, err error)
+	// Patch is used to update specific fields of existing records.
+	Patch(ctx context.Context, records []*databroker.Record, fields *fieldmaskpb.FieldMask) (serverVersion uint64, patchedRecords []*databroker.Record, err error)
 	// SetOptions sets the options for a type.
 	SetOptions(ctx context.Context, recordType string, options *databroker.Options) error
 	// Sync syncs record changes after the specified version.

pkg/storage/storagetest/storagetest.go

@@ -4,6 +4,11 @@ package storagetest
 
 import (
 	"context"
+	"fmt"
+	"os"
+	"runtime"
+	"strconv"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -18,15 +23,8 @@ import (
 	"github.com/pomerium/pomerium/pkg/storage"
 )
 
-// BackendWithPatch is a storage.Backend with an additional Patch() method.
-// TODO: delete this type once Patch() is added to the storage.Backend interface
-type BackendWithPatch interface {
-	storage.Backend
-	Patch(context.Context, []*databroker.Record, *fieldmaskpb.FieldMask) (uint64, []*databroker.Record, error)
-}
-
 // TestBackendPatch verifies the behavior of the backend Patch() method.
-func TestBackendPatch(t *testing.T, ctx context.Context, backend BackendWithPatch) { //nolint:revive
+func TestBackendPatch(t *testing.T, ctx context.Context, backend storage.Backend) { //nolint:revive
 	mkRecord := func(s *session.Session) *databroker.Record {
 		a, _ := anypb.New(s)
 		return &databroker.Record{
@@ -36,6 +34,7 @@ func TestBackendPatch(t *testing.T, ctx context.Context, backend BackendWithPatc
 		}
 	}
 
+	t.Run("basic", func(t *testing.T) {
 		// Populate an initial set of session records.
 		s1 := &session.Session{
 			Id:         "session-1",
@@ -113,6 +112,63 @@ func TestBackendPatch(t *testing.T, ctx context.Context, backend BackendWithPatc
 		r3, _ := backend.Get(ctx, "type.googleapis.com/session.Session", "session-3")
 		truncateTimestamps(updated[1].ModifiedAt, r3.ModifiedAt)
 		testutil.AssertProtoEqual(t, updated[1], r3)
+	})
+
+	t.Run("concurrent", func(t *testing.T) {
+		if n := gomaxprocs(); n < 2 {
+			t.Skipf("skipping concurrent test (GOMAXPROCS = %d)", n)
+		}
+
+		rs1 := make([]*databroker.Record, 1)
+		rs2 := make([]*databroker.Record, 1)
+
+		s1 := session.Session{Id: "concurrent", OauthToken: &session.OAuthToken{}}
+		s2 := session.Session{Id: "concurrent", OauthToken: &session.OAuthToken{}}
+
+		// Store an initial version of a session record.
+		rs1[0] = mkRecord(&s1)
+		_, err := backend.Put(ctx, rs1)
+		require.NoError(t, err)
+
+		fmAccessToken, err := fieldmaskpb.New(&session.Session{}, "oauth_token.access_token")
+		require.NoError(t, err)
+		fmRefreshToken, err := fieldmaskpb.New(&session.Session{}, "oauth_token.refresh_token")
+		require.NoError(t, err)
+
+		var wg sync.WaitGroup
+
+		// Repeatedly make Patch calls to update the session from two separate
+		// goroutines (one updating just the access token, the other updating
+		// just the refresh token.) Verify that no updates are lost.
+		for i := 0; i < 100; i++ {
+			access := fmt.Sprintf("access-%d", i)
+			s1.OauthToken.AccessToken = access
+			rs1[0] = mkRecord(&s1)
+
+			refresh := fmt.Sprintf("refresh-%d", i)
+			s2.OauthToken.RefreshToken = refresh
+			rs2[0] = mkRecord(&s2)
+
+			wg.Add(2)
+			go func() {
+				_, _, _ = backend.Patch(ctx, rs1, fmAccessToken)
+				wg.Done()
+			}()
+			go func() {
+				_, _, _ = backend.Patch(ctx, rs2, fmRefreshToken)
+				wg.Done()
+			}()
+			wg.Wait()
+
+			r, err := backend.Get(ctx, "type.googleapis.com/session.Session", "concurrent")
+			require.NoError(t, err)
+			data, err := r.Data.UnmarshalNew()
+			require.NoError(t, err)
+			s := data.(*session.Session)
+			require.Equal(t, access, s.OauthToken.AccessToken)
+			require.Equal(t, refresh, s.OauthToken.RefreshToken)
+		}
+	})
 }
 
 // truncateTimestamps truncates Timestamp messages to 1 µs precision.
@@ -121,3 +177,11 @@ func truncateTimestamps(ts ...*timestamppb.Timestamp) {
 		t.Nanos = (t.Nanos / 1000) * 1000
 	}
 }
+
+func gomaxprocs() int {
+	env := os.Getenv("GOMAXPROCS")
+	if n, err := strconv.Atoi(env); err == nil {
+		return n
+	}
+	return runtime.NumCPU()
+}