Ensure transferring the fee to the caller is done as the last step

pool-registry/src/main/java/org/aion/unity/PoolRegistry.java

@@ -83,7 +83,7 @@ public static void registerPool(Address signingAddress, int commissionRate, byte
 
         // step 2: create a staker in the staker registry
         /*
-        registerStaker(Address identityAddress, Address managementAddress, Address signingAddress, Address coinbaseAddress)
+        registerStaker(Address identityAddress, Address signingAddress, Address coinbaseAddress)
          */
         String methodName = "registerStaker";
         // encoded data is directly written to the byte array to reduce energy usage
@@ -362,9 +362,11 @@ public static void finalizeUndelegate(long id) {
         // At this point the StakerRegistry has transferred the fee amount to this contract in a re-entrant call.
         // This is a safe assumption since these two contracts are tightly coupled. Thus, the fee amount is not explicitly stored to save energy.
         assert reentrantValueTransferAmount != null;
-        // transfer the fee to the caller
-        secureCall(Blockchain.getCaller(), reentrantValueTransferAmount, new byte[0], Blockchain.getRemainingEnergy());
+        BigInteger fee = reentrantValueTransferAmount;
         reentrantValueTransferAmount = null;
+
+        // transfer the fee to the caller
+        secureCall(Blockchain.getCaller(), fee, new byte[0], Blockchain.getRemainingEnergy());
     }
 
     /**
@@ -391,10 +393,10 @@ public static void finalizeTransfer(long id) {
         // At this point the StakerRegistry has transferred the fee amount to this contract in a re-entrant call.
         // This is a safe assumption since these two contracts are tightly coupled. Thus, the fee amount is not explicitly stored to save energy.
         assert reentrantValueTransferAmount != null;
-        // transfer fee
-        secureCall(Blockchain.getCaller(), reentrantValueTransferAmount, new byte[0], Blockchain.getRemainingEnergy());
-        BigInteger remainingTransferValue = transfer.amount.subtract(reentrantValueTransferAmount);
+        BigInteger fee = reentrantValueTransferAmount;
         reentrantValueTransferAmount = null;
+        assert fee.compareTo(transfer.amount) <= 0;
+        BigInteger remainingTransferValue = transfer.amount.subtract(fee);
 
         // remove transfer
         PoolRegistryStorage.putPendingTransfer(id, null);
@@ -410,6 +412,9 @@ public static void finalizeTransfer(long id) {
         detectBlockRewards(stateMachine);
 
         delegate(transfer.initiator, transfer.toPool, remainingTransferValue, false, stateMachine, delegatorInfo);
+
+        // transfer fee
+        secureCall(Blockchain.getCaller(), fee, new byte[0], Blockchain.getRemainingEnergy());
     }
 
     /**
@@ -434,11 +439,6 @@ public static BigInteger withdrawRewards(Address pool) {
             amount = amount.add(stateMachine.onWithdrawOperator());
         }
 
-        // do a transfer if amount > 0
-        if (amount.signum() == 1) {
-            secureCall(caller, amount, new byte[0], Blockchain.getRemainingEnergy());
-        }
-
         // remove the delegator from storage if the stake is zero and all the rewards have been withdrawn
         if(delegatorInfo.stake.equals(BigInteger.ZERO)) {
             delegatorInfo = null;
@@ -448,6 +448,12 @@ public static BigInteger withdrawRewards(Address pool) {
         PoolRegistryStorage.putPoolRewards(pool, stateMachine.currentPoolRewards);
 
         PoolRegistryEvents.withdrew(caller, pool, amount);
+
+        // do a transfer if amount > 0
+        if (amount.signum() == 1) {
+            secureCall(caller, amount, new byte[0], Blockchain.getRemainingEnergy());
+        }
+
         return amount;
     }
 
@@ -521,10 +527,10 @@ public static void autoDelegateRewards(Address pool, Address delegator) {
 
             Blockchain.println("Auto delegation: fee = " + fee + ", remaining = " + remaining);
 
+            delegate(delegator, pool, remaining, true, stateMachine, delegatorInfo);
+
             // transfer fee to the caller
             secureCall(Blockchain.getCaller(), fee, new byte[0], Blockchain.getRemainingEnergy());
-
-            delegate(delegator, pool, remaining, true, stateMachine, delegatorInfo);
         }
     }
 

pool-registry/src/test/java/org/aion/unity/ReentrantTest.java

@@ -0,0 +1,226 @@
+package org.aion.unity;
+
+import avm.Address;
+import org.aion.avm.core.util.Helpers;
+import org.aion.avm.core.util.LogSizeUtils;
+import org.aion.avm.embed.AvmRule;
+import org.aion.avm.tooling.ABIUtil;
+import org.aion.avm.userlib.abi.ABIStreamingEncoder;
+import org.aion.kernel.TestingState;
+import org.aion.types.AionAddress;
+import org.aion.types.Log;
+import org.aion.unity.resources.ReentrantContract;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.spongycastle.util.encoders.Hex;
+
+import java.lang.reflect.Field;
+import java.math.BigInteger;
+import java.util.Scanner;
+
+/**
+ * This test ensures that reentrant calls to the contract, do not violate contract invariance:
+ * - Additional stake being generated
+ * - Additional rewards being withdrawn
+ */
+public class ReentrantTest {
+    private static BigInteger ENOUGH_BALANCE_TO_TRANSACT = BigInteger.TEN.pow(18 + 5);
+    private static BigInteger MIN_SELF_STAKE = new BigInteger("1000000000000000000000");
+    private static long COMMISSION_RATE_CHANGE_TIME_LOCK_PERIOD = 6 * 60 * 24 * 7;
+
+    @Rule
+    public AvmRule RULE = new AvmRule(false);
+
+    // default address with balance
+    private Address preminedAddress;
+
+    // contract address
+    private Address stakerRegistry;
+    private Address poolRegistry;
+
+    @Before
+    public void setup() {
+        preminedAddress = RULE.getRandomAddress(ENOUGH_BALANCE_TO_TRANSACT);
+
+        try (Scanner s = new Scanner(PoolRegistryTest.class.getResourceAsStream("StakerRegistry.txt"))) {
+            String contract = s.nextLine();
+            AvmRule.ResultWrapper result = RULE.deploy(preminedAddress, BigInteger.ZERO, Hex.decode(contract));
+            Assert.assertTrue(result.getReceiptStatus().isSuccess());
+            stakerRegistry = result.getDappAddress();
+        }
+
+        Address placeHolder = new Address(Helpers.hexStringToBytes("0000000000000000000000000000000000000000000000000000000000000000"));
+        byte[] coinbaseArguments = ABIUtil.encodeDeploymentArguments(placeHolder);
+        byte[] coinbaseBytes = RULE.getDappBytes(PoolCoinbase.class, coinbaseArguments, 1);
+
+        byte[] arguments = ABIUtil.encodeDeploymentArguments(stakerRegistry, MIN_SELF_STAKE, BigInteger.ONE, COMMISSION_RATE_CHANGE_TIME_LOCK_PERIOD, coinbaseBytes);
+        byte[] data = RULE.getDappBytes(PoolRegistry.class, arguments, 1, PoolStorageObjects.class, PoolRewardsStateMachine.class, PoolRegistryEvents.class, PoolRegistryStorage.class);
+
+        AvmRule.ResultWrapper result = RULE.deploy(preminedAddress, BigInteger.ZERO, data);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+
+        Assert.assertEquals(1, result.getLogs().size());
+        Log poolRegistryEvent = result.getLogs().get(0);
+        Assert.assertArrayEquals(LogSizeUtils.truncatePadTopic("ADSDeployed".getBytes()), poolRegistryEvent.copyOfTopics().get(0));
+        Assert.assertArrayEquals(stakerRegistry.toByteArray(), poolRegistryEvent.copyOfTopics().get(1));
+        Assert.assertEquals(MIN_SELF_STAKE, new BigInteger(poolRegistryEvent.copyOfTopics().get(2)));
+        Assert.assertEquals(BigInteger.ONE, new BigInteger(poolRegistryEvent.copyOfTopics().get(3)));
+        Assert.assertEquals(COMMISSION_RATE_CHANGE_TIME_LOCK_PERIOD, new BigInteger(poolRegistryEvent.copyOfData()).longValue());
+
+        poolRegistry = result.getDappAddress();
+    }
+
+    public Address setupNewPool(int fee) {
+        fee = fee * 10000;
+        Address newPool = RULE.getRandomAddress(ENOUGH_BALANCE_TO_TRANSACT);
+
+        // register a new pool
+        byte[] txData = new ABIStreamingEncoder()
+                .encodeOneString("registerPool")
+                .encodeOneAddress(newPool)
+                .encodeOneInteger(fee)
+                .encodeOneByteArray("https://".getBytes())
+                .encodeOneByteArray(new byte[32])
+                .toBytes();
+
+        AvmRule.ResultWrapper result = RULE.call(newPool, poolRegistry, nStake(1), txData, 2_000_000L, 1L);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+
+        // verify now
+        txData = new ABIStreamingEncoder()
+                .encodeOneString("isActive")
+                .encodeOneAddress(newPool)
+                .toBytes();
+        result = RULE.call(newPool, stakerRegistry, BigInteger.ZERO, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+        Assert.assertEquals(true, result.getDecodedReturnData());
+
+        return newPool;
+    }
+
+    @Test
+    public void reentrantAutoDelegateRewards() {
+        Address pool = setupNewPool(0);
+        BigInteger stake = nStake(1);
+
+        // deploy contract
+        byte[] arguments = ABIUtil.encodeDeploymentArguments(poolRegistry, pool);
+        byte[] data = RULE.getDappBytes(ReentrantContract.class, arguments, 1);
+
+        AvmRule.ResultWrapper result = RULE.deploy(preminedAddress, BigInteger.ZERO, data);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+        Address reentractContract = result.getDappAddress();
+
+        // delegate using the contract
+        byte[] txData = new ABIStreamingEncoder()
+                .encodeOneString("delegate")
+                .toBytes();
+        result = RULE.call(preminedAddress, reentractContract, stake, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+
+        txData = new ABIStreamingEncoder()
+                .encodeOneString("enableAutoRewardsDelegation")
+                .toBytes();
+        result = RULE.call(preminedAddress, reentractContract, BigInteger.ZERO, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+
+        generateBlock(pool, 1000000000);
+
+        txData = new ABIStreamingEncoder()
+                .encodeOneString("autoDelegateRewards")
+                .toBytes();
+        result = RULE.call(preminedAddress, reentractContract, BigInteger.ZERO, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+
+        txData = new ABIStreamingEncoder()
+                .encodeOneString("getStake")
+                .encodeOneAddress(pool)
+                .encodeOneAddress(reentractContract)
+                .toBytes();
+        result = RULE.call(preminedAddress, poolRegistry, BigInteger.ZERO, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+        BigInteger currentStake = (BigInteger) result.getDecodedReturnData();
+        BigInteger rewards = BigInteger.valueOf(1000000000 / 2);
+        BigInteger fee = rewards.divide(BigInteger.valueOf(1000000));
+        // stake = value of stake + half of the rewards - fee
+        // if the re-entrant call could cheat the system, this value would be greater
+        Assert.assertEquals(stake.add(rewards).subtract(fee), currentStake);
+
+        Assert.assertEquals(fee, RULE.kernel.getBalance(new AionAddress(reentractContract.toByteArray())));
+    }
+
+    @Test
+    public void reentrantGetRewards() {
+        Address pool = setupNewPool(0);
+        BigInteger stake = nStake(1);
+
+        // deploy contract
+        byte[] arguments = ABIUtil.encodeDeploymentArguments(poolRegistry, pool);
+        byte[] data = RULE.getDappBytes(ReentrantContract.class, arguments, 1);
+
+        AvmRule.ResultWrapper result = RULE.deploy(preminedAddress, BigInteger.ZERO, data);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+        Address reentractContract = result.getDappAddress();
+
+        // delegate using the contract
+        byte[] txData = new ABIStreamingEncoder()
+                .encodeOneString("delegate")
+                .toBytes();
+        result = RULE.call(preminedAddress, reentractContract, stake, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+
+        generateBlock(pool, 1000000000);
+
+        txData = new ABIStreamingEncoder()
+                .encodeOneString("withdrawRewards")
+                .toBytes();
+        result = RULE.call(preminedAddress, reentractContract, BigInteger.ZERO, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+
+        // total rewards/2
+        // if the re-entrant call could cheat the system, this value would be greater
+        Assert.assertEquals(1000000000 / 2, RULE.kernel.getBalance(new AionAddress(reentractContract.toByteArray())).intValue());
+    }
+
+    private BigInteger nStake(int n) {
+        return MIN_SELF_STAKE.multiply(BigInteger.valueOf(n));
+    }
+
+    private void generateBlock(Address pool, long blockRewards) {
+        AionAddress coinbaseAddress = new AionAddress(getCoinbaseAddress(pool).toByteArray());
+        RULE.kernel.adjustBalance(coinbaseAddress, BigInteger.valueOf(blockRewards));
+        incrementBlockNumber();
+    }
+
+    private long getBlockNumber() {
+        return RULE.kernel.getBlockNumber();
+    }
+
+    private void incrementBlockNumber() {
+        tweakBlockNumber(getBlockNumber() + 1);
+    }
+
+    private void tweakBlockNumber(long number) {
+        try {
+            Field f = TestingState.class.getDeclaredField("blockNumber");
+            f.setAccessible(true);
+
+            f.set(RULE.kernel, number);
+
+        } catch (NoSuchFieldException | IllegalAccessException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private Address getCoinbaseAddress(Address pool) {
+        byte[] txData = new ABIStreamingEncoder()
+                .encodeOneString("getCoinbaseAddress")
+                .encodeOneAddress(pool)
+                .toBytes();
+        AvmRule.ResultWrapper result = RULE.call(preminedAddress, stakerRegistry, BigInteger.ZERO, txData);
+        Assert.assertTrue(result.getReceiptStatus().isSuccess());
+        return (Address) result.getDecodedReturnData();
+    }
+}

pool-registry/src/test/java/org/aion/unity/resources/ReentrantContract.java

@@ -0,0 +1,79 @@
+package org.aion.unity.resources;
+
+import avm.Address;
+import avm.Blockchain;
+import avm.Result;
+import org.aion.avm.tooling.abi.Callable;
+import org.aion.avm.tooling.abi.Fallback;
+import org.aion.avm.tooling.abi.Initializable;
+import org.aion.avm.userlib.abi.ABIStreamingEncoder;
+
+import java.math.BigInteger;
+
+public class ReentrantContract {
+
+    @Initializable
+    public static Address poolRegistry;
+    @Initializable
+    public static Address pool;
+
+    private static boolean hasBeenCalled = false;
+    private static byte[] txData;
+
+    @Callable
+    public static void delegate() {
+        byte[] txData = new ABIStreamingEncoder()
+                .encodeOneString("delegate")
+                .encodeOneAddress(pool)
+                .toBytes();
+
+        Result result = Blockchain.call(poolRegistry, Blockchain.getValue(), txData, Blockchain.getRemainingEnergy());
+        Blockchain.require(result.isSuccess());
+    }
+
+    @Callable
+    public static void autoDelegateRewards() {
+        byte[] data = new ABIStreamingEncoder()
+                .encodeOneString("autoDelegateRewards")
+                .encodeOneAddress(pool)
+                .encodeOneAddress(Blockchain.getAddress())
+                .toBytes();
+        txData = data;
+        makeCall();
+    }
+
+    @Callable
+    public static void enableAutoRewardsDelegation() {
+        byte[] txData = new ABIStreamingEncoder()
+                .encodeOneString("enableAutoRewardsDelegation")
+                .encodeOneAddress(pool)
+                .encodeOneInteger(1)
+                .toBytes();
+        Result result = Blockchain.call(poolRegistry, BigInteger.ZERO, txData, Blockchain.getRemainingEnergy());
+        Blockchain.require(result.isSuccess());
+    }
+
+    @Callable
+    public static void withdrawRewards(){
+        byte[] data = new ABIStreamingEncoder()
+                .encodeOneString("withdrawRewards")
+                .encodeOneAddress(pool)
+                .toBytes();
+        txData = data;
+        makeCall();
+    }
+
+    @Fallback
+    public static void fallback() {
+        // call back to the pool registry, if this is the first call from PoolRegistry to this contract
+        if (!hasBeenCalled) {
+            hasBeenCalled = true;
+            makeCall();
+        }
+    }
+
+    private static void makeCall() {
+        Result result = Blockchain.call(poolRegistry, BigInteger.ZERO, txData, Blockchain.getRemainingEnergy());
+        Blockchain.require(result.isSuccess());
+    }
+}