From 76ec1fbb7db0a9c6a4f7f522ff50691e854bbbd0 Mon Sep 17 00:00:00 2001
From: Lorenz Herzberger <lorenzherzberger@gmail.com>
Date: Mon, 22 Jul 2024 10:58:40 +0200
Subject: [PATCH] feat: add migration for activated ta counter

Signed-off-by: Lorenz Herzberger <lorenzherzberger@gmail.com>
---
 app/app.go                         |  3 +++
 x/machine/keeper/migrations.go     |  5 +++++
 x/machine/keeper/trust_anchor.go   |  8 ++++----
 x/machine/migrations/v3/migrate.go | 29 +++++++++++++++++++++++++++++
 x/machine/module.go                |  3 +++
 x/machine/types/keys.go            |  2 +-
 6 files changed, 45 insertions(+), 5 deletions(-)
 create mode 100644 x/machine/migrations/v3/migrate.go

diff --git a/app/app.go b/app/app.go
index b26ecbe..943df6a 100644
--- a/app/app.go
+++ b/app/app.go
@@ -1002,4 +1002,7 @@ func (app *App) setupUpgradeHandlers() {
 	app.UpgradeKeeper.SetUpgradeHandler("v0.10.5", func(ctx sdk.Context, _ upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 		return app.mm.RunMigrations(ctx, app.configurator, fromVM)
 	})
+	app.UpgradeKeeper.SetUpgradeHandler("v0.10.6", func(ctx sdk.Context, _ upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		return app.mm.RunMigrations(ctx, app.configurator, fromVM)
+	})
 }
diff --git a/x/machine/keeper/migrations.go b/x/machine/keeper/migrations.go
index decb5a1..668d453 100644
--- a/x/machine/keeper/migrations.go
+++ b/x/machine/keeper/migrations.go
@@ -3,6 +3,7 @@ package keeper
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	v2 "github.com/planetmint/planetmint-go/x/machine/migrations/v2"
+	v3 "github.com/planetmint/planetmint-go/x/machine/migrations/v3"
 )
 
 // Migrator is a struct for handling in-place store migrations.
@@ -19,3 +20,7 @@ func NewMigrator(keeper Keeper) Migrator {
 func (m Migrator) Migrate1to2(ctx sdk.Context) error {
 	return v2.MigrateStore(ctx, m.keeper.storeKey, m.keeper.cdc)
 }
+
+func (m Migrator) Migrate2to3(ctx sdk.Context) error {
+	return v3.MigrateStore(ctx, m.keeper.taStoreKey, m.keeper.cdc)
+}
diff --git a/x/machine/keeper/trust_anchor.go b/x/machine/keeper/trust_anchor.go
index 07fed74..b40f6a1 100644
--- a/x/machine/keeper/trust_anchor.go
+++ b/x/machine/keeper/trust_anchor.go
@@ -51,15 +51,15 @@ func (k Keeper) GetTrustAnchor(ctx sdk.Context, pubKey string) (val types.TrustA
 }
 
 func (k Keeper) setActivatedTACount(ctx sdk.Context, counter uint64) {
-	taCounterStore := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefix(types.TAIndexKey))
+	taCounterStore := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefix(types.ActivatedTACounterPrefix))
 	bz := make([]byte, 8)
 	binary.BigEndian.PutUint64(bz, counter)
-	taCounterStore.Set([]byte(types.ActivatedTACounterKey), bz)
+	taCounterStore.Set([]byte{1}, bz)
 }
 
 func (k Keeper) GetActivatedTACount(ctx sdk.Context) (counter uint64) {
-	taCounterStore := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefix(types.TAIndexKey))
-	bz := taCounterStore.Get([]byte(types.ActivatedTACounterKey))
+	taCounterStore := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefix(types.ActivatedTACounterPrefix))
+	bz := taCounterStore.Get([]byte{1})
 	if bz == nil {
 		return 0
 	}
diff --git a/x/machine/migrations/v3/migrate.go b/x/machine/migrations/v3/migrate.go
new file mode 100644
index 0000000..fbccedb
--- /dev/null
+++ b/x/machine/migrations/v3/migrate.go
@@ -0,0 +1,29 @@
+package v3
+
+import (
+	"encoding/binary"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/planetmint/planetmint-go/x/machine/types"
+)
+
+func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.BinaryCodec) error {
+	store := prefix.NewStore(ctx.KVStore(storeKey), types.KeyPrefix(types.ActivatedTACounterPrefix))
+
+	count := uint64(0)
+
+	iterator := store.Iterator(nil, nil)
+	defer iterator.Close()
+	for ; iterator.Valid(); iterator.Next() {
+		count++
+	}
+
+	bz := make([]byte, 8)
+	binary.BigEndian.PutUint64(bz, count)
+	store.Set([]byte{1}, bz)
+
+	return nil
+}
diff --git a/x/machine/module.go b/x/machine/module.go
index c028031..db9a783 100644
--- a/x/machine/module.go
+++ b/x/machine/module.go
@@ -121,6 +121,9 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 	if err := cfg.RegisterMigration(types.ModuleName, 2, m.Migrate1to2); err != nil {
 		panic(fmt.Errorf("failed to register migration of %s to v2: %w", types.ModuleName, err))
 	}
+	if err := cfg.RegisterMigration(types.ModuleName, 3, m.Migrate2to3); err != nil {
+		panic(fmt.Errorf("failed to register migration of %s to v3: %w", types.ModuleName, err))
+	}
 }
 
 // RegisterInvariants registers the invariants of the module. If an invariant deviates from its predicted value, the InvariantRegistry triggers appropriate logic (most often the chain will be halted)
diff --git a/x/machine/types/keys.go b/x/machine/types/keys.go
index 30fd9c6..0fb0932 100644
--- a/x/machine/types/keys.go
+++ b/x/machine/types/keys.go
@@ -29,7 +29,7 @@ const (
 
 	ParamsKey = "Machine/Params"
 
-	ActivatedTACounterKey = "ActivatedTACounter"
+	ActivatedTACounterPrefix = "ActivatedTACounter"
 )
 
 func KeyPrefix(p string) []byte {