Migrated sanction module to collections for state storage.

.changelog/unreleased/features/2498-sanction-module-collections.md

@@ -0,0 +1 @@
+* Migrated sanction module to collections for state storage [#2498](https://github.com/provenance-io/provenance/issues/2498).

app/app.go

@@ -728,7 +728,7 @@ func New(
 		unsanctionableAddrs = append(unsanctionableAddrs, authtypes.NewModuleAddress(mName))
 	}
 	unsanctionableAddrs = append(unsanctionableAddrs, authtypes.NewModuleAddress(quarantine.ModuleName))
-	app.SanctionKeeper = sanctionkeeper.NewKeeper(appCodec, keys[sanction.StoreKey],
+	app.SanctionKeeper = sanctionkeeper.NewKeeper(appCodec, runtime.NewKVStoreService(keys[sanction.StoreKey]),
 		app.BankKeeper, &app.GovKeeper,
 		govAuthority, unsanctionableAddrs)
 

x/sanction/codec.go

@@ -1,6 +1,12 @@
 package sanction
 
 import (
+	"encoding/binary"
+	"encoding/json"
+	fmt "fmt"
+
+	"cosmossdk.io/collections"
+
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/msgservice"
@@ -15,3 +21,256 @@ func RegisterInterfaces(registry codectypes.InterfaceRegistry) {
 
 	msgservice.RegisterMsgServiceDesc(registry, &_Msg_serviceDesc)
 }
+
+// AccAddressKey implements KeyCodec[sdk.AccAddress] for length-prefixed addresses
+type AccAddressKey struct{}
+
+func (a AccAddressKey) Encode(buffer []byte, addr sdk.AccAddress) (int, error) {
+	size := a.Size(addr)
+	if len(buffer) < size {
+		return 0, fmt.Errorf("buffer too small: need %d, got %d", size, len(buffer))
+	}
+
+	buffer[0] = byte(len(addr))
+	copy(buffer[1:], addr)
+	return size, nil
+}
+
+func (a AccAddressKey) Decode(buffer []byte) (int, sdk.AccAddress, error) {
+	if len(buffer) == 0 {
+		return 0, nil, fmt.Errorf("empty buffer")
+	}
+
+	addrLen := int(buffer[0])
+	size := 1 + addrLen
+
+	if len(buffer) < size {
+		return 0, nil, fmt.Errorf("insufficient bytes: need %d, got %d", size, len(buffer))
+	}
+
+	addr := sdk.AccAddress(buffer[1 : 1+addrLen])
+	return size, addr, nil
+}
+
+func (a AccAddressKey) Size(addr sdk.AccAddress) int {
+	return 1 + len(addr)
+}
+
+func (a AccAddressKey) EncodeJSON(addr sdk.AccAddress) ([]byte, error) {
+	return json.Marshal(addr.String())
+}
+
+func (a AccAddressKey) DecodeJSON(b []byte) (sdk.AccAddress, error) {
+	var s string
+	if err := json.Unmarshal(b, &s); err != nil {
+		return nil, err
+	}
+	return sdk.AccAddressFromBech32(s)
+}
+
+func (a AccAddressKey) Stringify(addr sdk.AccAddress) string {
+	return addr.String()
+}
+
+func (a AccAddressKey) KeyType() string {
+	return "sdk.AccAddress"
+}
+
+func (a AccAddressKey) EncodeNonTerminal(buffer []byte, addr sdk.AccAddress) (int, error) {
+	return a.Encode(buffer, addr)
+}
+
+func (a AccAddressKey) DecodeNonTerminal(buffer []byte) (int, sdk.AccAddress, error) {
+	return a.Decode(buffer)
+}
+
+func (a AccAddressKey) SizeNonTerminal(addr sdk.AccAddress) int {
+	return a.Size(addr)
+}
+
+// TemporaryKeyCodec implements KeyCodec[collections.Pair[sdk.AccAddress, uint64]]
+// Format: <addr len (1 byte)><addr><gov prop id (8 bytes)>
+type TemporaryKeyCodec struct{}
+
+func (t TemporaryKeyCodec) Encode(buffer []byte, key collections.Pair[sdk.AccAddress, uint64]) (int, error) {
+	addr := key.K1()
+	propID := key.K2()
+
+	size := t.Size(key)
+	if len(buffer) < size {
+		return 0, fmt.Errorf("buffer too small: need %d, got %d", size, len(buffer))
+	}
+
+	buffer[0] = byte(len(addr))
+	copy(buffer[1:], addr)
+	binary.BigEndian.PutUint64(buffer[1+len(addr):], propID)
+
+	return size, nil
+}
+
+func (t TemporaryKeyCodec) Decode(buffer []byte) (int, collections.Pair[sdk.AccAddress, uint64], error) {
+	if len(buffer) < 1+8 {
+		return 0, collections.Pair[sdk.AccAddress, uint64]{}, fmt.Errorf("buffer too short: got %d bytes", len(buffer))
+	}
+
+	addrLen := int(buffer[0])
+	size := 1 + addrLen + 8
+
+	if len(buffer) < size {
+		return 0, collections.Pair[sdk.AccAddress, uint64]{}, fmt.Errorf("invalid address length: need %d, got %d", size, len(buffer))
+	}
+
+	addr := sdk.AccAddress(buffer[1 : 1+addrLen])
+	propID := binary.BigEndian.Uint64(buffer[1+addrLen:])
+
+	return size, collections.Join(addr, propID), nil
+}
+
+func (t TemporaryKeyCodec) Size(key collections.Pair[sdk.AccAddress, uint64]) int {
+	return 1 + len(key.K1()) + 8
+}
+
+func (t TemporaryKeyCodec) EncodeJSON(key collections.Pair[sdk.AccAddress, uint64]) ([]byte, error) {
+	type jsonKey struct {
+		Address string `json:"address"`
+		PropID  uint64 `json:"prop_id"`
+	}
+	return json.Marshal(jsonKey{
+		Address: key.K1().String(),
+		PropID:  key.K2(),
+	})
+}
+
+func (t TemporaryKeyCodec) DecodeJSON(b []byte) (collections.Pair[sdk.AccAddress, uint64], error) {
+	var data struct {
+		Address string `json:"address"`
+		PropID  uint64 `json:"prop_id"`
+	}
+	if err := json.Unmarshal(b, &data); err != nil {
+		return collections.Pair[sdk.AccAddress, uint64]{}, err
+	}
+
+	addr, err := sdk.AccAddressFromBech32(data.Address)
+	if err != nil {
+		return collections.Pair[sdk.AccAddress, uint64]{}, err
+	}
+
+	return collections.Join(addr, data.PropID), nil
+}
+
+func (t TemporaryKeyCodec) Stringify(key collections.Pair[sdk.AccAddress, uint64]) string {
+	return fmt.Sprintf("%s/%d", key.K1().String(), key.K2())
+}
+
+func (t TemporaryKeyCodec) KeyType() string {
+	return "sanction.TemporaryKey"
+}
+
+func (t TemporaryKeyCodec) EncodeNonTerminal(buffer []byte, key collections.Pair[sdk.AccAddress, uint64]) (int, error) {
+	return t.Encode(buffer, key)
+}
+
+func (t TemporaryKeyCodec) DecodeNonTerminal(buffer []byte) (int, collections.Pair[sdk.AccAddress, uint64], error) {
+	return t.Decode(buffer)
+}
+
+func (t TemporaryKeyCodec) SizeNonTerminal(key collections.Pair[sdk.AccAddress, uint64]) int {
+	return t.Size(key)
+}
+
+func (t TemporaryKeyCodec) EncodePrefix(addr sdk.AccAddress) []byte {
+	prefix := make([]byte, 1+len(addr))
+	prefix[0] = byte(len(addr))
+	copy(prefix[1:], addr)
+	return prefix
+}
+
+// ProposalIndexKeyCodec implements KeyCodec[collections.Pair[uint64, sdk.AccAddress]]
+// Format: <proposal id (8 bytes)><addr len (1 byte)><addr>
+type ProposalIndexKeyCodec struct{}
+
+func (p ProposalIndexKeyCodec) Encode(buffer []byte, key collections.Pair[uint64, sdk.AccAddress]) (int, error) {
+	propID := key.K1()
+	addr := key.K2()
+
+	size := p.Size(key)
+	if len(buffer) < size {
+		return 0, fmt.Errorf("buffer too small: need %d, got %d", size, len(buffer))
+	}
+
+	binary.BigEndian.PutUint64(buffer[0:8], propID)
+	buffer[8] = byte(len(addr))
+	copy(buffer[9:], addr)
+
+	return size, nil
+}
+
+func (p ProposalIndexKeyCodec) Decode(buffer []byte) (int, collections.Pair[uint64, sdk.AccAddress], error) {
+	if len(buffer) < 8+1 {
+		return 0, collections.Pair[uint64, sdk.AccAddress]{}, fmt.Errorf("buffer too short: got %d bytes", len(buffer))
+	}
+
+	propID := binary.BigEndian.Uint64(buffer[0:8])
+	addrLen := int(buffer[8])
+	size := 8 + 1 + addrLen
+
+	if len(buffer) < size {
+		return 0, collections.Pair[uint64, sdk.AccAddress]{}, fmt.Errorf("invalid address length: need %d, got %d", size, len(buffer))
+	}
+
+	addr := sdk.AccAddress(buffer[9 : 9+addrLen])
+
+	return size, collections.Join(propID, addr), nil
+}
+
+func (p ProposalIndexKeyCodec) Size(key collections.Pair[uint64, sdk.AccAddress]) int {
+	return 8 + 1 + len(key.K2())
+}
+
+func (p ProposalIndexKeyCodec) EncodeJSON(key collections.Pair[uint64, sdk.AccAddress]) ([]byte, error) {
+	type jsonKey struct {
+		PropID  uint64 `json:"prop_id"`
+		Address string `json:"address"`
+	}
+	return json.Marshal(jsonKey{
+		PropID:  key.K1(),
+		Address: key.K2().String(),
+	})
+}
+
+func (p ProposalIndexKeyCodec) DecodeJSON(b []byte) (collections.Pair[uint64, sdk.AccAddress], error) {
+	var data struct {
+		PropID  uint64 `json:"prop_id"`
+		Address string `json:"address"`
+	}
+	if err := json.Unmarshal(b, &data); err != nil {
+		return collections.Pair[uint64, sdk.AccAddress]{}, err
+	}
+
+	addr, err := sdk.AccAddressFromBech32(data.Address)
+	if err != nil {
+		return collections.Pair[uint64, sdk.AccAddress]{}, err
+	}
+
+	return collections.Join(data.PropID, addr), nil
+}
+
+func (p ProposalIndexKeyCodec) Stringify(key collections.Pair[uint64, sdk.AccAddress]) string {
+	return fmt.Sprintf("%d/%s", key.K1(), key.K2().String())
+}
+
+func (p ProposalIndexKeyCodec) KeyType() string {
+	return "sanction.ProposalIndexKey"
+}
+
+func (p ProposalIndexKeyCodec) EncodeNonTerminal(buffer []byte, key collections.Pair[uint64, sdk.AccAddress]) (int, error) {
+	return p.Encode(buffer, key)
+}
+
+func (p ProposalIndexKeyCodec) DecodeNonTerminal(buffer []byte) (int, collections.Pair[uint64, sdk.AccAddress], error) {
+	return p.Decode(buffer)
+}
+
+func (p ProposalIndexKeyCodec) SizeNonTerminal(key collections.Pair[uint64, sdk.AccAddress]) int {
+	return p.Size(key)
+}

x/sanction/keeper/export_test.go

@@ -3,8 +3,7 @@ package keeper
 import (
 	"context"
 
-	storetypes "cosmossdk.io/store/types"
-
+	"cosmossdk.io/core/store"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
@@ -45,9 +44,9 @@ func (k Keeper) WithUnsanctionableAddrs(unsanctionableAddrs map[string]bool) Kee
 	return k
 }
 
-// StoreKey, for unit tests, exposes this keeper's storekey.
-func (k Keeper) StoreKey() storetypes.StoreKey {
-	return k.storeKey
+// KVStore, for unit tests, returns a KVStore for the given context from the keeper's StoreService.
+func (k Keeper) KVStore(ctx sdk.Context) store.KVStore {
+	return k.StoreService.OpenKVStore(ctx)
 }
 
 // MsgSanctionTypeURL, for unit tests, exposes this keeper's msgSanctionTypeURL.
@@ -71,23 +70,23 @@ func (k Keeper) GetParamAsCoinsOrDefault(ctx sdk.Context, name string, dflt sdk.
 }
 
 // GetLatestTempEntry, for unit tests, exposes this keeper's getLatestTempEntry function.
-func (k Keeper) GetLatestTempEntry(store storetypes.KVStore, addr sdk.AccAddress) []byte {
-	return k.getLatestTempEntry(store, addr)
+func (k Keeper) GetLatestTempEntry(ctx sdk.Context, addr sdk.AccAddress) []byte {
+	return k.getLatestTempEntry(ctx, addr)
 }
 
 // GetParam, for unit tests, exposes this keeper's getParam function.
-func (k Keeper) GetParam(store storetypes.KVStore, name string) (string, bool) {
-	return k.getParam(store, name)
+func (k Keeper) GetParam(ctx sdk.Context, name string) (string, bool) {
+	return k.getParam(ctx, name)
 }
 
 // SetParam, for unit tests, exposes this keeper's setParam function.
-func (k Keeper) SetParam(store storetypes.KVStore, name, value string) {
-	k.setParam(store, name, value)
+func (k Keeper) SetParam(ctx sdk.Context, name, value string) {
+	k.setParam(ctx, name, value)
 }
 
 // DeleteParam, for unit tests, exposes this keeper's deleteParam function.
-func (k Keeper) DeleteParam(store storetypes.KVStore, name string) {
-	k.deleteParam(store, name)
+func (k Keeper) DeleteParam(ctx sdk.Context, name string) {
+	k.deleteParam(ctx, name)
 }
 
 // ProposalGovHook, for unit tests, exposes this keeper's proposalGovHook function.

x/sanction/keeper/genesis.go

@@ -65,18 +65,20 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *sanction.GenesisState {
 // This is designed for use with ExportGenesis. See also IterateSanctionedAddresses.
 func (k Keeper) GetAllSanctionedAddresses(ctx sdk.Context) []string {
 	var rv []string
-	k.IterateSanctionedAddresses(ctx, func(addr sdk.AccAddress) bool {
+	if err := k.IterateSanctionedAddresses(ctx, func(addr sdk.AccAddress) bool {
 		rv = append(rv, addr.String())
 		return false
-	})
+	}); err != nil {
+		panic(fmt.Errorf("failed to iterate sanctioned addresses: %w", err))
+	}
 	return rv
 }
 
 // GetAllTemporaryEntries gets all the Temporary entries.
 // This is designed for use with ExportGenesis. See also IterateTemporaryEntries.
 func (k Keeper) GetAllTemporaryEntries(ctx sdk.Context) []*sanction.TemporaryEntry {
 	var rv []*sanction.TemporaryEntry
-	k.IterateTemporaryEntries(ctx, nil, func(addr sdk.AccAddress, id uint64, isSanction bool) bool {
+	if err := k.IterateTemporaryEntries(ctx, nil, func(addr sdk.AccAddress, id uint64, isSanction bool) bool {
 		status := sanction.TEMP_STATUS_SANCTIONED
 		if !isSanction {
 			status = sanction.TEMP_STATUS_UNSANCTIONED
@@ -87,6 +89,8 @@ func (k Keeper) GetAllTemporaryEntries(ctx sdk.Context) []*sanction.TemporaryEnt
 			Status:     status,
 		})
 		return false
-	})
+	}); err != nil {
+		panic(fmt.Errorf("failed to iterate temporary entries: %w", err))
+	}
 	return rv
 }

x/sanction/keeper/genesis_test.go

@@ -532,7 +532,7 @@ func (s *GenesisTestSuite) TestKeeper_ExportGenesis() {
 
 		store := s.GetStore()
 		s.Require().NotPanics(func() {
-			s.Keeper.DeleteParam(store, keeper.ParamNameImmediateUnsanctionMinDeposit)
+			s.Keeper.DeleteParam(s.SdkCtx, keeper.ParamNameImmediateUnsanctionMinDeposit)
 		}, "deleting the ParamNameImmediateUnsanctionMinDeposit param entry")
 		// Note: Not using UnsanctionAddresses here since I want to keep the temp entries for these addrs.
 		s.Require().NotPanics(func() {

x/sanction/keeper/gov_hooks.go

@@ -97,7 +97,9 @@ func (k Keeper) proposalGovHook(goCtx context.Context, proposalID uint64) error
 		// C) The extra processing from calling DeleteGovPropTempEntries is probably on par with what's needed
 		//    to not always call it.
 		// D) There's no risk of this deleting anything that shouldn't be deleted.
-		k.DeleteGovPropTempEntries(ctx, proposalID)
+		if err := k.DeleteGovPropTempEntries(ctx, proposalID); err != nil {
+			panic(fmt.Errorf("failed to delete temporary entries for proposal %d: %w", proposalID, err))
+		}
 	case govv1.StatusPassed:
 		// Nothing to do. The processing of the proposal message does everything that's needed.
 	default: