Add optional settings to the admin server

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

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) {