diff --git a/keeper/keeper.go b/keeper/keeper.go index e67f50d..9640d90 100644 --- a/keeper/keeper.go +++ b/keeper/keeper.go @@ -46,8 +46,8 @@ type Keeper struct { EntitlementsOwner collections.Item[string] Paused collections.Item[bool] PublicCapabilities collections.Map[string, bool] - RoleCapabilities collections.Map[collections.Pair[string, uint64], bool] - UserRoles collections.Map[collections.Pair[[]byte, uint64], bool] + RoleCapabilities collections.Map[[]byte, bool] + UserRoles collections.Map[[]byte, bool] addressCodec address.Codec accountKeeper types.AccountKeeper @@ -89,8 +89,8 @@ func NewKeeper( EntitlementsOwner: collections.NewItem(builder, entitlements.OwnerKey, "entitlements_owner", collections.StringValue), 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.PairKeyCodec(collections.StringKey, collections.Uint64Key), collections.BoolValue), - UserRoles: collections.NewMap(builder, entitlements.UserPrefix, "entitlements_user_roles", collections.PairKeyCodec(collections.BytesKey, collections.Uint64Key), 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.BytesKey, collections.BoolValue), accountKeeper: accountKeeper, bankKeeper: bankKeeper, diff --git a/keeper/msg_server_entitlements_test.go b/keeper/msg_server_entitlements_test.go index 7b2c74e..c90b230 100644 --- 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. @@ -464,7 +464,7 @@ func TestEntitlementsUserCapability(t *testing.T) { tmpRole := k.RoleCapabilities k.RoleCapabilities = collections.NewMap( collections.NewSchemaBuilder(mocks.FailingStore(mocks.Set, utils.GetKVStore(ctx, types.ModuleName))), - entitlements.CapabilityPrefix, "entitlements_role_capabilities", collections.PairKeyCodec(collections.StringKey, collections.Uint64Key), collections.BoolValue, + entitlements.CapabilityPrefix, "entitlements_role_capabilities", collections.BytesKey, collections.BoolValue, ) // ACT: Attempt set role capability with failing RoleCapabilities collection store. diff --git a/keeper/state_entitlements.go b/keeper/state_entitlements.go index e5b1644..f6b1cf0 100644 --- a/keeper/state_entitlements.go +++ b/keeper/state_entitlements.go @@ -8,6 +8,7 @@ package keeper import ( "context" + "encoding/binary" "slices" "cosmossdk.io/collections" @@ -80,13 +81,19 @@ func (k *Keeper) SetPublicCapability(ctx context.Context, method string, enabled func (k *Keeper) GetCapabilityRoles(ctx context.Context, method string) []entitlements.Role { var roles []entitlements.Role - _ = k.RoleCapabilities.Walk(ctx, collections.NewPrefixedPairRange[string, uint64](method), func(key collections.Pair[string, uint64], enabled bool) (stop bool, err error) { + itr, _ := k.RoleCapabilities.Iterate(ctx, new(collections.Range[[]byte]).Prefix([]byte(method))) + + 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 } @@ -94,10 +101,10 @@ func (k *Keeper) GetCapabilityRoles(ctx context.Context, method string) []entitl func (k *Keeper) GetAllCapabilityRoles(ctx context.Context) []entitlements.RoleCapability { var capabilityRoles []entitlements.RoleCapability - _ = k.RoleCapabilities.Walk(ctx, nil, func(key collections.Pair[string, uint64], enabled bool) (stop bool, err error) { + _ = k.RoleCapabilities.Walk(ctx, nil, func(key []byte, enabled bool) (stop bool, err error) { capabilityRoles = append(capabilityRoles, entitlements.RoleCapability{ - Method: key.K1(), - Role: entitlements.Role(key.K2()), + Method: string(key[:len(key)-8]), + Role: entitlements.Role(binary.BigEndian.Uint64(key[len(key)-8:])), Enabled: enabled, }) @@ -108,21 +115,27 @@ func (k *Keeper) GetAllCapabilityRoles(ctx context.Context) []entitlements.RoleC } func (k *Keeper) SetRoleCapability(ctx context.Context, method string, role entitlements.Role, enabled bool) error { - return k.RoleCapabilities.Set(ctx, collections.Join(method, uint64(role)), enabled) + return k.RoleCapabilities.Set(ctx, entitlements.CapabilityRoleKey(method, role), enabled) } // -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 } @@ -130,11 +143,11 @@ func (k *Keeper) GetUserRoles(ctx context.Context, user []byte) []entitlements.R 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, }) @@ -145,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 index da4c019..b9c3d70 100644 --- 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 index 9a44116..abce0e8 100644 --- a/types/entitlements/keys.go +++ b/types/entitlements/keys.go @@ -8,8 +8,6 @@ package entitlements import "encoding/binary" -const SubmoduleName = "halo-entitlements" - var ( OwnerKey = []byte("entitlements/owner") PausedKey = []byte("entitlements/paused") @@ -18,26 +16,14 @@ var ( UserPrefix = []byte("entitlements/user/") ) -func PublicKey(method string) []byte { - return append(PublicPrefix, []byte(method)...) -} - -func CapabilityKey(method string) []byte { - return append(CapabilityPrefix, []byte(method)...) -} - func CapabilityRoleKey(method string, role Role) []byte { bz := make([]byte, 8) binary.BigEndian.PutUint64(bz, uint64(role)) - return append(CapabilityKey(method), bz...) -} - -func UserKey(address []byte) []byte { - return append(UserPrefix, address...) + return append([]byte(method), bz...) } func UserRoleKey(address []byte, role Role) []byte { bz := make([]byte, 8) binary.BigEndian.PutUint64(bz, uint64(role)) - return append(UserKey(address), bz...) + return append(address, bz...) } diff --git a/types/keys.go b/types/keys.go index def872f..3589d4c 100644 --- 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...) -}