fix(contract): compare token to vault token

pooltogether · Jul 15, 2021 · ec57ada · ec57ada
1 parent 85d05ca
commit ec57ada
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/contracts/yield-source/YearnV2YieldSource.sol b/contracts/yield-source/YearnV2YieldSource.sol
@@ -40,6 +40,7 @@ contract YearnV2YieldSource is IYieldSource, ERC20Upgradeable, OwnableUpgradeabl
     /// @notice Emitted when the yield source is initialized
     event YearnV2YieldSourceInitialized(
         IYVaultV2 vault,
+        IERC20Upgradeable token,
         uint8 decimals,
         string symbol,
         string name
@@ -100,6 +101,13 @@ contract YearnV2YieldSource is IYieldSource, ERC20Upgradeable, OwnableUpgradeabl
         vault = _vault;
 
         require(address(_token) != address(0), "YearnV2YieldSource/token-not-zero-address");
+
+        address _vaultToken = _vault.token();
+
+        if (_vaultToken != address(0)) {
+            require(_vaultToken == address(_token), "YearnV2YieldSource/token-address-different");
+        }
+
         token = _token;
 
         __Ownable_init();
@@ -113,6 +121,7 @@ contract YearnV2YieldSource is IYieldSource, ERC20Upgradeable, OwnableUpgradeabl
 
         emit YearnV2YieldSourceInitialized(
             _vault,
+            _token,
             _decimals,
             _symbol,
             _name

diff --git a/test/YearnV2YieldSource.test.ts b/test/YearnV2YieldSource.test.ts
@@ -106,6 +106,7 @@ describe('yearnV2YieldSource', () => {
           .to.emit(yearnV2YieldSource, 'YearnV2YieldSourceInitialized')
           .withArgs(
             vault.address,
+            underlyingToken.address,
             yearnV2YieldSourceTokenDecimals,
             yearnV2YieldSourceTokenSymbol,
             yearnV2YieldSourceTokenName,
@@ -143,6 +144,18 @@ describe('yearnV2YieldSource', () => {
       ).to.be.revertedWith('YearnV2YieldSource/token-not-zero-address');
     });
 
+    it('should fail if token address is different from vault token address', async () => {
+      const randomWallet = ethers.Wallet.createRandom();
+
+      await expect(
+        initializeYearnV2YieldSource(
+          vault.address,
+          randomWallet.address,
+          yearnV2YieldSourceTokenDecimals
+        ),
+      ).to.be.revertedWith('YearnV2YieldSource/token-address-different');
+    });
+
     it('should fail if token decimal is not greater than 0', async () => {
       await vault.mock.apiVersion.returns(compatibleVersions[0]);