stackrox · dhaus67 · Dec 8, 2022 · Dec 8, 2022 · Dec 8, 2022 · Dec 8, 2022
@@ -0,0 +1,165 @@
+package m112tom113
+
+import (
+	"strings"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/pkg/errors"
+	"github.com/stackrox/rox/generated/storage"
+	"github.com/stackrox/rox/migrator/log"
+	"github.com/stackrox/rox/migrator/migrations"
+	"github.com/stackrox/rox/migrator/types"
+	"github.com/stackrox/rox/pkg/errorhelpers"
+	"github.com/stackrox/rox/pkg/uuid"
+	bolt "go.etcd.io/bbolt"
+)
+
+type groupEntry struct {
+	key   []byte
+	value []byte
+}
+
+var (
+	bucketName = []byte("groups2")
+
+	migration = types.Migration{
+		StartingSeqNum: 112,
+		VersionAfter:   &storage.Version{SeqNum: 113},
+		Run: func(databases *types.Databases) error {
+			return recreateGroupsBucket(databases.BoltDB)
+		},
+	}
+)
+
+func init() {
+	migrations.MustRegisterMigration(migration)
+}
+
+func recreateGroupsBucket(db *bolt.DB) error {
+	// Short-circuit if the bucket does not exist.
+	exists, err := checkGroupBucketExists(db)
+	if err != nil {
+		return errors.Wrap(err, "error checking if groups bucket exists")
+	}
+	if !exists {
+		log.WriteToStderr("groups bucket did not exist, hence no re-creation of the groups bucket was done.")
+		return nil
+	}
+
+	groupEntries, err := fetchGroupsBucket(db)
+	if err != nil {
+		return errors.Wrap(err, "error fetching groups to recreate")
+	}
+
+	// Drop the bucket.
+	if err := dropGroupsBucket(db); err != nil {
+		return errors.Wrap(err, "error dropping groups bucket")
+	}
+
+	// Create groups bucket and filter out invalid entries.
+	if err := createGroupsBucket(db, groupEntries); err != nil {
+		return errors.Wrap(err, "error recreating groups bucket")
+	}
+
+	return nil
+}
+
+func fetchGroupsBucket(db *bolt.DB) (groupEntries []groupEntry, err error) {
+	err = db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(bucketName)
+		// We previously checked that the bucket should be available, but still add this check here.
+		if bucket == nil {
+			return nil
+		}
+		return bucket.ForEach(func(k, v []byte) error {
+			groupEntries = append(groupEntries, groupEntry{key: k, value: v})
+			return nil
+		})
+	})
+	if err != nil {
+		return nil, err
+	}
+	return groupEntries, nil
+}
+
+func dropGroupsBucket(db *bolt.DB) error {
+	return db.Update(func(tx *bolt.Tx) error {
+		return tx.DeleteBucket(bucketName)
+	})
+}
+
+func createGroupsBucket(db *bolt.DB, groupEntries []groupEntry) (err error) {
+	err = db.Update(func(tx *bolt.Tx) error {
+		// Explicitly use the CreateBucket here instead of CreateBucketIfNotExists, as we require the bucket to be
+		// previously dropped.
+		bucket, err := tx.CreateBucket(bucketName)
+		if err != nil {
+			return err
+		}
+
+		var putGroupErrs errorhelpers.ErrorList
+		for _, entry := range groupEntries {
+			// After migration 105_to_106, we can assume that the key will be a UUID and the value will be the group
+			// proto message.
+			// Here, we will check that the key will be a string and can be parsed as a UUID.
+			// If that's the case, the entry is valid, and we will add it to the re-created bucket.
+			// If not, we will log the invalid entry that will be dropped.
+			if !verifyKeyValuePair(entry.key, entry.value) {
+				log.WriteToStderrf("Invalid group entry found in groups bucket (key=%s, value=%s). This entry"+
+					" will be dropped.",
+					entry.key, entry.value)
+				continue
+			}
+
+			if err := bucket.Put(entry.key, entry.value); err != nil {
+				putGroupErrs.AddError(err)
+			}
+		}
+
+		return putGroupErrs.ToError()
+	})
+	return err
+}
+
+func checkGroupBucketExists(db *bolt.DB) (exists bool, err error) {
+	exists = true
+	err = db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(bucketName)
+		if bucket == nil {
+			exists = false
+		}
+		return nil
+	})
+	return exists, err
+}
+
+const (
+	// Value has been taken from:
+	//	https://github.com/stackrox/stackrox/blob/6a702b26d66dcc2236a742907809071249187070/central/group/datastore/validate.go#L13
+	groupIDPrefix = "io.stackrox.authz.group."
+	// Value has been taken from:
+	//	https://github.com/stackrox/stackrox/blob/1bd8c26d4918c3b530ad4fd713244d9cf71e786d/migrator/migrations/m_105_to_m_106_group_id/migration.go#L134
+	groupMigratedIDPrefix = "io.stackrox.authz.group.migrated."
+)
+
+func verifyKeyValuePair(key, value []byte) bool {
+	stringKey := string(key)
+
+	// The key should be a string ID, with a constant prefix and a UUID.
+	if !strings.HasPrefix(stringKey, groupIDPrefix) && !strings.HasPrefix(stringKey, groupMigratedIDPrefix) {
+		return false
+	}
+	stringKey = strings.TrimPrefix(stringKey, groupMigratedIDPrefix)
+	stringKey = strings.TrimPrefix(stringKey, groupIDPrefix)
+	_, err := uuid.FromString(stringKey)
+	if err != nil {
+		return false
+	}
+
+	// The value should be a storage.Group with ID set.
+	var group storage.Group
+	if err := proto.Unmarshal(value, &group); err != nil {
+		return false
+	}
+	return group.GetProps().GetId() != ""
+}
@@ -0,0 +1,167 @@
+package m112tom113
+
+import (
+	"testing"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/stackrox/rox/generated/storage"
+	"github.com/stackrox/rox/migrator/bolthelpers"
+	"github.com/stackrox/rox/pkg/testutils"
+	"github.com/stackrox/rox/pkg/uuid"
+	"github.com/stretchr/testify/suite"
+	bolt "go.etcd.io/bbolt"
+)
+
+func TestMigration(t *testing.T) {
+	suite.Run(t, new(recreateGroupsBucketMigrationTestSuite))
+}
+
+type recreateGroupsBucketMigrationTestSuite struct {
+	suite.Suite
+
+	db *bolt.DB
+}
+
+func (suite *recreateGroupsBucketMigrationTestSuite) SetupTest() {
+	db, err := bolthelpers.NewTemp(testutils.DBFileName(suite))
+	suite.Require().NoError(err, "failed to make BoltDB")
+	suite.db = db
+}
+
+func (suite *recreateGroupsBucketMigrationTestSuite) TearDownTest() {
+	testutils.TearDownDB(suite.db)
+}
+
+func (suite *recreateGroupsBucketMigrationTestSuite) TestMigrate() {
+	// existingGroup should not be dropped during re-creation.
+	existingGroup := &storage.Group{
+		Props: &storage.GroupProperties{
+			Id:             "io.stackrox.authz.group." + uuid.NewV4().String(),
+			AuthProviderId: "some-value",
+		},
+		RoleName: "some-value",
+	}
+	rawExistingGroup, err := proto.Marshal(existingGroup)
+	suite.NoError(err)
+
+	// migratedGroup should not be dropped during re-creation.
+	migratedGroup := &storage.Group{
+		Props: &storage.GroupProperties{
+			Id:             "io.stackrox.authz.group.migrated." + uuid.NewV4().String(),
+			AuthProviderId: "some-value",
+		},
+		RoleName: "some-value",
+	}
+	rawMigratedGroup, err := proto.Marshal(migratedGroup)
+	suite.NoError(err)
+
+	// invalidGroup should be dropped during re-creation.
+	invalidGroup := &storage.Group{
+		Props: &storage.GroupProperties{
+			Key: "some-value",
+		},
+		RoleName: "",
+	}
+	rawInvalidGroup, err := proto.Marshal(invalidGroup)
+	suite.NoError(err)
+
+	// The following cases represent the entries within the groups bucket _before_ migration.
+	// After migration, note that:
+	// - existing-group should not have been dropped, due to having an ID.
+	// - migrated-group should not have been dropped, due to having an ID.
+	// - invalid-group should have been dropped, due to no ID.
+	// - invalid-bytes should have been dropped, due to some weird data and no group proto message.
+	cases := map[string]struct {
+		entry                groupEntry
+		existsAfterMigration bool
+	}{
+		"existing-group": {
+			entry: groupEntry{
+				key:   []byte(existingGroup.GetProps().GetId()),
+				value: rawExistingGroup,
+			},
+			existsAfterMigration: true,
+		},
+		"migrated-group": {
+			entry: groupEntry{
+				key:   []byte(migratedGroup.GetProps().GetId()),
+				value: rawMigratedGroup,
+			},
+			existsAfterMigration: true,
+		},
+		"invalid-group": {
+			entry: groupEntry{
+				key:   []byte("some-random-key"),
+				value: rawInvalidGroup,
+			},
+		},
+		"invalid-group-stored-by-id": {
+			entry: groupEntry{
+				key:   []byte(existingGroup.GetProps().GetId() + "make-it-unique"),
+				value: rawInvalidGroup,
+			},
+		},
+		"invalid-bytes": {
+			entry: groupEntry{
+				key:   []byte("some-random-bytes-no-one-knows"),
+				value: []byte("some-other-random-bytes"),
+			},
+		},
+	}
+
+	var expectedEntriesAfterMigration int
+	for _, c := range cases {
+		if c.existsAfterMigration {
+			expectedEntriesAfterMigration++
+		}
+	}
+
+	// 1. Migration should succeed if the bucket does not exist.
+	suite.NoError(recreateGroupsBucket(suite.db))
+
+	// 2. Add the groups to the groups bucket before running the migration.
+	err = suite.db.Update(func(tx *bolt.Tx) error {
+		bucket, err := tx.CreateBucketIfNotExists(bucketName)
+		suite.NoError(err)
+		for _, c := range cases {
+			suite.NoError(bucket.Put(c.entry.key, c.entry.value))
+		}
+		return nil
+	})
+	suite.NoError(err)
+
+	// 3. Run the migration to re-create the groups bucket and remove invalid entries.
+	suite.NoError(recreateGroupsBucket(suite.db))
+
+	// 4. Verify that all entries are as expected.
+	err = suite.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(bucketName)
+
+		for _, c := range cases {
+			// In case the entry should not exist, it shouldn't be possible to retrieve any values from the given key.
+			if !c.existsAfterMigration {
+				suite.Empty(bucket.Get(c.entry.key))
+			} else {
+				// In case the entry should exist, it should match the expected value.
+				value := bucket.Get(c.entry.key)
+				suite.NotEmpty(value)
+				suite.Equal(c.entry.value, value)
+			}
+		}
+		return nil
+	})
+	suite.NoError(err)
+
+	// 5. Verify that the entries count matches.
+	var actualEntriesCount int
+	err = suite.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(bucketName)
+
+		return bucket.ForEach(func(k, v []byte) error {
+			actualEntriesCount++
+			return nil
+		})
+	})
+	suite.NoError(err)
+	suite.Equal(expectedEntriesAfterMigration, actualEntriesCount)
+}
@@ -14,6 +14,7 @@ import (
 	_ "github.com/stackrox/rox/migrator/migrations/m_109_to_m_110_networkpolicy_guidance_2"
 	_ "github.com/stackrox/rox/migrator/migrations/m_110_to_m_111_replace_deprecated_resources"
 	_ "github.com/stackrox/rox/migrator/migrations/m_111_to_m_112_groups_invalid_values"
+	_ "github.com/stackrox/rox/migrator/migrations/m_112_to_m_113_recreate_groups_bucket"
 	_ "github.com/stackrox/rox/migrator/migrations/m_55_to_m_56_node_scanning_empty"
 	_ "github.com/stackrox/rox/migrator/migrations/m_56_to_m_57_compliance_policy_categories"
 	_ "github.com/stackrox/rox/migrator/migrations/m_57_to_m_58_update_run_secrets_volume_policy_regex"

diff --git a/pkg/migrations/internal/seq_num.go b/pkg/migrations/internal/seq_num.go
@@ -4,7 +4,7 @@ var (
 	// CurrentDBVersionSeqNum is the current DB version number.
 	// This must be incremented every time we write a migration.
 	// It is a shared constant between central and the migrator binary.
-	CurrentDBVersionSeqNum = 112
+	CurrentDBVersionSeqNum = 113
 	// PostgresDBVersionPlus is the current DB version number with Postgres DB data migration.
 	PostgresDBVersionPlus = 56
 )