databroker: add support for querying the databroker (#1443)

* databroker: add support for querying the databroker

* remove query method, use getall so encryption works

* add test

* return early

pkg/storage/redis/redis_test.go

@@ -19,6 +19,7 @@ import (
 	"google.golang.org/protobuf/types/known/anypb"
 
 	"github.com/pomerium/pomerium/pkg/cryptutil"
+	"github.com/pomerium/pomerium/pkg/grpc/directory"
 )
 
 var db *DB
@@ -120,6 +121,11 @@ func testDB(t *testing.T) {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	defer cancelFunc()
 
+	users := []*directory.User{
+		{Id: "u1", GroupIds: []string{"test", "admin"}},
+		{Id: "u2"},
+		{Id: "u3", GroupIds: []string{"test"}},
+	}
 	ids := []string{"a", "b", "c"}
 	id := ids[0]
 	c := db.pool.Get()
@@ -133,13 +139,13 @@ func testDB(t *testing.T) {
 		assert.Nil(t, record)
 	})
 	t.Run("get record", func(t *testing.T) {
-		data := new(anypb.Any)
+		data, _ := anypb.New(users[0])
 		assert.NoError(t, db.Put(ctx, id, data))
 		record, err := db.Get(ctx, id)
 		require.NoError(t, err)
 		if assert.NotNil(t, record) {
 			assert.NotNil(t, record.CreatedAt)
-			assert.Equal(t, data, record.Data)
+			assert.NotEmpty(t, record.Data)
 			assert.Nil(t, record.DeletedAt)
 			assert.Equal(t, "a", record.Id)
 			assert.NotNil(t, record.ModifiedAt)
@@ -163,9 +169,9 @@ func testDB(t *testing.T) {
 		records, err := db.GetAll(ctx)
 		assert.NoError(t, err)
 		assert.Len(t, records, 0)
-		data := new(anypb.Any)
 
-		for _, id := range ids {
+		for i, id := range ids {
+			data, _ := anypb.New(users[i])
 			assert.NoError(t, db.Put(ctx, id, data))
 		}
 		records, err = db.GetAll(ctx)

pkg/storage/storage.go

@@ -3,10 +3,13 @@ package storage
 
 import (
 	"context"
+	"strings"
 	"time"
 
+	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/known/anypb"
 
+	"github.com/pomerium/pomerium/internal/log"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 )
 
@@ -38,3 +41,74 @@ type Backend interface {
 	// the channel.
 	Watch(ctx context.Context) <-chan struct{}
 }
+
+// MatchAny searches any data with a query.
+func MatchAny(any *anypb.Any, query string) bool {
+	if any == nil {
+		return false
+	}
+
+	msg, err := any.UnmarshalNew()
+	if err != nil {
+		// ignore invalid any types
+		log.Error().Err(err).Msg("storage: invalid any type")
+		return false
+	}
+
+	// search by query
+	return matchProtoMessage(msg.ProtoReflect(), query)
+}
+
+func matchProtoMessage(msg protoreflect.Message, query string) bool {
+	md := msg.Descriptor()
+	fds := md.Fields()
+	for i := 0; i < fds.Len(); i++ {
+		fd := fds.Get(i)
+		if !msg.Has(fd) {
+			continue
+		}
+		if matchProtoValue(fd, msg.Get(fd), query) {
+			return true
+		}
+	}
+	return false
+}
+
+func matchProtoValue(fd protoreflect.FieldDescriptor, v protoreflect.Value, query string) bool {
+	switch {
+	case fd.IsList():
+		return matchProtoListValue(fd, v.List(), query)
+	case fd.IsMap():
+		return matchProtoMapValue(fd, v.Map(), query)
+	default:
+		return matchProtoSingularValue(fd, v, query)
+	}
+}
+
+func matchProtoSingularValue(fd protoreflect.FieldDescriptor, v protoreflect.Value, query string) bool {
+	switch fd.Kind() {
+	case protoreflect.MessageKind:
+		return matchProtoMessage(v.Message(), query)
+	case protoreflect.StringKind:
+		return strings.Contains(strings.ToLower(v.String()), query)
+	}
+	return false
+}
+
+func matchProtoListValue(fd protoreflect.FieldDescriptor, l protoreflect.List, query string) bool {
+	for i := 0; i < l.Len(); i++ {
+		if matchProtoSingularValue(fd, l.Get(i), query) {
+			return true
+		}
+	}
+	return false
+}
+
+func matchProtoMapValue(fd protoreflect.FieldDescriptor, m protoreflect.Map, query string) bool {
+	matches := false
+	m.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool {
+		matches = matches || matchProtoSingularValue(fd, v, query)
+		return !matches
+	})
+	return matches
+}

pkg/storage/storage_test.go

@@ -2,11 +2,14 @@ package storage
 
 import (
 	"context"
+	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/types/known/anypb"
 
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/grpc/user"
 )
 
 type mockBackend struct {
@@ -16,6 +19,7 @@ type mockBackend struct {
 	list         func(ctx context.Context, sinceVersion string) ([]*databroker.Record, error)
 	delete       func(ctx context.Context, id string) error
 	clearDeleted func(ctx context.Context, cutoff time.Time)
+	query        func(ctx context.Context, query string, offset, limit int) ([]*databroker.Record, int, error)
 	watch        func(ctx context.Context) <-chan struct{}
 }
 
@@ -47,6 +51,20 @@ func (m *mockBackend) ClearDeleted(ctx context.Context, cutoff time.Time) {
 	m.clearDeleted(ctx, cutoff)
 }
 
+func (m *mockBackend) Query(ctx context.Context, query string, offset, limit int) ([]*databroker.Record, int, error) {
+	return m.query(ctx, query, offset, limit)
+}
+
 func (m *mockBackend) Watch(ctx context.Context) <-chan struct{} {
 	return m.watch(ctx)
 }
+
+func TestMatchAny(t *testing.T) {
+	u := &user.User{Id: "id", Name: "name", Email: "email"}
+	data, _ := anypb.New(u)
+	assert.True(t, MatchAny(data, ""))
+	assert.True(t, MatchAny(data, "id"))
+	assert.True(t, MatchAny(data, "name"))
+	assert.True(t, MatchAny(data, "email"))
+	assert.False(t, MatchAny(data, "nope"))
+}