chore: cleaning up some functions (#8387)

* chore: cleaning up some functions * attempt more performant --------- Co-authored-by: Adam Tucker <[email protected]>
osmosis-labs · Jun 13, 2024 · 2f19aac · 2f19aac
1 parent f508ff1
commit 2f19aac
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 6 deletions.
diff --git a/x/tokenfactory/keeper/bankactions.go b/x/tokenfactory/keeper/bankactions.go
@@ -44,6 +44,10 @@ func (k Keeper) burnFrom(ctx sdk.Context, amount sdk.Coin, burnFrom string) erro
 		return err
 	}
 
+	if k.IsModuleAcc(ctx, addr) {
+		return types.ErrBurnFromModuleAccount
+	}
+
 	err = k.bankKeeper.SendCoinsFromAccountToModule(ctx,
 		addr,
 		types.ModuleName,
@@ -96,3 +100,8 @@ func (k Keeper) forceTransfer(ctx sdk.Context, amount sdk.Coin, fromAddr string,
 
 	return k.bankKeeper.SendCoins(ctx, fromSdkAddr, toSdkAddr, sdk.NewCoins(amount))
 }
+
+// IsModuleAcc checks if a given address is restricted
+func (k Keeper) IsModuleAcc(ctx sdk.Context, addr sdk.AccAddress) bool {
+	return k.permAddrMap[addr.String()]
+}
diff --git a/x/tokenfactory/keeper/keeper.go b/x/tokenfactory/keeper/keeper.go
@@ -18,8 +18,9 @@ import (
 
 type (
 	Keeper struct {
-		storeKey  storetypes.StoreKey
-		permAddrs map[string]authtypes.PermissionsForAddress
+		storeKey    storetypes.StoreKey
+		permAddrs   map[string]authtypes.PermissionsForAddress
+		permAddrMap map[string]bool
 
 		paramSpace paramtypes.Subspace
 
@@ -45,14 +46,18 @@ func NewKeeper(
 	}
 
 	permAddrs := make(map[string]authtypes.PermissionsForAddress)
+	permAddrMap := make(map[string]bool)
 	for name, perms := range maccPerms {
-		permAddrs[name] = authtypes.NewPermissionsForAddress(name, perms)
+		permsForAddr := authtypes.NewPermissionsForAddress(name, perms)
+		permAddrs[name] = permsForAddr
+		permAddrMap[permsForAddr.GetAddress().String()] = true
 	}
 
 	return Keeper{
-		storeKey:   storeKey,
-		paramSpace: paramSpace,
-		permAddrs:  permAddrs,
+		storeKey:    storeKey,
+		paramSpace:  paramSpace,
+		permAddrs:   permAddrs,
+		permAddrMap: permAddrMap,
 
 		accountKeeper:       accountKeeper,
 		bankKeeper:          bankKeeper,

diff --git a/x/tokenfactory/keeper/msg_server_test.go b/x/tokenfactory/keeper/msg_server_test.go
@@ -104,6 +104,25 @@ func (s *KeeperTestSuite) TestBurnDenomMsg() {
 	}
 }
 
+func (s *KeeperTestSuite) TestBurnFromMsg() {
+	// Create a denom
+	s.CreateDefaultDenom()
+
+	s.Run(fmt.Sprintf("test burn from"), func() {
+		mintAmt := sdk.NewInt64Coin(s.defaultDenom, 10)
+
+		_, err := s.msgServer.Mint(sdk.WrapSDKContext(s.Ctx), types.NewMsgMint(s.TestAccs[0].String(), mintAmt))
+
+		govModAcc := s.App.AccountKeeper.GetModuleAccount(s.Ctx, govtypes.ModuleName)
+
+		err = s.App.BankKeeper.SendCoins(s.Ctx, s.TestAccs[0], govModAcc.GetAddress(), sdk.NewCoins(mintAmt))
+		s.Require().NoError(err)
+
+		_, err = s.msgServer.Burn(s.Ctx, types.NewMsgBurnFrom(s.TestAccs[0].String(), mintAmt, govModAcc.GetAddress().String()))
+		s.Require().ErrorContains(err, "burning from Module Account is not allowed")
+	})
+}
+
 func (s *KeeperTestSuite) TestForceTransferMsg() {
 	// Create a denom
 	s.CreateDefaultDenom()