From ffd8a3c8710d06bb473c3c7a26fcb1a67203e27f Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Fri, 8 Jul 2022 13:13:21 +0200
Subject: [PATCH] Prevent migration to a restricted code

---
 x/wasm/keeper/contract_keeper.go |  2 +-
 x/wasm/keeper/keeper.go          |  4 ++++
 x/wasm/keeper/keeper_test.go     | 13 +++++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/x/wasm/keeper/contract_keeper.go b/x/wasm/keeper/contract_keeper.go
index 9dc0a4be9d..a6a0adf0e1 100644
--- a/x/wasm/keeper/contract_keeper.go
+++ b/x/wasm/keeper/contract_keeper.go
@@ -75,7 +75,7 @@ func (p PermissionedKeeper) UnpinCode(ctx sdk.Context, codeID uint64) error {
 	return p.nested.unpinCode(ctx, codeID)
 }
 
-// SetExtraContractAttributes updates the extra attributes that can be stored with the contract info
+// SetContractInfoExtension updates the extra attributes that can be stored with the contract info
 func (p PermissionedKeeper) SetContractInfoExtension(ctx sdk.Context, contract sdk.AccAddress, extra types.ContractInfoExtension) error {
 	return p.nested.setContractInfoExtension(ctx, contract, extra)
 }
diff --git a/x/wasm/keeper/keeper.go b/x/wasm/keeper/keeper.go
index 3fdcbec483..d261a8b9ce 100644
--- a/x/wasm/keeper/keeper.go
+++ b/x/wasm/keeper/keeper.go
@@ -395,6 +395,10 @@ func (k Keeper) migrate(ctx sdk.Context, contractAddress sdk.AccAddress, caller
 		return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, "unknown code")
 	}
 
+	if !authZ.CanInstantiateContract(newCodeInfo.InstantiateConfig, caller) {
+		return nil, sdkerrors.Wrap(sdkerrors.ErrUnauthorized, "to use new code")
+	}
+
 	// check for IBC flag
 	switch report, err := k.wasmVM.AnalyzeCode(newCodeInfo.CodeHash); {
 	case err != nil:
diff --git a/x/wasm/keeper/keeper_test.go b/x/wasm/keeper/keeper_test.go
index bb20a39267..44527ac97c 100644
--- a/x/wasm/keeper/keeper_test.go
+++ b/x/wasm/keeper/keeper_test.go
@@ -875,6 +875,10 @@ func TestMigrate(t *testing.T) {
 	ibcCodeID := StoreIBCReflectContract(t, ctx, keepers).CodeID
 	require.NotEqual(t, originalCodeID, newCodeID)
 
+	restrictedCodeID := StoreHackatomExampleContract(t, ctx, keepers).CodeID
+	keeper.SetAccessConfig(ctx, restrictedCodeID, types.AllowNobody)
+	require.NotEqual(t, originalCodeID, restrictedCodeID)
+
 	anyAddr := RandomAccountAddress(t)
 	newVerifierAddr := RandomAccountAddress(t)
 	initMsgBz := HackatomExampleInitMsg{
@@ -952,6 +956,15 @@ func TestMigrate(t *testing.T) {
 			toCodeID:   originalCodeID,
 			expErr:     sdkerrors.ErrUnauthorized,
 		},
+		"prevent migration when new code is restricted": {
+			admin:      creator,
+			caller:     creator,
+			initMsg:    initMsgBz,
+			fromCodeID: originalCodeID,
+			toCodeID:   restrictedCodeID,
+			migrateMsg: migMsgBz,
+			expErr:     sdkerrors.ErrUnauthorized,
+		},
 		"fail with non existing code id": {
 			admin:      creator,
 			caller:     creator,