mirror of
https://github.com/pomerium/pomerium.git
synced 2025-07-04 10:28:10 +02:00
5e42bfbb34
Add a new Patch() method that updates specific fields of an existing record's data, based on a field mask. Extract some logic from the existing Get() and Put() methods so it can be shared with the new Patch() method.
342 lines
9 KiB
Go
342 lines
9 KiB
Go
package inmemory
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/sync/errgroup"
|
|
"google.golang.org/protobuf/proto"
|
|
"google.golang.org/protobuf/types/known/anypb"
|
|
"google.golang.org/protobuf/types/known/fieldmaskpb"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
|
|
"github.com/pomerium/pomerium/internal/testutil"
|
|
"github.com/pomerium/pomerium/pkg/grpc/databroker"
|
|
"github.com/pomerium/pomerium/pkg/grpc/session"
|
|
"github.com/pomerium/pomerium/pkg/storage"
|
|
)
|
|
|
|
func TestBackend(t *testing.T) {
|
|
ctx := context.Background()
|
|
backend := New()
|
|
defer func() { _ = backend.Close() }()
|
|
t.Run("get missing record", func(t *testing.T) {
|
|
record, err := backend.Get(ctx, "TYPE", "abcd")
|
|
require.Error(t, err)
|
|
assert.Nil(t, record)
|
|
})
|
|
t.Run("get record", func(t *testing.T) {
|
|
data := new(anypb.Any)
|
|
sv, err := backend.Put(ctx, []*databroker.Record{
|
|
{
|
|
Type: "TYPE",
|
|
Id: "a",
|
|
Data: data,
|
|
},
|
|
{
|
|
Type: "TYPE",
|
|
Id: "b",
|
|
Data: data,
|
|
},
|
|
{
|
|
Type: "TYPE",
|
|
Id: "c",
|
|
Data: data,
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, backend.serverVersion, sv)
|
|
for i, id := range []string{"a", "b", "c"} {
|
|
record, err := backend.Get(ctx, "TYPE", id)
|
|
require.NoError(t, err)
|
|
if assert.NotNil(t, record) {
|
|
assert.Equal(t, data, record.Data)
|
|
assert.Nil(t, record.DeletedAt)
|
|
assert.Equal(t, id, record.Id)
|
|
assert.NotNil(t, record.ModifiedAt)
|
|
assert.Equal(t, "TYPE", record.Type)
|
|
assert.Equal(t, uint64(i+1), record.Version)
|
|
}
|
|
}
|
|
})
|
|
t.Run("delete record", func(t *testing.T) {
|
|
sv, err := backend.Put(ctx, []*databroker.Record{{
|
|
Type: "TYPE",
|
|
Id: "abcd",
|
|
DeletedAt: timestamppb.Now(),
|
|
}})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, backend.serverVersion, sv)
|
|
record, err := backend.Get(ctx, "TYPE", "abcd")
|
|
assert.Error(t, err)
|
|
assert.Nil(t, record)
|
|
})
|
|
}
|
|
|
|
func TestPatch(t *testing.T) {
|
|
ctx := context.Background()
|
|
backend := New()
|
|
t.Cleanup(func() { _ = backend.Close() })
|
|
|
|
mkRecord := func(s *session.Session) *databroker.Record {
|
|
a, _ := anypb.New(s)
|
|
return &databroker.Record{
|
|
Type: a.TypeUrl,
|
|
Id: s.Id,
|
|
Data: a,
|
|
}
|
|
}
|
|
|
|
// Populate an initial set of session records.
|
|
s1 := &session.Session{
|
|
Id: "session-1",
|
|
IdToken: &session.IDToken{Issuer: "issuer-1"},
|
|
OauthToken: &session.OAuthToken{AccessToken: "access-token-1"},
|
|
}
|
|
s2 := &session.Session{
|
|
Id: "session-2",
|
|
IdToken: &session.IDToken{Issuer: "issuer-2"},
|
|
OauthToken: &session.OAuthToken{AccessToken: "access-token-2"},
|
|
}
|
|
s3 := &session.Session{
|
|
Id: "session-3",
|
|
IdToken: &session.IDToken{Issuer: "issuer-3"},
|
|
OauthToken: &session.OAuthToken{AccessToken: "access-token-3"},
|
|
}
|
|
initial := []*databroker.Record{mkRecord(s1), mkRecord(s2), mkRecord(s3)}
|
|
|
|
_, err := backend.Put(ctx, initial)
|
|
require.NoError(t, err)
|
|
|
|
// Now patch just the oauth_token field.
|
|
u1 := &session.Session{
|
|
Id: "session-1",
|
|
OauthToken: &session.OAuthToken{AccessToken: "access-token-1-new"},
|
|
}
|
|
u2 := &session.Session{
|
|
Id: "session-4-does-not-exist",
|
|
OauthToken: &session.OAuthToken{AccessToken: "access-token-4-new"},
|
|
}
|
|
u3 := &session.Session{
|
|
Id: "session-3",
|
|
OauthToken: &session.OAuthToken{AccessToken: "access-token-3-new"},
|
|
}
|
|
|
|
mask, _ := fieldmaskpb.New(&session.Session{}, "oauth_token")
|
|
|
|
_, updated, err := backend.Patch(
|
|
ctx, []*databroker.Record{mkRecord(u1), mkRecord(u2), mkRecord(u3)}, mask)
|
|
require.NoError(t, err)
|
|
|
|
// The OAuthToken message should be updated but the IDToken message should
|
|
// be unchanged, as it was not included in the field mask. The results
|
|
// should indicate that only two records were updated (one did not exist).
|
|
assert.Equal(t, 2, len(updated))
|
|
assert.Greater(t, updated[0].Version, initial[0].Version)
|
|
assert.Greater(t, updated[1].Version, initial[2].Version)
|
|
testutil.AssertProtoJSONEqual(t, `{
|
|
"@type": "type.googleapis.com/session.Session",
|
|
"id": "session-1",
|
|
"idToken": {
|
|
"issuer": "issuer-1"
|
|
},
|
|
"oauthToken": {
|
|
"accessToken": "access-token-1-new"
|
|
}
|
|
}`, updated[0].Data)
|
|
testutil.AssertProtoJSONEqual(t, `{
|
|
"@type": "type.googleapis.com/session.Session",
|
|
"id": "session-3",
|
|
"idToken": {
|
|
"issuer": "issuer-3"
|
|
},
|
|
"oauthToken": {
|
|
"accessToken": "access-token-3-new"
|
|
}
|
|
}`, updated[1].Data)
|
|
|
|
// Verify that the updates will indeed be seen by a subsequent Get().
|
|
r1, _ := backend.Get(ctx, "type.googleapis.com/session.Session", "session-1")
|
|
testutil.AssertProtoEqual(t, updated[0], r1)
|
|
r3, _ := backend.Get(ctx, "type.googleapis.com/session.Session", "session-3")
|
|
testutil.AssertProtoEqual(t, updated[1], r3)
|
|
}
|
|
|
|
func TestExpiry(t *testing.T) {
|
|
ctx := context.Background()
|
|
backend := New(WithExpiry(0))
|
|
defer func() { _ = backend.Close() }()
|
|
|
|
for i := 0; i < 1000; i++ {
|
|
sv, err := backend.Put(ctx, []*databroker.Record{{
|
|
Type: "TYPE",
|
|
Id: fmt.Sprint(i),
|
|
}})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, backend.serverVersion, sv)
|
|
}
|
|
stream, err := backend.Sync(ctx, "", backend.serverVersion, 0)
|
|
require.NoError(t, err)
|
|
var records []*databroker.Record
|
|
for stream.Next(false) {
|
|
records = append(records, stream.Record())
|
|
}
|
|
_ = stream.Close()
|
|
require.Len(t, records, 1000)
|
|
|
|
backend.removeChangesBefore(time.Now().Add(time.Second))
|
|
|
|
stream, err = backend.Sync(ctx, "", backend.serverVersion, 0)
|
|
require.NoError(t, err)
|
|
records = nil
|
|
for stream.Next(false) {
|
|
records = append(records, stream.Record())
|
|
}
|
|
_ = stream.Close()
|
|
require.Len(t, records, 0)
|
|
}
|
|
|
|
func TestConcurrency(t *testing.T) {
|
|
ctx := context.Background()
|
|
backend := New()
|
|
|
|
eg, ctx := errgroup.WithContext(ctx)
|
|
eg.Go(func() error {
|
|
for i := 0; i < 1000; i++ {
|
|
_, _ = backend.Get(ctx, "", fmt.Sprint(i))
|
|
}
|
|
return nil
|
|
})
|
|
eg.Go(func() error {
|
|
for i := 0; i < 1000; i++ {
|
|
_, _ = backend.Put(ctx, []*databroker.Record{{
|
|
Id: fmt.Sprint(i),
|
|
}})
|
|
}
|
|
return nil
|
|
})
|
|
assert.NoError(t, eg.Wait())
|
|
}
|
|
|
|
func TestStream(t *testing.T) {
|
|
ctx := context.Background()
|
|
backend := New()
|
|
defer func() { _ = backend.Close() }()
|
|
|
|
stream, err := backend.Sync(ctx, "TYPE", backend.serverVersion, 0)
|
|
require.NoError(t, err)
|
|
defer func() { _ = stream.Close() }()
|
|
|
|
eg, ctx := errgroup.WithContext(ctx)
|
|
eg.Go(func() error {
|
|
for i := 0; i < 10000; i++ {
|
|
assert.True(t, stream.Next(true))
|
|
assert.Nil(t, stream.Err())
|
|
assert.Equal(t, "TYPE", stream.Record().GetType())
|
|
assert.Equal(t, fmt.Sprint(i), stream.Record().GetId())
|
|
assert.Equal(t, uint64(i+1), stream.Record().GetVersion())
|
|
}
|
|
return nil
|
|
})
|
|
eg.Go(func() error {
|
|
for i := 0; i < 10000; i++ {
|
|
_, err := backend.Put(ctx, []*databroker.Record{{
|
|
Type: "TYPE",
|
|
Id: fmt.Sprint(i),
|
|
}})
|
|
assert.NoError(t, err)
|
|
}
|
|
return nil
|
|
})
|
|
require.NoError(t, eg.Wait())
|
|
}
|
|
|
|
func TestStreamClose(t *testing.T) {
|
|
ctx := context.Background()
|
|
t.Run("by backend", func(t *testing.T) {
|
|
backend := New()
|
|
stream, err := backend.Sync(ctx, "", backend.serverVersion, 0)
|
|
require.NoError(t, err)
|
|
require.NoError(t, backend.Close())
|
|
assert.False(t, stream.Next(true))
|
|
assert.Error(t, stream.Err())
|
|
})
|
|
t.Run("by stream", func(t *testing.T) {
|
|
backend := New()
|
|
stream, err := backend.Sync(ctx, "", backend.serverVersion, 0)
|
|
require.NoError(t, err)
|
|
require.NoError(t, stream.Close())
|
|
assert.False(t, stream.Next(true))
|
|
assert.Error(t, stream.Err())
|
|
})
|
|
t.Run("by context", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
backend := New()
|
|
stream, err := backend.Sync(ctx, "", backend.serverVersion, 0)
|
|
require.NoError(t, err)
|
|
cancel()
|
|
assert.False(t, stream.Next(true))
|
|
assert.Error(t, stream.Err())
|
|
})
|
|
}
|
|
|
|
func TestCapacity(t *testing.T) {
|
|
ctx := context.Background()
|
|
backend := New()
|
|
defer func() { _ = backend.Close() }()
|
|
|
|
err := backend.SetOptions(ctx, "EXAMPLE", &databroker.Options{
|
|
Capacity: proto.Uint64(3),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
for i := 0; i < 10; i++ {
|
|
_, err = backend.Put(ctx, []*databroker.Record{{
|
|
Type: "EXAMPLE",
|
|
Id: fmt.Sprint(i),
|
|
}})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
_, _, stream, err := backend.SyncLatest(ctx, "EXAMPLE", nil)
|
|
require.NoError(t, err)
|
|
|
|
records, err := storage.RecordStreamToList(stream)
|
|
require.NoError(t, err)
|
|
assert.Len(t, records, 3)
|
|
|
|
var ids []string
|
|
for _, r := range records {
|
|
ids = append(ids, r.GetId())
|
|
}
|
|
assert.Equal(t, []string{"7", "8", "9"}, ids, "should contain recent records")
|
|
}
|
|
|
|
func TestLease(t *testing.T) {
|
|
ctx := context.Background()
|
|
backend := New()
|
|
{
|
|
ok, err := backend.Lease(ctx, "test", "a", time.Second*30)
|
|
require.NoError(t, err)
|
|
assert.True(t, ok, "expected a to acquire the lease")
|
|
}
|
|
{
|
|
ok, err := backend.Lease(ctx, "test", "b", time.Second*30)
|
|
require.NoError(t, err)
|
|
assert.False(t, ok, "expected b to fail to acquire the lease")
|
|
}
|
|
{
|
|
ok, err := backend.Lease(ctx, "test", "a", 0)
|
|
require.NoError(t, err)
|
|
assert.False(t, ok, "expected a to clear the lease")
|
|
}
|
|
{
|
|
ok, err := backend.Lease(ctx, "test", "b", time.Second*30)
|
|
require.NoError(t, err)
|
|
assert.True(t, ok, "expected b to to acquire the lease")
|
|
}
|
|
}
|