diff --git a/fleetspeak/src/server/admin/admin.go b/fleetspeak/src/server/admin/admin.go index 11dc31ad..5ac9aac6 100644 --- a/fleetspeak/src/server/admin/admin.go +++ b/fleetspeak/src/server/admin/admin.go @@ -41,15 +41,27 @@ const ( maxMsgSize = 2 << 20 // 2MiB ) +// Options are optional settings for the admin server. +type Options struct { + RestrictPayloadRetrieval bool // Whether to allow retrieving payloads of pending messages. +} + // NewServer returns an admin_grpc.AdminServer which performs operations using // the provided db.Store. func NewServer(s db.Store, n notifications.Notifier) sgrpc.AdminServer { + return NewServerWithOptions(s, n, Options{}) +} + +// NewServerWithOptions returns an admin_grpc.AdminServer which performs +// operations using the provided db.Store and options. +func NewServerWithOptions(s db.Store, n notifications.Notifier, o Options) sgrpc.AdminServer { if n == nil { n = inotifications.NoopNotifier{} } return adminServer{ store: s, notifier: n, + options: o, } } @@ -59,6 +71,7 @@ type adminServer struct { store db.Store notifier notifications.Notifier + options Options } func (s adminServer) CreateBroadcast(ctx context.Context, req *spb.CreateBroadcastRequest) (*fspb.EmptyMessage, error) { @@ -289,6 +302,10 @@ func (s adminServer) DeletePendingMessages(ctx context.Context, r *spb.DeletePen } func (s adminServer) GetPendingMessages(ctx context.Context, r *spb.GetPendingMessagesRequest) (*spb.GetPendingMessagesResponse, error) { + if r.WantData && s.options.RestrictPayloadRetrieval { + return nil, errors.New("retrieving payloads of pending messages is not permitted") + } + ids, err := s.bytesToClientIds(r.ClientIds) if err != nil { return nil, err diff --git a/fleetspeak/src/server/servertests/admin_test.go b/fleetspeak/src/server/servertests/admin_test.go index a66116ba..4b280de4 100644 --- a/fleetspeak/src/server/servertests/admin_test.go +++ b/fleetspeak/src/server/servertests/admin_test.go @@ -488,6 +488,45 @@ func TestPendingMessages(t *testing.T) { } }) + t.Run("GetPendingMessages/RestrictPayloadRetrieval", func(t *testing.T) { + restrictedServer := admin.NewServerWithOptions(ts.DS, nil, admin.Options{ + RestrictPayloadRetrieval: true, + }) + + t.Run("WantData=true", func(t *testing.T) { + greq := &spb.GetPendingMessagesRequest{ + ClientIds: [][]byte{id.Bytes()}, + WantData: true, + } + if _, err := restrictedServer.GetPendingMessages(ctx, greq); err == nil { + t.Fatal("GetPendingMessages(WantData=true) with RestrictPayloadRetrieval should fail, but got nil error") + } else if !strings.Contains(err.Error(), "retrieving payloads of pending messages is not permitted") { + t.Fatalf("Unexpected error message: %v", err) + } + }) + + t.Run("WantData=false", func(t *testing.T) { + greq := &spb.GetPendingMessagesRequest{ + ClientIds: [][]byte{id.Bytes()}, + WantData: false, + } + gresp, err := restrictedServer.GetPendingMessages(ctx, greq) + if err != nil { + t.Fatalf("GetPendingMessages returned error: %v", err) + } + if len(gresp.Messages) != len(msgs) { + t.Fatalf("Bad size of returned messages. Expected: %v. Got: %v.", len(msgs), len(gresp.Messages)) + } + for i, msg := range msgs { + expectedMessage := proto.Clone(msg).(*fspb.Message) + expectedMessage.Data = nil + if !proto.Equal(gresp.Messages[i], expectedMessage) { + t.Fatalf("Got bad message. Expected: [%v]. Got: [%v].", expectedMessage, gresp.Messages[i]) + } + } + }) + }) + // Delete the message from the pending messages list. t.Run("DeletePendingMessages", func(t *testing.T) {