feat: run post migration script in plan to delete invalid keys

Author: ironman0x7b2

x/plan/migrations/expected.go

@@ -7,11 +7,17 @@ import (
 	"github.com/sentinel-official/sentinelhub/v12/x/plan/types/v3"
 )
 
+type LeaseKeeper interface {
+	HasAnyLeaseForNodeByProvider(ctx sdk.Context, nodeAddr types.NodeAddress, provAddr types.ProvAddress) bool
+}
+
 type NodeKeeper interface {
+	DeleteNodeForPlan(ctx sdk.Context, id uint64, addr types.NodeAddress)
 	Store(ctx sdk.Context) sdk.KVStore
 }
 
 type PlanKeeper interface {
+	DeletePlanForNodeByProvider(ctx sdk.Context, nodeAddr types.NodeAddress, provAddr types.ProvAddress, id uint64)
 	GetPlan(ctx sdk.Context, id uint64) (v3.Plan, bool)
 	SetPlan(ctx sdk.Context, plan v3.Plan)
 	SetPlanForNodeByProvider(ctx sdk.Context, nodeAddr types.NodeAddress, provAddr types.ProvAddress, id uint64)

x/plan/migrations/migrator.go

@@ -1,6 +1,7 @@
 package migrations
 
 import (
+	"encoding/binary"
 	"fmt"
 
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -9,22 +10,25 @@ import (
 
 	base "github.com/sentinel-official/sentinelhub/v12/types"
 	"github.com/sentinel-official/sentinelhub/v12/types/v1"
-	"github.com/sentinel-official/sentinelhub/v12/x/node/types"
+	nodetypes "github.com/sentinel-official/sentinelhub/v12/x/node/types"
+	"github.com/sentinel-official/sentinelhub/v12/x/plan/types"
 	"github.com/sentinel-official/sentinelhub/v12/x/plan/types/v2"
 	"github.com/sentinel-official/sentinelhub/v12/x/plan/types/v3"
 )
 
 type Migrator struct {
-	cdc  codec.BinaryCodec
-	node NodeKeeper
-	plan PlanKeeper
+	cdc   codec.BinaryCodec
+	lease LeaseKeeper
+	node  NodeKeeper
+	plan  PlanKeeper
 }
 
-func NewMigrator(cdc codec.BinaryCodec, node NodeKeeper, plan PlanKeeper) Migrator {
+func NewMigrator(cdc codec.BinaryCodec, lease LeaseKeeper, node NodeKeeper, plan PlanKeeper) Migrator {
 	return Migrator{
-		cdc:  cdc,
-		node: node,
-		plan: plan,
+		cdc:   cdc,
+		lease: lease,
+		node:  node,
+		plan:  plan,
 	}
 }
 
@@ -87,7 +91,7 @@ func (k *Migrator) migratePlans(ctx sdk.Context) {
 }
 
 func (k *Migrator) setPlanForNodeByProviderKeys(ctx sdk.Context) {
-	store := prefix.NewStore(k.node.Store(ctx), types.NodeForPlanKeyPrefix)
+	store := prefix.NewStore(k.node.Store(ctx), nodetypes.NodeForPlanKeyPrefix)
 
 	it := store.Iterator(nil, nil)
 	defer it.Close()
@@ -121,3 +125,26 @@ func (k *Migrator) setPlanForProviderKeys(ctx sdk.Context, keys ...[]byte) {
 		k.plan.SetPlanForProvider(ctx, addr, id)
 	}
 }
+
+func (k *Migrator) PostMigrate(ctx sdk.Context) error {
+	store := prefix.NewStore(k.plan.Store(ctx), types.PlanForNodeKeyPrefix)
+
+	it := store.Iterator(nil, nil)
+	defer it.Close()
+
+	for ; it.Valid(); it.Next() {
+		key := it.Key()
+		nodeAddrLen := int(key[0])
+		nodeAddr := base.NodeAddress(key[1 : 1+nodeAddrLen])
+		provAddrLen := int(key[1+nodeAddrLen])
+		provAddr := base.ProvAddress(key[1+nodeAddrLen+1 : 1+nodeAddrLen+1+provAddrLen])
+		id := binary.BigEndian.Uint64(key[1+nodeAddrLen+1+provAddrLen:])
+
+		if !k.lease.HasAnyLeaseForNodeByProvider(ctx, nodeAddr, provAddr) {
+			k.node.DeleteNodeForPlan(ctx, id, nodeAddr)
+			k.plan.DeletePlanForNodeByProvider(ctx, nodeAddr, provAddr, id)
+		}
+	}
+
+	return nil
+}

x/subscription/migrations/migrator.go

@@ -160,6 +160,7 @@ func (k *Migrator) migrateSubscriptions(ctx sdk.Context) {
 					if !item.Status.Equal(v1base.StatusActive) {
 						refund = true
 					}
+
 					if k.lease.HasAnyLeaseForNodeByProvider(ctx, nodeAddr, provAddr) {
 						refund = true
 					}
@@ -168,6 +169,7 @@ func (k *Migrator) migrateSubscriptions(ctx sdk.Context) {
 					if !found {
 						panic(fmt.Errorf("node %s does not exist", nodeAddr))
 					}
+
 					if !node.Status.Equal(v1base.StatusActive) {
 						refund = true
 					}
@@ -176,6 +178,7 @@ func (k *Migrator) migrateSubscriptions(ctx sdk.Context) {
 					if !found {
 						panic(fmt.Errorf("provider %s does not exist", provAddr))
 					}
+
 					if !provider.Status.Equal(v1base.StatusActive) {
 						refund = true
 					}
@@ -253,6 +256,17 @@ func (k *Migrator) migrateSubscriptions(ctx sdk.Context) {
 				k.subscription.SetSubscriptionForPlan(ctx, subscription.PlanID, subscription.ID)
 				k.subscription.SetSubscriptionForInactiveAt(ctx, subscription.InactiveAt, subscription.ID)
 				k.subscription.SetSubscriptionForRenewalAt(ctx, subscription.RenewalAt(), subscription.ID)
+
+				k.subscription.IterateAllocationsForSubscription(ctx, item.ID, func(_ int, item v2.Allocation) bool {
+					accAddr, err := sdk.AccAddressFromBech32(item.Address)
+					if err != nil {
+						panic(err)
+					}
+
+					k.subscription.SetSubscriptionForAccount(ctx, accAddr, subscription.ID)
+
+					return false
+				})
 			}
 		}
 	}

x/vpn/migrations/migrator.go

@@ -30,7 +30,7 @@ func NewMigrator(cdc codec.BinaryCodec, k keeper.Keeper) Migrator {
 		lease:    lease.NewMigrator(cdc, &k.Lease),
 		provider: provider.NewMigrator(cdc, &k.Provider),
 		node:     node.NewMigrator(cdc, &k.Node),
-		plan:     plan.NewMigrator(cdc, &k.Node, &k.Plan),
+		plan:     plan.NewMigrator(cdc, &k.Lease, &k.Node, &k.Plan),
 		subscription: subscription.NewMigrator(
 			cdc, &k.Deposit, &k.Lease, &k.Node, &k.Plan, &k.Provider, &k.Subscription,
 		),
@@ -67,5 +67,9 @@ func (k *Migrator) Migrate(ctx sdk.Context) error {
 		return err
 	}
 
+	if err := k.plan.PostMigrate(ctx); err != nil {
+		return err
+	}
+
 	return nil
 }