From d38d4cdd4797bd15f402cb17b02ef02beecf2557 Mon Sep 17 00:00:00 2001
From: Paul <108695806+pxrl@users.noreply.github.com>
Date: Wed, 4 Jan 2023 21:30:54 +0100
Subject: [PATCH] fix: [L-02] Lack of input validation (#47)

The OZ December 2022 audit identified a lack of input validation on
the following functions:
 - stakeFor() (missing validation of beneficiary address).
 - unstake()  (missing validation on the amount input).
 - _stake()   (missing validation on the amount input).

This change applies the missing validation in accordance with the
recommendations outlined in the finding.
---
 contracts/AcceleratingDistributor.sol |  4 ++
 test/AcceleratingDistributor.Admin.ts | 67 +++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/contracts/AcceleratingDistributor.sol b/contracts/AcceleratingDistributor.sol
index 4067583..a497a95 100644
--- a/contracts/AcceleratingDistributor.sol
+++ b/contracts/AcceleratingDistributor.sol
@@ -199,6 +199,7 @@ contract AcceleratingDistributor is ReentrancyGuard, Ownable, Multicall {
         uint256 amount,
         address beneficiary
     ) external nonReentrant onlyEnabled(stakedToken) {
+        require(beneficiary != address(0), "Invalid beneficiary");
         _stake(stakedToken, amount, beneficiary);
     }
 
@@ -208,6 +209,8 @@ contract AcceleratingDistributor is ReentrancyGuard, Ownable, Multicall {
      * @param amount The amount of the token to withdraw.
      */
     function unstake(address stakedToken, uint256 amount) public nonReentrant onlyInitialized(stakedToken) {
+        require(amount > 0, "Invalid amount");
+
         _updateReward(stakedToken, msg.sender);
         UserDeposit storage userDeposit = stakingTokens[stakedToken].stakingBalances[msg.sender];
 
@@ -400,6 +403,7 @@ contract AcceleratingDistributor is ReentrancyGuard, Ownable, Multicall {
         uint256 amount,
         address staker
     ) internal {
+        require(amount > 0, "Invalid amount");
         _updateReward(stakedToken, staker);
 
         UserDeposit storage userDeposit = stakingTokens[stakedToken].stakingBalances[staker];
diff --git a/test/AcceleratingDistributor.Admin.ts b/test/AcceleratingDistributor.Admin.ts
index ee196e2..8114fa2 100644
--- a/test/AcceleratingDistributor.Admin.ts
+++ b/test/AcceleratingDistributor.Admin.ts
@@ -5,6 +5,8 @@ import { baseEmissionRate, maxMultiplier, secondsToMaxMultiplier } from "./const
 let timer: Contract, acrossToken: Contract, distributor: Contract, lpToken1: Contract;
 let owner: SignerWithAddress, rando: SignerWithAddress;
 
+const zeroAddress = ethers.constants.AddressZero;
+
 describe("AcceleratingDistributor: Admin Functions", async function () {
   beforeEach(async function () {
     [owner, rando] = await ethers.getSigners();
@@ -187,4 +189,69 @@ describe("AcceleratingDistributor: Admin Functions", async function () {
     await expect(distributor.connect(owner).withdrawReward(lpToken1.address)).to.not.be.reverted;
     await expect(distributor.connect(owner).exit(lpToken1.address)).to.not.be.reverted;
   });
+
+  it("Input validation on staking-related methods", async function () {
+    await lpToken1.mint(owner.address, toWei(69));
+    await lpToken1.connect(owner).approve(distributor.address, toWei(69));
+
+    // Modifiers take precedence when staking is disabled.
+    for (const stakingEnabled of [true, false]) {
+      await distributor.configureStakingToken(
+        lpToken1.address,
+        stakingEnabled,
+        baseEmissionRate,
+        maxMultiplier,
+        secondsToMaxMultiplier
+      );
+
+      await expect(distributor.connect(owner).stake(lpToken1.address, toWei(0))).to.be.revertedWith(
+        stakingEnabled ? "Invalid amount" : "stakedToken not enabled"
+      );
+      await expect(distributor.connect(owner).stakeFor(lpToken1.address, toWei(0), zeroAddress)).to.be.revertedWith(
+        stakingEnabled ? "Invalid beneficiary" : "stakedToken not enabled"
+      );
+      await expect(distributor.connect(owner).stakeFor(lpToken1.address, toWei(1), zeroAddress)).to.be.revertedWith(
+        stakingEnabled ? "Invalid beneficiary" : "stakedToken not enabled"
+      );
+
+      if (stakingEnabled) {
+        await expect(distributor.connect(owner).stake(lpToken1.address, toWei(1))).to.not.be.reverted;
+        await expect(distributor.connect(owner).unstake(lpToken1.address, toWei(0))).to.be.revertedWith(
+          "Invalid amount"
+        );
+        await expect(distributor.connect(owner).unstake(lpToken1.address, toWei(1))).to.not.be.reverted;
+      } else {
+        await expect(distributor.connect(owner).stake(lpToken1.address, toWei(1))).to.be.revertedWith(
+          "stakedToken not enabled"
+        );
+        await expect(distributor.connect(owner).unstake(lpToken1.address, toWei(0))).to.be.revertedWith(
+          "Invalid amount"
+        );
+
+        // Staked balance is 0.
+        await expect(distributor.connect(owner).unstake(lpToken1.address, toWei(1))).to.be.reverted;
+      }
+    }
+
+    // Validate withdrawal guards when staking is disabled.
+    await distributor.configureStakingToken(
+      lpToken1.address,
+      true,
+      baseEmissionRate,
+      maxMultiplier,
+      secondsToMaxMultiplier
+    );
+    await expect(distributor.connect(owner).stake(lpToken1.address, toWei(2))).to.not.be.reverted;
+
+    await distributor.configureStakingToken(
+      lpToken1.address,
+      false,
+      baseEmissionRate,
+      maxMultiplier,
+      secondsToMaxMultiplier
+    );
+    await expect(distributor.connect(owner).unstake(lpToken1.address, toWei(0))).to.be.revertedWith("Invalid amount");
+    await expect(distributor.connect(owner).unstake(lpToken1.address, toWei(1))).to.not.be.reverted;
+    await expect(distributor.connect(owner).exit(lpToken1.address)).to.not.be.reverted;
+  });
 });