fix: ensure state compatiblity for user roles

noble-assets · Oct 29, 2024 · 0719754 · 0719754
1 parent 741981e
commit 0719754
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 28 deletions.
diff --git a/keeper/keeper.go b/keeper/keeper.go
@@ -47,7 +47,7 @@ type Keeper struct {
 	Paused             collections.Item[bool]
 	PublicCapabilities collections.Map[string, bool]
 	RoleCapabilities   collections.Map[[]byte, bool]
-	UserRoles          collections.Map[collections.Pair[[]byte, uint64], bool]
+	UserRoles          collections.Map[[]byte, bool]
 
 	addressCodec      address.Codec
 	accountKeeper     types.AccountKeeper
@@ -90,7 +90,7 @@ func NewKeeper(
 		Paused:             collections.NewItem(builder, entitlements.PausedKey, "entitlements_paused", collections.BoolValue),
 		PublicCapabilities: collections.NewMap(builder, entitlements.PublicPrefix, "entitlements_public_capabilities", collections.StringKey, collections.BoolValue),
 		RoleCapabilities:   collections.NewMap(builder, entitlements.CapabilityPrefix, "entitlements_role_capabilities", collections.BytesKey, collections.BoolValue),
-		UserRoles:          collections.NewMap(builder, entitlements.UserPrefix, "entitlements_user_roles", collections.PairKeyCodec(collections.BytesKey, collections.Uint64Key), collections.BoolValue),
+		UserRoles:          collections.NewMap(builder, entitlements.UserPrefix, "entitlements_user_roles", collections.BytesKey, collections.BoolValue),
 
 		accountKeeper:     accountKeeper,
 		bankKeeper:        bankKeeper,

diff --git a/keeper/msg_server_entitlements_test.go b/keeper/msg_server_entitlements_test.go
@@ -319,7 +319,7 @@ func TestEntitlementsUserRoles(t *testing.T) {
 	tmp := k.UserRoles
 	k.UserRoles = collections.NewMap(
 		collections.NewSchemaBuilder(mocks.FailingStore(mocks.Set, utils.GetKVStore(ctx, types.ModuleName))),
-		entitlements.UserPrefix, "entitlements_user_roles", collections.PairKeyCodec(collections.BytesKey, collections.Uint64Key), collections.BoolValue,
+		entitlements.UserPrefix, "entitlements_user_roles", collections.BytesKey, collections.BoolValue,
 	)
 
 	// ACT: Attempt set user role with failing UserRoles collection store.

diff --git a/keeper/state_entitlements.go b/keeper/state_entitlements.go
@@ -120,28 +120,34 @@ func (k *Keeper) SetRoleCapability(ctx context.Context, method string, role enti
 
 //
 
-func (k *Keeper) GetUserRoles(ctx context.Context, user []byte) []entitlements.Role {
+func (k *Keeper) GetUserRoles(ctx context.Context, address []byte) []entitlements.Role {
 	var roles []entitlements.Role
 
-	_ = k.UserRoles.Walk(ctx, collections.NewPrefixedPairRange[[]byte, uint64](user), func(key collections.Pair[[]byte, uint64], enabled bool) (stop bool, err error) {
+	itr, _ := k.UserRoles.Iterate(ctx, new(collections.Range[[]byte]).Prefix(address))
+
+	defer itr.Close()
+
+	for ; itr.Valid(); itr.Next() {
+		key, _ := itr.Key()
+		enabled, _ := itr.Value()
+
 		if enabled {
-			roles = append(roles, entitlements.Role(key.K2()))
+			role := binary.BigEndian.Uint64(key[len(key)-8:])
+			roles = append(roles, entitlements.Role(role))
 		}
-
-		return false, nil
-	})
+	}
 
 	return roles
 }
 
 func (k *Keeper) GetAllUserRoles(ctx context.Context) []entitlements.UserRole {
 	var userRoles []entitlements.UserRole
 
-	_ = k.UserRoles.Walk(ctx, nil, func(key collections.Pair[[]byte, uint64], enabled bool) (stop bool, err error) {
-		address, _ := k.addressCodec.BytesToString(key.K1())
+	_ = k.UserRoles.Walk(ctx, nil, func(key []byte, enabled bool) (stop bool, err error) {
+		address, _ := k.addressCodec.BytesToString(key[:len(key)-8])
 		userRoles = append(userRoles, entitlements.UserRole{
 			User:    address,
-			Role:    entitlements.Role(key.K2()),
+			Role:    entitlements.Role(binary.BigEndian.Uint64(key[len(key)-8:])),
 			Enabled: enabled,
 		})
 
@@ -152,10 +158,10 @@ func (k *Keeper) GetAllUserRoles(ctx context.Context) []entitlements.UserRole {
 }
 
 func (k *Keeper) HasRole(ctx context.Context, address []byte, role entitlements.Role) bool {
-	enabled, _ := k.UserRoles.Get(ctx, collections.Join(address, uint64(role)))
+	enabled, _ := k.UserRoles.Get(ctx, entitlements.UserRoleKey(address, role))
 	return enabled
 }
 
 func (k *Keeper) SetUserRole(ctx context.Context, address []byte, role entitlements.Role, enabled bool) error {
-	return k.UserRoles.Set(ctx, collections.Join(address, uint64(role)), enabled)
+	return k.UserRoles.Set(ctx, entitlements.UserRoleKey(address, role), enabled)
 }
diff --git a/types/aggregator/keys.go b/types/aggregator/keys.go
@@ -6,19 +6,9 @@
 
 package aggregator
 
-import "encoding/binary"
-
-const SubmoduleName = "halo-aggregator"
-
 var (
 	OwnerKey       = []byte("aggregator/owner")
 	LastRoundIDKey = []byte("aggregator/last_round_id")
 	NextPriceKey   = []byte("aggregator/next_price")
 	RoundPrefix    = []byte("aggregator/round/")
 )
-
-func RoundKey(id uint64) []byte {
-	bz := make([]byte, 8)
-	binary.BigEndian.PutUint64(bz, id)
-	return append(RoundPrefix, bz...)
-}
diff --git a/types/entitlements/keys.go b/types/entitlements/keys.go
@@ -21,3 +21,9 @@ func CapabilityRoleKey(method string, role Role) []byte {
 	binary.BigEndian.PutUint64(bz, uint64(role))
 	return append([]byte(method), bz...)
 }
+
+func UserRoleKey(address []byte, role Role) []byte {
+	bz := make([]byte, 8)
+	binary.BigEndian.PutUint64(bz, uint64(role))
+	return append(address, bz...)
+}
diff --git a/types/keys.go b/types/keys.go
@@ -16,7 +16,3 @@ var (
 	OwnerKey    = []byte("owner")
 	NoncePrefix = []byte("nonce/")
 )
-
-func NonceKey(address []byte) []byte {
-	return append(NoncePrefix, address...)
-}