Add hard coded proxyId into each AssetProxy

6 files changed, 76 insertions, 5 deletions

diff --git a/packages/contracts/src/contracts/current/protocol/AssetProxy/IAssetProxy.sol b/packages/contracts/src/contracts/current/protocol/AssetProxy/IAssetProxy.sol
index 60e74723d..43f45d200 100644
--- a/packages/contracts/src/contracts/current/protocol/AssetProxy/IAssetProxy.sol
+++ b/packages/contracts/src/contracts/current/protocol/AssetProxy/IAssetProxy.sol
@@ -33,4 +33,11 @@ contract IAssetProxy is IAuthorizable {
         address to,
         uint256 amount)
         external;
+    
+    /// @dev Gets the proxy id associated with the proxy address.
+    /// @return Proxy id.
+    function getProxyId()
+        external
+        view
+        returns (uint8);
 }
diff --git a/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC20Proxy.sol b/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC20Proxy.sol
index eef3a39db..c3cfd8c2e 100644
--- a/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC20Proxy.sol
+++ b/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC20Proxy.sol
@@ -29,6 +29,8 @@ contract ERC20Proxy is
     IAssetProxy
 {
 
+    uint8 constant PROXY_ID = 1;
+
     /// @dev Transfers ERC20 tokens. Either succeeds or throws.
     /// @param assetMetadata ERC20-encoded byte array.
     /// @param from Address to transfer token from.
@@ -42,9 +44,25 @@ contract ERC20Proxy is
         external
         onlyAuthorized
     {
+        // Data must be intended for this proxy.
+        require(uint8(assetMetadata[0]) == PROXY_ID);
+
+        // Decode metadata.
         require(assetMetadata.length == 21);
         address token = readAddress(assetMetadata, 1);
+
+        // Transfer tokens.
         bool success = IERC20Token(token).transferFrom(from, to, amount);
         require(success == true);
     }
+
+    /// @dev Gets the proxy id associated with the proxy address.
+    /// @return Proxy id.
+    function getProxyId()
+        external
+        view
+        returns (uint8)
+    {
+        return PROXY_ID;
+    }
 }
diff --git a/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC721Proxy.sol b/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC721Proxy.sol
index eeb7e6710..e11de6744 100644
--- a/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC721Proxy.sol
+++ b/packages/contracts/src/contracts/current/protocol/AssetProxy/proxies/ERC721Proxy.sol
@@ -29,6 +29,8 @@ contract ERC721Proxy is
     IAssetProxy
 {
 
+    uint8 constant PROXY_ID = 2;
+
     /// @dev Transfers ERC721 tokens. Either succeeds or throws.
     /// @param assetMetadata ERC721-encoded byte array
     /// @param from Address to transfer token from.
@@ -42,17 +44,31 @@ contract ERC721Proxy is
         external
         onlyAuthorized
     {
+        // Data must be intended for this proxy.
+        require(uint8(assetMetadata[0]) == PROXY_ID);
+
         // There exists only 1 of each token.
         require(amount == 1);
 
-        // Decode metadata
+        // Decode metadata.
         require(assetMetadata.length == 53);
         address token = readAddress(assetMetadata, 1);
         uint256 tokenId = readUint256(assetMetadata, 21);
 
+        // Transfer token.
         // Either succeeds or throws.
         // @TODO: Call safeTransferFrom if there is additional
         //        data stored in `assetMetadata`.
         ERC721Token(token).transferFrom(from, to, tokenId);
     }
+
+    /// @dev Gets the proxy id associated with the proxy address.
+    /// @return Proxy id.
+    function getProxyId()
+        external
+        view
+        returns (uint8)
+    {
+        return PROXY_ID;
+    }
 }
diff --git a/packages/contracts/src/contracts/current/protocol/Exchange/MixinAssetProxyDispatcher.sol b/packages/contracts/src/contracts/current/protocol/Exchange/MixinAssetProxyDispatcher.sol
index b03bf464a..ed0bacdf8 100644
--- a/packages/contracts/src/contracts/current/protocol/Exchange/MixinAssetProxyDispatcher.sol
+++ b/packages/contracts/src/contracts/current/protocol/Exchange/MixinAssetProxyDispatcher.sol
@@ -43,7 +43,7 @@ contract MixinAssetProxyDispatcher is
     {
         // Do nothing if no amount should be transferred.
         if (amount > 0) {
-            // Lookup asset proxy
+            // Lookup asset proxy.
             require(assetMetadata.length >= 1);
             uint8 assetProxyId = uint8(assetMetadata[0]);
             IAssetProxy assetProxy = assetProxies[assetProxyId];
@@ -66,11 +66,19 @@ contract MixinAssetProxyDispatcher is
         external
         onlyOwner
     {
-        // Ensure the existing asset proxy is not unintentionally overwritten
+        // Ensure the existing asset proxy is not unintentionally overwritten.
         require(oldAssetProxy == address(assetProxies[assetProxyId]));
 
-        // Add asset proxy and log registration
-        assetProxies[assetProxyId] = IAssetProxy(newAssetProxy);
+        IAssetProxy assetProxy = IAssetProxy(newAssetProxy);
+
+        // Ensure that the id of newAssetProxy matches the passed in assetProxyId, unless it is being reset to 0.
+        if (newAssetProxy != address(0)) {
+            uint8 newAssetProxyId = assetProxy.getProxyId();
+            require(newAssetProxyId == assetProxyId);
+        }
+
+        // Add asset proxy and log registration.
+        assetProxies[assetProxyId] = assetProxy;
         emit AssetProxySet(assetProxyId, newAssetProxy, oldAssetProxy);
     }
 
diff --git a/packages/contracts/test/asset_proxy/proxies.ts b/packages/contracts/test/asset_proxy/proxies.ts
index 6aefadee7..bdd0a8696 100644
--- a/packages/contracts/test/asset_proxy/proxies.ts
+++ b/packages/contracts/test/asset_proxy/proxies.ts
@@ -145,6 +145,11 @@ describe('Asset Transfer Proxies', () => {
                 }),
             ).to.be.rejectedWith(constants.REVERT);
         });
+
+        it('should have an id of 1', async () => {
+            const proxyId = await erc20Proxy.getProxyId.callAsync();
+            expect(proxyId).to.equal(1);
+        });
     });
 
     describe('Transfer Proxy - ERC721', () => {
@@ -240,5 +245,10 @@ describe('Asset Transfer Proxies', () => {
                 ),
             ).to.be.rejectedWith(constants.REVERT);
         });
+
+        it('should have an id of 2', async () => {
+            const proxyId = await erc721Proxy.getProxyId.callAsync();
+            expect(proxyId).to.equal(2);
+        });
     });
 });
diff --git a/packages/contracts/test/exchange/dispatcher.ts b/packages/contracts/test/exchange/dispatcher.ts
index 1c346c93f..31f5f5dbb 100644
--- a/packages/contracts/test/exchange/dispatcher.ts
+++ b/packages/contracts/test/exchange/dispatcher.ts
@@ -194,6 +194,18 @@ describe('AssetProxyDispatcher', () => {
                 ),
             ).to.be.rejectedWith(constants.REVERT);
         });
+
+        it('should throw if attempting to register a proxy to the incorrect id', async () => {
+            const prevProxyAddress = ZeroEx.NULL_ADDRESS;
+            return expect(
+                assetProxyDispatcher.registerAssetProxy.sendTransactionAsync(
+                    AssetProxyId.ERC721,
+                    erc20Proxy.address,
+                    prevProxyAddress,
+                    { from: owner },
+                ),
+            ).to.be.rejectedWith(constants.REVERT);
+        });
     });
 
     describe('getAssetProxy', () => {