Disallow calling processWithdrawCommitment twice

contracts/Vault.sol

@@ -4,7 +4,6 @@ pragma experimental ABIEncoderV2;
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { Initializable } from "@openzeppelin/contracts/proxy/Initializable.sol";
-import { Bitmap } from "./libs/Bitmap.sol";
 import { Types } from "./libs/Types.sol";
 import { MerkleTree } from "./libs/MerkleTree.sol";
 import { ImmutableOwnable } from "./libs/ImmutableOwnable.sol";
@@ -16,13 +15,14 @@ contract Vault is Initializable, ImmutableOwnable {
     using Types for Types.MassMigrationCommitment;
     using Types for Types.Batch;
 
+    // withdrawRoot => bool
+    mapping(bytes32 => bool) private approvedWithdrawals;
+
     // Can't be immutable yet. Since the rollup is deployed after Vault
     Rollup public rollup;
     SpokeRegistry public immutable spokes;
     ITokenRegistry public immutable tokenRegistry;
 
-    mapping(uint256 => uint256) private bitmap;
-
     constructor(ITokenRegistry _tokenRegistry, SpokeRegistry _spokes) public {
         tokenRegistry = _tokenRegistry;
         spokes = _spokes;
@@ -37,14 +37,15 @@ contract Vault is Initializable, ImmutableOwnable {
         rollup = _rollup;
     }
 
-    function isBatchApproved(uint256 batchID) public view returns (bool) {
-        return Bitmap.isClaimed(batchID, bitmap);
-    }
-
     function requestApproval(
         uint256 batchID,
         Types.MMCommitmentInclusionProof memory commitmentMP
     ) public {
+        require(
+            !approvedWithdrawals[commitmentMP.commitment.body.withdrawRoot],
+            "Vault: commitment was already approved for withdrawal"
+        );
+
         require(
             msg.sender ==
                 spokes.getSpokeAddress(commitmentMP.commitment.body.spokeID),
@@ -66,10 +67,13 @@ contract Vault is Initializable, ImmutableOwnable {
             ),
             "Vault: Commitment is not present in batch"
         );
+
+        approvedWithdrawals[commitmentMP.commitment.body.withdrawRoot] = true;
+
         (address addr, uint256 l2Unit) =
             tokenRegistry.safeGetRecord(commitmentMP.commitment.body.tokenID);
-        Bitmap.setClaimed(batchID, bitmap);
         uint256 l1Amount = commitmentMP.commitment.body.amount * l2Unit;
+
         require(
             IERC20(addr).approve(msg.sender, l1Amount),
             "Vault: Token approval failed"

contracts/WithdrawManager.sol

@@ -15,7 +15,7 @@ import { Vault } from "./Vault.sol";
 
 contract WithdrawManager {
     using Tx for bytes;
-    using Types for Types.UserState;
+    using Types for Types.MMUserState;
     using SafeERC20 for IERC20;
 
     // withdrawRoot => a bitmap of whether a publicIndex owner has the token claimed

contracts/client/FrontendMassMigration.sol

@@ -120,10 +120,12 @@ contract FrontendMassMigration {
             sender
         );
         if (result != Types.Result.Ok) return (sender.encode(), "", result);
-        newReceiver = Transition.createState(
+        newReceiver = Transition.createMMState(
+            _tx.fromIndex,
             sender.pubkeyID,
             tokenID,
-            _tx.amount
+            _tx.amount,
+            sender.nonce
         );
         return (sender.encode(), newReceiver, result);
     }

contracts/libs/Transition.sol

@@ -13,6 +13,7 @@ import { Tx } from "./Tx.sol";
 library Transition {
     using SafeMath for uint256;
     using Types for Types.UserState;
+    using Types for Types.MMUserState;
 
     function processTransfer(
         bytes32 stateRoot,
@@ -63,7 +64,13 @@ library Transition {
             from
         );
         if (result != Types.Result.Ok) return (newRoot, "", result);
-        freshState = createState(from.state.pubkeyID, tokenID, _tx.amount);
+        freshState = createMMState(
+            _tx.fromIndex,
+            from.state.pubkeyID,
+            tokenID,
+            _tx.amount,
+            from.state.nonce
+        );
 
         return (newRoot, freshState, Types.Result.Ok);
     }
@@ -223,4 +230,22 @@ library Transition {
             });
         return state.encode();
     }
+
+    function createMMState(
+        uint256 stateID,
+        uint256 pubkeyID,
+        uint256 tokenID,
+        uint256 amount,
+        uint256 nonce
+    ) internal pure returns (bytes memory stateEncoded) {
+        Types.MMUserState memory state =
+            Types.MMUserState({
+                stateID: stateID,
+                pubkeyID: pubkeyID,
+                tokenID: tokenID,
+                balance: amount,
+                nonce: nonce
+            });
+        return state.encode();
+    }
 }

contracts/libs/Types.sol

@@ -223,13 +223,36 @@ library Types {
             .decode(encoded, (uint256, uint256, uint256, uint256));
     }
 
+    struct MMUserState {
+        uint256 stateID;
+        uint256 pubkeyID;
+        uint256 tokenID;
+        uint256 balance;
+        uint256 nonce;
+    }
+
+    function encode(MMUserState memory state)
+        internal
+        pure
+        returns (bytes memory)
+    {
+        return
+            abi.encodePacked(
+                state.stateID,
+                state.pubkeyID,
+                state.tokenID,
+                state.balance,
+                state.nonce
+            );
+    }
+
     struct StateMerkleProof {
         UserState state;
         bytes32[] witness;
     }
 
     struct StateMerkleProofWithPath {
-        UserState state;
+        MMUserState state;
         uint256 path;
         bytes32[] witness;
     }

test/integration.test.ts

@@ -331,8 +331,7 @@ describe("Integration Test", function() {
             const tree = migrationTrees[i];
             const { user, index } = group.pickRandom();
             const signature = user.signRaw(withdrawerAddress).sol;
-            // The new stateID in the migration tree is the position user in the group
-            const withdrawProof = tree.getWithdrawProof(index);
+            const withdrawProof = tree.getWithdrawProof(user.stateID, index);
             await withdrawManager
                 .connect(withdrawer)
                 .claimTokens(