EOLink 0.2.5: security test with pytest parametrize

If you like this content consider tipping in ETH, LINK or an ERC20 stablecoin at this address: 0xf9f6849230cBe10200B43dB103511b778898e71C

Code associated with this post is available on Github

One of the main security concerns with a smart contract is the ability for malicious actors to execute functions in order to reroute funds or otherwise manipulate the contract activity. This risk is mitigated by assigning specific access modifiers to functions that restrict access to certain groups or individuals as represented by their wallet addresses, or under certain conditions. The project testing has so far assumed that functions have been accessed by the correct wallets only, meaning the contract’s response to requests from other wallets has not been explicitly checked. This raises two interesting challenges for the testing scripts: 1) some tests will now be designed to fail; 2) some tests need to be run multiple times with systematically changing test conditions.

These two new requirements can be met with two nested pytest wrappers – mark.xfail and parametrize. The former wrapper instructs pytest that the expected and correct behaviour for a specific test is to fail. This is the case with some of the security tests we want to carry out – using an unauthorized wallet is supposed to fail as it violates the rules of the access modifier. Without wrapping the test in mark.xfail this would cause pytest to return a test failure because the function has not executed, but it is actually a success from a security auditing perspective. After wrapping the test in mark.xfail, the test report returns a yellow “X” to signify that a test failed as expected.

Screenshot showing the output from a testing script that includes tests designed to pass (green dots) and tests expected to fail (yellow X)

Parametrize is used to iterate through tests. This is achieved by passing an iterator variable name and the various values to iterate through as arguments in the pytest wrapper @pytest.mark.parameterize. The variable name is then used inside the function to represent the variable that changes in each iteration.

Thankfully, instead of wrapping the test function twice, mark.xfail can be passed as an argument to a specific variable value in the iterator wrapper, meaning that pytest can be instructed to expect passes for some values and fails for others. In our example, this means we can expect the authorized wallet to pass and the others to fail.

There is one other nuance to testing in this way. A failed test may change the state of the blockchain in some way that alters the initial conditions for the next test iteration, for example if the test fails after making a transaction, the balances of the various wallets will change between each iteration, sometimes causing unintended fails in tests that should pass and vice-versa, or just creating uncontrolled test conditions that could be unrepresentative of the contract in the wild. This can only really be mitigated by careful programming, likely including transaction reverts or explicit rebalancing of the state of the blockchain in each iteration to ensure the initial conditions are identical in each iteration.

An example of this is included below for the “test_withdrawal_from_contract” function. This is the key function in the contract that redistributes funds to external wallets, so it is the function with the most obvious security risk. It is only supposed to be accessible to the contract owner. In the code below pytest.mark.parametrize is used to iterate through three wallets, and pytest.mark.xfail is used to indicate that 2/3 of the tests should fail due to the requests coming from invalid wallet addresses.

@pytest.mark.parametrize("wallet",
    [
     pytest.param('donor', marks=pytest.mark.xfail(reason="wallet does not have access")),
     pytest.param('customer', marks=pytest.mark.xfail(reason="wallet does not have access)),
     'owner'
    ])
def test_withdrawal_from_contract(set_deposit_amount, getDeployedContract, wallet, load_owner, load_customer, load_donor, set_threshold):
    
    """
    ensure withdrawal of DAI from contract to recipient executs correctly.
    If oracle data > threshold, send to customer
    if oracle data < threshold, send to donor
    parametrized test: only "owner" has permission, so 2/3 expected to fail
    """    
    
    if wallet == 'donor':
        wallet = load_donor
    elif wallet == 'customer':
        wallet = load_customer
    elif wallet == 'owner':
        wallet = load_owner
    else:
        raise("INVALID WALLET")


    dai = interface.IERC20('0xFf795577d9AC8bD7D90Ee22b6C1703490b6512FD')
    adai = interface.IERC20('0xdCf0aF9e59C002FA3AA091a46196b37530FD48a8')
    contract = getDeployedContract
    trigger = contract.viewValueFromOracle()

    assert adai.balanceOf(contract) == 0
    assert dai.balanceOf(contract) > set_deposit_amount

    initialOwnerBalance = dai.balanceOf(load_owner)
    initialCustomerBalance = dai.balanceOf(load_customer)
    initialDonorBalance = dai.balanceOf(load_donor)

    contract.retrieveDAI({'from':wallet})
    
    time.sleep(5)

    if trigger > set_threshold:

        assert adai.balanceOf(contract) == 0
        assert dai.balanceOf(contract) == 0
        assert dai.balanceOf(load_owner) > initialOwnerBalance
        assert dai.balanceOf(load_customer) == initialCustomerBalance+set_deposit_amount
        assert dai.balanceOf(load_donor) == initialDonorBalance
    
    else:
        assert adai.balanceOf(contract) == 0
        assert dai.balanceOf(contract) == 0
        assert dai.balanceOf(load_owner) > initialOwnerBalance
        assert dai.balanceOf(load_customer) == initialCustomerBalance
        assert dai.balanceOf(load_donor) == initialDonorBalance+set_deposit_amount

    return

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s