Fix claim contracts implementation

Author: ilitteri

claim_contracts/src/AlignedToken.sol

@@ -3,24 +3,60 @@ pragma solidity ^0.8.19;
 
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/token/ERC20/ERC20Upgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/token/ERC20/extensions/ERC20PermitUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
+
+contract AlignedToken is
+    Initializable,
+    ERC20Upgradeable,
+    EIP712Upgradeable,
+    ERC20PermitUpgradeable,
+    Ownable2StepUpgradeable
+{
+    /// @notice Name of the token.
+    string public constant NAME = "Aligned Token";
+
+    /// @notice Symbol of the token.
+    string public constant SYMBOL = "ALIGN";
+
+    /// @notice Version of the token implementation.
+    string public constant VERSION = "1";
+
+    /// @notice Total supply of the token.
+    uint256 public constant TOTAL_SUPPLY = 10_000_000_000e18; // 10 billion
 
-contract AlignedToken is Initializable, ERC20Upgradeable {
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
     }
 
+    /// @notice Initializes the contract.
+    /// @dev This initializer should be called only once.
+    /// @param _owner address of the owner of the token.
+    /// @param _foundation address of the foundation.
+    /// @param _claim address of the claim.
     function initialize(
-        address _beneficiary1,
-        uint256 _beneficiary1Part,
-        address _beneficiary2,
-        uint256 _beneficiary2Part,
-        address _beneficiary3,
-        uint256 _beneficiary3Part
+        address _owner,
+        address _foundation,
+        address _claim
     ) public initializer {
-        __ERC20_init("AlignedToken", "ALI");
-        _mint(_beneficiary1, _beneficiary1Part);
-        _mint(_beneficiary2, _beneficiary2Part);
-        _mint(_beneficiary3, _beneficiary3Part);
+        require(
+            _foundation != address(0) &&
+                _claim != address(0) &&
+                _owner != address(0),
+            "Invalid _foundation or _claim or _owner"
+        );
+        __ERC20_init(NAME, SYMBOL);
+        __EIP712_init(NAME, VERSION);
+        __ERC20Permit_init(NAME);
+        __Ownable2Step_init(); // default is msg.sender
+        _transferOwnership(_owner);
+        _mint(_foundation, 7_300_000_000e18); // 7.3 billion
+        _mint(_claim, 2_700_000_000e18); // 2.7 billion
+    }
+
+    /// @notice Prevents the owner from renouncing ownership.
+    function renounceOwnership() public view override onlyOwner {
+        revert("Cannot renounce ownership");
     }
 }

claim_contracts/src/ClaimableAirdrop.sol

@@ -6,56 +6,77 @@ import "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol";
 import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
-import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 
 contract ClaimableAirdrop is
     ReentrancyGuard,
     Initializable,
     PausableUpgradeable,
-    OwnableUpgradeable
+    Ownable2StepUpgradeable
 {
-    address public tokenContractAddress;
-    address public tokenOwnerAddress;
+    /// @notice Address of the token contract to claim.
+    address public tokenProxy;
+
+    /// @notice Address of the wallet that has the tokens to distribute to the claimants.
+    address public claimSupplier;
+
+    /// @notice Timestamp until which the claimants can claim the tokens.
     uint256 public limitTimestampToClaim;
+
+    /// @notice Merkle root of the claimants.
     bytes32 public claimMerkleRoot;
 
+    /// @notice Mapping of the claimants that have claimed the tokens.
+    /// @dev true if the claimant has claimed the tokens.
     mapping(address => bool) public hasClaimed;
 
+    /// @notice Event emitted when a claimant claims the tokens.
     event TokenClaimed(address indexed to, uint256 indexed amount);
 
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
     }
 
+    /// @notice Initializes the contract.
+    /// @dev This initializer should be called only once.
+    /// @param _owner address of the owner of the token.
+    /// @param _tokenProxy address of the token contract.
+    /// @param _claimSupplier address of the wallet that has the tokens to distribute to the claimants.
+    /// @param _limitTimestampToClaim timestamp until which the claimants can claim the tokens.
+    /// @param _claimMerkleRoot Merkle root of the claimants.
     function initialize(
-        address _multisig,
-        address _tokenContractAddress,
-        address _tokenOwnerAddress,
+        address _owner,
+        address _tokenProxy,
+        address _claimSupplier,
         uint256 _limitTimestampToClaim,
         bytes32 _claimMerkleRoot
     ) public initializer {
-        __Ownable_init(_multisig);
-        __Pausable_init();
-
-        require(_multisig != address(0), "Invalid multisig address");
+        require(_owner != address(0), "Invalid owner address");
         require(
-            _tokenContractAddress != address(0),
+            _tokenProxy != address(0) && _tokenProxy != address(this),
             "Invalid token contract address"
         );
         require(
-            _tokenOwnerAddress != address(0),
+            _claimSupplier != address(0) && _claimSupplier != address(this),
             "Invalid token owner address"
         );
         require(_limitTimestampToClaim > block.timestamp, "Invalid timestamp");
         require(_claimMerkleRoot != 0, "Invalid Merkle root");
 
-        tokenContractAddress = _tokenContractAddress;
-        tokenOwnerAddress = _tokenOwnerAddress;
+        __Ownable2Step_init(); // default is msg.sender
+        _transferOwnership(_owner);
+        __Pausable_init();
+
+        tokenProxy = _tokenProxy;
+        claimSupplier = _claimSupplier;
         limitTimestampToClaim = _limitTimestampToClaim;
         claimMerkleRoot = _claimMerkleRoot;
     }
 
+    /// @notice Claim the tokens.
+    /// @param amount amount of tokens to claim.
+    /// @param merkleProof Merkle proof of the claim.
     function claim(
         uint256 amount,
         bytes32[] calldata merkleProof
@@ -77,15 +98,13 @@ contract ClaimableAirdrop is
         require(verifies, "Invalid Merkle proof");
 
         require(
-            IERC20(tokenContractAddress).allowance(
-                tokenOwnerAddress,
-                address(this)
-            ) >= amount,
+            IERC20(tokenProxy).allowance(claimSupplier, address(this)) >=
+                amount,
             "Insufficient token allowance"
         );
 
-        bool success = IERC20(tokenContractAddress).transferFrom(
-            tokenOwnerAddress,
+        bool success = IERC20(tokenProxy).transferFrom(
+            claimSupplier,
             msg.sender,
             amount
         );
@@ -97,11 +116,33 @@ contract ClaimableAirdrop is
         emit TokenClaimed(msg.sender, amount);
     }
 
-    function pause() public onlyOwner {
+    /// @notice Update the Merkle root.
+    /// @param newRoot new Merkle root.
+    function updateMerkleRoot(bytes32 newRoot) external whenPaused onlyOwner {
+        require(newRoot != 0 && newRoot != claimMerkleRoot, "Invalid root");
+        claimMerkleRoot = newRoot;
+    }
+
+    /// @notice Extend the claim period.
+    /// @param newTimestamp new timestamp until which the claimants can claim the tokens.
+    function extendClaimPeriod(
+        uint256 newTimestamp
+    ) external whenPaused onlyOwner {
+        require(
+            newTimestamp > limitTimestampToClaim &&
+                newTimestamp > block.timestamp,
+            "Can only extend from current timestamp"
+        );
+        limitTimestampToClaim = newTimestamp;
+    }
+
+    /// @notice Pause the contract.
+    function pause() public whenNotPaused onlyOwner {
         _pause();
     }
 
-    function unpause() public onlyOwner {
+    /// @notice Unpause the contract.
+    function unpause() public whenPaused onlyOwner {
         _unpause();
     }
 }