Fix clearing of ownedTokensIndex and allTokensIndex in ERC721Token.sol

It fixes #839.
OpenZeppelin · Apr 23, 2018 · 53426fd · 53426fd
1 parent 9ee09da
commit 53426fd
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 11 deletions.
diff --git a/contracts/mocks/ERC721TokenMock.sol b/contracts/mocks/ERC721TokenMock.sol
@@ -24,4 +24,12 @@ contract ERC721TokenMock is ERC721Token {
   function setTokenURI(uint256 _tokenId, string _uri) public {
     super._setTokenURI(_tokenId, _uri);
   }
+
+  function ownedTokenIndex(uint256 _tokenId) public view returns(uint256) {
+    return ownedTokensIndex[_tokenId];
+  }
+
+  function allTokenIndex(uint256 _tokenId) public view returns(uint256) {
+    return allTokensIndex[_tokenId];
+  }
 }
diff --git a/contracts/token/ERC721/ERC721Token.sol b/contracts/token/ERC721/ERC721Token.sol
@@ -138,8 +138,11 @@ contract ERC721Token is ERC721, ERC721BasicToken {
     // the lastToken to the first position, and then dropping the element placed in the last position of the list
 
     ownedTokens[_from].length--;
-    ownedTokensIndex[_tokenId] = 0;
-    ownedTokensIndex[lastToken] = tokenIndex;
+    delete ownedTokensIndex[_tokenId];
+
+    if (_tokenId != lastToken) {
+      ownedTokensIndex[lastToken] = tokenIndex;
+    }
   }
 
   /**
@@ -178,8 +181,10 @@ contract ERC721Token is ERC721, ERC721BasicToken {
     allTokens[lastTokenIndex] = 0;
 
     allTokens.length--;
-    allTokensIndex[_tokenId] = 0;
-    allTokensIndex[lastToken] = tokenIndex;
-  }
+    delete allTokensIndex[_tokenId];
 
+    if (_tokenId != lastToken) {
+      allTokensIndex[lastToken] = tokenIndex;
+    }
+  }
 }
diff --git a/test/token/ERC721/ERC721Token.test.js b/test/token/ERC721/ERC721Token.test.js
@@ -51,25 +51,38 @@ contract('ERC721Token', function (accounts) {
     });
 
     describe('burn', function () {
-      const tokenId = firstTokenId;
+      const tokenId = secondTokenId;
       const sender = creator;
 
       beforeEach(async function () {
         await this.token.burn(tokenId, { from: sender });
       });
 
       it('removes that token from the token list of the owner', async function () {
-        const token = await this.token.tokenOfOwnerByIndex(sender, 0);
-        token.toNumber().should.be.equal(secondTokenId);
+        const firstToken = await this.token.tokenOfOwnerByIndex(sender, 0);
+        firstToken.toNumber().should.be.equal(firstTokenId);
+        await assertRevert(this.token.tokenOfOwnerByIndex(sender, 1));
       });
 
       it('adjusts all tokens list', async function () {
-        const token = await this.token.tokenByIndex(0);
-        token.toNumber().should.be.equal(secondTokenId);
+        const firstToken = await this.token.tokenByIndex(0);
+        firstToken.toNumber().should.be.equal(firstTokenId);
+        await assertRevert(this.token.tokenByIndex(1));
+      });
+
+      ['ownedTokenIndex', 'allTokenIndex'].forEach(fnName => {
+        const mappingName = fnName.replace('Token', 'Tokens');
+
+        it(`adjusts ${mappingName}`, async function () {
+          const firstTokenIndex = await this.token[fnName](firstTokenId);
+          const secondTokenIndex = await this.token[fnName](secondTokenId);
+          firstTokenIndex.toNumber().should.be.equal(0);
+          secondTokenIndex.toNumber().should.be.equal(0);
+        });
       });
 
       it('burns all tokens', async function () {
-        await this.token.burn(secondTokenId, { from: sender });
+        await this.token.burn(firstTokenId, { from: sender });
         const total = await this.token.totalSupply();
         total.toNumber().should.be.equal(0);
         await assertRevert(this.token.tokenByIndex(0));