SupportBundles: Do not allow multiple collectors with the same UID (#60581)

* unexport service methods

* fix typos

* make constants

* do not allow double register of support bundle collectors

* fix get collectors response

pkg/infra/supportbundles/supportbundlesimpl/api.go

@@ -36,7 +36,7 @@ func (s *Service) registerAPIEndpoints(routeRegister routing.RouteRegister) {
 }
 
 func (s *Service) handleList(ctx *models.ReqContext) response.Response {
-	bundles, err := s.List(ctx.Req.Context())
+	bundles, err := s.list(ctx.Req.Context())
 	if err != nil {
 		return response.Error(http.StatusInternalServerError, "failed to list bundles", err)
 	}
@@ -59,7 +59,7 @@ func (s *Service) handleCreate(ctx *models.ReqContext) response.Response {
 		return response.Error(http.StatusBadRequest, "failed to parse request", err)
 	}
 
-	bundle, err := s.Create(context.Background(), c.Collectors, ctx.SignedInUser)
+	bundle, err := s.create(context.Background(), c.Collectors, ctx.SignedInUser)
 	if err != nil {
 		return response.Error(http.StatusInternalServerError, "failed to create support bundle", err)
 	}
@@ -74,7 +74,7 @@ func (s *Service) handleCreate(ctx *models.ReqContext) response.Response {
 
 func (s *Service) handleDownload(ctx *models.ReqContext) {
 	uid := web.Params(ctx.Req)[":uid"]
-	bundle, err := s.Get(ctx.Req.Context(), uid)
+	bundle, err := s.get(ctx.Req.Context(), uid)
 	if err != nil {
 		ctx.Redirect("/admin/support-bundles")
 		return
@@ -102,7 +102,7 @@ func (s *Service) handleDownload(ctx *models.ReqContext) {
 
 func (s *Service) handleRemove(ctx *models.ReqContext) response.Response {
 	uid := web.Params(ctx.Req)[":uid"]
-	err := s.Remove(ctx.Req.Context(), uid)
+	err := s.remove(ctx.Req.Context(), uid)
 	if err != nil {
 		return response.Error(http.StatusInternalServerError, "failed to remove bundle", err)
 	}
@@ -111,5 +111,11 @@ func (s *Service) handleRemove(ctx *models.ReqContext) response.Response {
 }
 
 func (s *Service) handleGetCollectors(ctx *models.ReqContext) response.Response {
-	return response.JSON(http.StatusOK, s.collectors)
+	collectors := make([]supportbundles.Collector, 0, len(s.collectors))
+
+	for _, c := range s.collectors {
+		collectors = append(collectors, c)
+	}
+
+	return response.JSON(http.StatusOK, collectors)
 }

pkg/infra/supportbundles/supportbundlesimpl/models.go

@@ -35,7 +35,7 @@ var (
 	}
 )
 
-func DeclareFixedRoles(ac accesscontrol.Service) error {
+func declareFixedRoles(ac accesscontrol.Service) error {
 	bundleReader := accesscontrol.RoleRegistration{
 		Role:   bundleReaderRole,
 		Grants: []string{string(org.RoleAdmin)},

pkg/infra/supportbundles/supportbundlesimpl/service.go

@@ -21,9 +21,14 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 )
 
+const (
+	cleanUpInterval       = 24 * time.Hour
+	bundleCreationTimeout = 20 * time.Minute
+)
+
 type Service struct {
 	cfg            *setting.Cfg
-	store          *store
+	store          bundleStore
 	pluginStore    plugins.Store
 	pluginSettings pluginsettings.Service
 	accessControl  ac.AccessControl
@@ -31,7 +36,7 @@ type Service struct {
 
 	log log.Logger
 
-	collectors []supportbundles.Collector
+	collectors map[string]supportbundles.Collector
 }
 
 func ProvideService(cfg *setting.Cfg,
@@ -54,6 +59,7 @@ func ProvideService(cfg *setting.Cfg,
 		accessControl:  accessControl,
 		features:       features,
 		log:            log.New("supportbundle.service"),
+		collectors:     make(map[string]supportbundles.Collector),
 	}
 
 	if !features.IsEnabled(featuremgmt.FlagSupportBundles) {
@@ -61,7 +67,7 @@ func ProvideService(cfg *setting.Cfg,
 	}
 
 	if !accessControl.IsDisabled() {
-		if err := DeclareFixedRoles(accesscontrolService); err != nil {
+		if err := declareFixedRoles(accesscontrolService); err != nil {
 			return nil, err
 		}
 	}
@@ -79,14 +85,39 @@ func ProvideService(cfg *setting.Cfg,
 	return s, nil
 }
 
-func (s *Service) Create(ctx context.Context, collectors []string, usr *user.SignedInUser) (*supportbundles.Bundle, error) {
+func (s *Service) RegisterSupportItemCollector(collector supportbundles.Collector) {
+	if _, ok := s.collectors[collector.UID]; ok {
+		s.log.Warn("Support bundle collector with the same UID already registered", "uid", collector.UID)
+	}
+
+	s.collectors[collector.UID] = collector
+}
+
+func (s *Service) Run(ctx context.Context) error {
+	if !s.features.IsEnabled(featuremgmt.FlagSupportBundles) {
+		return nil
+	}
+
+	ticker := time.NewTicker(cleanUpInterval)
+	defer ticker.Stop()
+	s.cleanup(ctx)
+	select {
+	case <-ticker.C:
+		s.cleanup(ctx)
+	case <-ctx.Done():
+		break
+	}
+	return ctx.Err()
+}
+
+func (s *Service) create(ctx context.Context, collectors []string, usr *user.SignedInUser) (*supportbundles.Bundle, error) {
 	bundle, err := s.store.Create(ctx, usr)
 	if err != nil {
 		return nil, err
 	}
 
 	go func(uid string, collectors []string) {
-		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
+		ctx, cancel := context.WithTimeout(context.Background(), bundleCreationTimeout)
 		defer cancel()
 		s.startBundleWork(ctx, collectors, uid)
 	}(bundle.UID, collectors)
@@ -94,24 +125,24 @@ func (s *Service) Create(ctx context.Context, collectors []string, usr *user.Sig
 	return bundle, nil
 }
 
-func (s *Service) Get(ctx context.Context, uid string) (*supportbundles.Bundle, error) {
+func (s *Service) get(ctx context.Context, uid string) (*supportbundles.Bundle, error) {
 	return s.store.Get(ctx, uid)
 }
 
-func (s *Service) List(ctx context.Context) ([]supportbundles.Bundle, error) {
+func (s *Service) list(ctx context.Context) ([]supportbundles.Bundle, error) {
 	return s.store.List()
 }
 
-func (s *Service) Remove(ctx context.Context, uid string) error {
+func (s *Service) remove(ctx context.Context, uid string) error {
 	// Remove the data
 	bundle, err := s.store.Get(ctx, uid)
 	if err != nil {
-		return fmt.Errorf("could not retrieve support bundle with uid %s: %w", uid, err)
+		return fmt.Errorf("could not retrieve support bundle with UID %s: %w", uid, err)
 	}
 
 	// TODO handle cases when bundles aren't complete yet
 	if bundle.State == supportbundles.StatePending {
-		return fmt.Errorf("could not remove a support bundle with uid %s as it is still beign cteated", uid)
+		return fmt.Errorf("could not remove a support bundle with uid %s as it is still being created", uid)
 	}
 
 	if bundle.FilePath != "" {
@@ -124,30 +155,8 @@ func (s *Service) Remove(ctx context.Context, uid string) error {
 	return s.store.Remove(ctx, uid)
 }
 
-func (s *Service) RegisterSupportItemCollector(collector supportbundles.Collector) {
-	// FIXME: add check for duplicate UIDs
-	s.collectors = append(s.collectors, collector)
-}
-
-func (s *Service) Run(ctx context.Context) error {
-	if !s.features.IsEnabled(featuremgmt.FlagSupportBundles) {
-		return nil
-	}
-
-	ticker := time.NewTicker(24 * time.Hour)
-	defer ticker.Stop()
-	s.cleanup(ctx)
-	select {
-	case <-ticker.C:
-		s.cleanup(ctx)
-	case <-ctx.Done():
-		break
-	}
-	return ctx.Err()
-}
-
 func (s *Service) cleanup(ctx context.Context) {
-	bundles, err := s.List(ctx)
+	bundles, err := s.list(ctx)
 	if err != nil {
 		s.log.Error("failed to list bundles to clean up", "error", err)
 	}
@@ -155,7 +164,7 @@ func (s *Service) cleanup(ctx context.Context) {
 	if err == nil {
 		for _, b := range bundles {
 			if time.Now().Unix() >= b.ExpiresAt {
-				if err := s.Remove(ctx, b.UID); err != nil {
+				if err := s.remove(ctx, b.UID); err != nil {
 					s.log.Error("failed to cleanup bundle", "error", err)
 				}
 			}

pkg/infra/supportbundles/supportbundlesimpl/service_test.go

@@ -0,0 +1,44 @@
+package supportbundlesimpl
+
+import (
+	"context"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/supportbundles"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/require"
+)
+
+func TestService_RegisterSupportItemCollector(t *testing.T) {
+	s := &Service{
+		cfg:            &setting.Cfg{},
+		store:          nil,
+		pluginStore:    nil,
+		pluginSettings: nil,
+		accessControl:  nil,
+		features:       nil,
+		log:            log.NewNopLogger(),
+		collectors:     map[string]supportbundles.Collector{},
+	}
+	collector := supportbundles.Collector{
+		UID:               "test",
+		DisplayName:       "test",
+		Description:       "test",
+		IncludedByDefault: true,
+		Default:           true,
+		Fn: func(context.Context) (*supportbundles.SupportItem, error) {
+			return nil, nil
+		},
+	}
+
+	t.Run("should register collector", func(t *testing.T) {
+		s.RegisterSupportItemCollector(collector)
+		require.Len(t, s.collectors, 1)
+	})
+
+	t.Run("should not register collector with same UID", func(t *testing.T) {
+		s.RegisterSupportItemCollector(collector)
+		require.Len(t, s.collectors, 1)
+	})
+}

pkg/infra/supportbundles/supportbundlesimpl/store.go

@@ -14,6 +14,10 @@ import (
 	"github.com/grafana/grafana/pkg/services/user"
 )
 
+const (
+	defaultBundleExpiration = 72 * time.Hour // 72h
+)
+
 func newStore(kv kvstore.KVStore) *store {
 	return &store{kv: kvstore.WithNamespace(kv, 0, "supportbundle")}
 }
@@ -22,6 +26,14 @@ type store struct {
 	kv *kvstore.NamespacedKVStore
 }
 
+type bundleStore interface {
+	Create(ctx context.Context, usr *user.SignedInUser) (*supportbundles.Bundle, error)
+	Get(ctx context.Context, uid string) (*supportbundles.Bundle, error)
+	List() ([]supportbundles.Bundle, error)
+	Remove(ctx context.Context, uid string) error
+	Update(ctx context.Context, uid string, state supportbundles.State, filePath string) error
+}
+
 func (s *store) Create(ctx context.Context, usr *user.SignedInUser) (*supportbundles.Bundle, error) {
 	uid, err := uuid.NewRandom()
 	if err != nil {
@@ -33,7 +45,7 @@ func (s *store) Create(ctx context.Context, usr *user.SignedInUser) (*supportbun
 		State:     supportbundles.StatePending,
 		Creator:   usr.Login,
 		CreatedAt: time.Now().Unix(),
-		ExpiresAt: time.Now().Add(24 * time.Hour).Unix(),
+		ExpiresAt: time.Now().Add(defaultBundleExpiration).Unix(),
 	}
 
 	if err := s.set(ctx, &bundle); err != nil {