Added validation to yield source address

pooltogether · Apr 9, 2021 · ea6a68f · ea6a68f
1 parent 73e17c1
commit ea6a68f
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 3 deletions.
diff --git a/contracts/prize-pool/yield-source/YieldSourcePrizePool.sol b/contracts/prize-pool/yield-source/YieldSourcePrizePool.sol
@@ -28,6 +28,7 @@ contract YieldSourcePrizePool is PrizePool {
     public
     initializer
   {
+    require(address(_yieldSource) != address(0), "YieldSourcePrizePool/yield-source-zero");
     PrizePool.initialize(
       _reserveRegistry,
       _controlledTokens,
@@ -36,7 +37,11 @@ contract YieldSourcePrizePool is PrizePool {
     );
     yieldSource = _yieldSource;
 
-    emit YieldSourcePrizePoolInitialized(address(yieldSource));
+    // A hack to determine whether it's an actual yield source
+    (bool succeeded,) = address(_yieldSource).staticcall(abi.encode(_yieldSource.depositToken.selector));
+    require(succeeded, "YieldSourcePrizePool/invalid-yield-source");
+
+    emit YieldSourcePrizePoolInitialized(address(_yieldSource));
   }
 
   /// @notice Determines whether the passed token can be transferred out as an external award.

diff --git a/test/YieldSourcePrizePool.test.js b/test/YieldSourcePrizePool.test.js
@@ -12,7 +12,7 @@ let overrides = { gasLimit: 9500000 }
 describe('YieldSourcePrizePool', function() {
   let wallet, wallet2
 
-  let prizePool, erc20token, prizeStrategy, reserveRegistry, yieldSource
+  let prizePool, erc20token, prizeStrategy, reserveRegistry, yieldSource, YieldSourcePrizePoolHarness
 
   let poolMaxExitFee = toWei('0.5')
   let poolMaxTimelockDuration = 10000
@@ -43,7 +43,7 @@ describe('YieldSourcePrizePool', function() {
     reserveRegistry = await deployMockContract(wallet, RegistryInterface.abi, overrides)
 
     debug('deploying YieldSourcePrizePoolHarness...')
-    const YieldSourcePrizePoolHarness =  await hre.ethers.getContractFactory("YieldSourcePrizePoolHarness", wallet, overrides)
+    YieldSourcePrizePoolHarness =  await hre.ethers.getContractFactory("YieldSourcePrizePoolHarness", wallet, overrides)
     prizePool = await YieldSourcePrizePoolHarness.deploy()
 
     const ControlledToken = await hre.artifacts.readArtifact("ControlledToken")
@@ -73,6 +73,31 @@ describe('YieldSourcePrizePool', function() {
 
       expect(await prizePool.yieldSource()).to.equal(yieldSource.address)
     })
+
+    it('should require the yield source', async () => {
+      prizePool = await YieldSourcePrizePoolHarness.deploy()
+
+      await expect(prizePool.initializeYieldSourcePrizePool(
+        reserveRegistry.address,
+        [ticket.address],
+        poolMaxExitFee,
+        poolMaxTimelockDuration,
+        ethers.constants.AddressZero
+      )).to.be.revertedWith("YieldSourcePrizePool/yield-source-zero")
+    })
+
+    it('should require a valid yield source', async () => {
+      prizePool = await YieldSourcePrizePoolHarness.deploy()
+      await ticket.mock.controller.returns(prizePool.address)
+
+      await expect(prizePool.initializeYieldSourcePrizePool(
+        reserveRegistry.address,
+        [ticket.address],
+        poolMaxExitFee,
+        poolMaxTimelockDuration,
+        prizePool.address
+      )).to.be.revertedWith("YieldSourcePrizePool/invalid-yield-source")
+    })
   })
 
   describe('supply()', async () => {