feat: add proxy contract and upgradeability

treenteq · Feb 21, 2025 · 1517595 · 1517595
1 parent 1ef6fa8
commit 1517595
Show file tree

Hide file tree

Showing 8 changed files with 212 additions and 37 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -4,3 +4,6 @@
 [submodule "lib/openzeppelin-contracts"]
 	path = lib/openzeppelin-contracts
 	url = https://github.com/OpenZeppelin/openzeppelin-contracts
+[submodule "lib/openzeppelin-contracts-upgradeable"]
+	path = lib/openzeppelin-contracts-upgradeable
+	url = https://github.com/OpenZeppelin/openzeppelin-contracts-upgradeable
diff --git a/foundry.toml b/foundry.toml
@@ -2,7 +2,10 @@
 src = "src"
 out = "out"
 libs = ["lib"]
-remappings = ["@openzeppelin/=lib/openzeppelin-contracts/"]
+remappings = [
+    "@openzeppelin/contracts/=lib/openzeppelin-contracts/contracts/",
+    "@openzeppelin/contracts-upgradeable/=lib/openzeppelin-contracts-upgradeable/contracts/"
+]
 via-ir = true
 
 # See more config options https://github.com/foundry-rs/foundry/blob/master/crates/config/README.md#all-options

diff --git a/lib/openzeppelin-contracts-upgradeable b/lib/openzeppelin-contracts-upgradeable
diff --git a/script/DeployBondingCurve.s.sol b/script/DeployBondingCurve.s.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.18;
 import {Script} from "forge-std/Script.sol";
 import {DatasetBondingCurve} from "../src/DatasetBondingCurve.sol";
 import {DatasetToken} from "../src/DeployDataset.sol";
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 
 contract DeployBondingCurve is Script {
     function run() external returns (DatasetBondingCurve, DatasetToken) {
@@ -13,20 +14,42 @@ contract DeployBondingCurve is Script {
 
         vm.startBroadcast(deployerPrivateKey);
 
-        // First deploy the DatasetToken contract
-        DatasetToken datasetToken = new DatasetToken(
+        // Deploy implementation contracts
+        DatasetToken datasetTokenImpl = new DatasetToken();
+        DatasetBondingCurve bondingCurveImpl = new DatasetBondingCurve();
+
+        // Prepare initialization data for DatasetToken
+        bytes memory datasetTokenData = abi.encodeWithSelector(
+            DatasetToken.initialize.selector,
             "ipfs://",
             deployerAddress,
             usdcAddress
         );
 
-        // Then deploy the BondingCurve contract
-        DatasetBondingCurve bondingCurve = new DatasetBondingCurve(
+        // Deploy DatasetToken proxy
+        ERC1967Proxy datasetTokenProxy = new ERC1967Proxy(
+            address(datasetTokenImpl),
+            datasetTokenData
+        );
+        DatasetToken datasetToken = DatasetToken(address(datasetTokenProxy));
+
+        // Prepare initialization data for BondingCurve
+        bytes memory bondingCurveData = abi.encodeWithSelector(
+            DatasetBondingCurve.initialize.selector,
             address(datasetToken),
             usdcAddress,
             deployerAddress
         );
 
+        // Deploy BondingCurve proxy
+        ERC1967Proxy bondingCurveProxy = new ERC1967Proxy(
+            address(bondingCurveImpl),
+            bondingCurveData
+        );
+        DatasetBondingCurve bondingCurve = DatasetBondingCurve(
+            address(bondingCurveProxy)
+        );
+
         // Set the bonding curve in the dataset token contract
         datasetToken.setBondingCurve(address(bondingCurve));
 
@@ -35,3 +58,35 @@ contract DeployBondingCurve is Script {
         return (bondingCurve, datasetToken);
     }
 }
+
+// Upgrade script for future upgrades
+contract UpgradeContracts is Script {
+    function run(address proxyAddress, string memory contractType) external {
+        uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
+
+        vm.startBroadcast(deployerPrivateKey);
+
+        if (
+            keccak256(bytes(contractType)) == keccak256(bytes("DatasetToken"))
+        ) {
+            // Deploy new DatasetToken implementation
+            DatasetToken newImpl = new DatasetToken();
+
+            // Upgrade proxy to new implementation
+            DatasetToken(proxyAddress).upgradeToAndCall(address(newImpl), "");
+        } else if (
+            keccak256(bytes(contractType)) == keccak256(bytes("BondingCurve"))
+        ) {
+            // Deploy new BondingCurve implementation
+            DatasetBondingCurve newImpl = new DatasetBondingCurve();
+
+            // Upgrade proxy to new implementation
+            DatasetBondingCurve(proxyAddress).upgradeToAndCall(
+                address(newImpl),
+                ""
+            );
+        }
+
+        vm.stopBroadcast();
+    }
+}
diff --git a/src/DatasetBondingCurve.sol b/src/DatasetBondingCurve.sol
@@ -1,13 +1,19 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.18;
 
-import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import {Math} from "@openzeppelin/contracts/utils/math/Math.sol";
 import {DatasetToken} from "./DeployDataset.sol";
 import {IUSDC} from "./interfaces/IUSDC.sol";
+import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import {console2} from "forge-std/console2.sol";
 
-contract DatasetBondingCurve is Ownable {
+contract DatasetBondingCurve is
+    Initializable,
+    OwnableUpgradeable,
+    UUPSUpgradeable
+{
     using Math for uint256;
 
     // The Dataset Token contract
@@ -33,20 +39,42 @@ contract DatasetBondingCurve is Ownable {
     event PriceCalculated(uint256 indexed tokenId, uint256 price);
     event InitialPriceSet(uint256 indexed tokenId, uint256 initialPrice);
 
-    constructor(
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
+    /**
+     * @dev Initializer function to replace constructor
+     * @param datasetTokenAddress The address of the dataset token contract
+     * @param usdcAddress The address of the USDC token contract
+     * @param initialOwner The address of the initial owner
+     */
+    function initialize(
         address datasetTokenAddress,
         address usdcAddress,
         address initialOwner
-    ) Ownable(initialOwner) {
+    ) public initializer {
         require(
             datasetTokenAddress != address(0),
             "Invalid dataset token address"
         );
         require(usdcAddress != address(0), "Invalid USDC address");
+
+        __Ownable_init(initialOwner);
+        __UUPSUpgradeable_init();
+
         datasetToken = DatasetToken(datasetTokenAddress);
         usdc = IUSDC(usdcAddress);
     }
 
+    /**
+     * @dev Required override for UUPS proxy upgrade authorization
+     */
+    function _authorizeUpgrade(
+        address newImplementation
+    ) internal override onlyOwner {}
+
     /**
      * @dev Set the initial price for a token's bonding curve
      * @param tokenId The ID of the token

diff --git a/src/DeployDataset.sol b/src/DeployDataset.sol
@@ -1,13 +1,21 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.18;
 
-import {ERC1155} from "@openzeppelin/contracts/token/ERC1155/ERC1155.sol";
-import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
-import {ReentrancyGuard} from "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
+import {ERC1155Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC1155/ERC1155Upgradeable.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {ReentrancyGuardUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import {DatasetBondingCurve} from "./DatasetBondingCurve.sol";
 import {IUSDC} from "./interfaces/IUSDC.sol";
 
-contract DatasetToken is ERC1155, Ownable, ReentrancyGuard {
+contract DatasetToken is
+    Initializable,
+    ERC1155Upgradeable,
+    OwnableUpgradeable,
+    ReentrancyGuardUpgradeable,
+    UUPSUpgradeable
+{
     // Token ID counter
     uint256 private _currentTokenId;
 
@@ -68,29 +76,38 @@ contract DatasetToken is ERC1155, Ownable, ReentrancyGuard {
 
     event BondingCurveUpdated(address indexed newBondingCurve);
 
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
     /**
-     * @dev Constructor to initialize the contract.
-     * @param uri The base URI for metadata.
-     * @param initialOwner The address of the initial owner.
-     * @param usdcAddress The address of the USDC contract.
+     * @dev Initializer function to replace constructor
+     * @param uri The base URI for metadata
+     * @param initialOwner The address of the initial owner
+     * @param usdcAddress The address of the USDC contract
      */
-    constructor(
+    function initialize(
         string memory uri,
         address initialOwner,
         address usdcAddress
-    ) ERC1155(uri) Ownable(initialOwner) {
+    ) public initializer {
         require(usdcAddress != address(0), "Invalid USDC address");
+
+        __ERC1155_init(uri);
+        __Ownable_init(initialOwner);
+        __ReentrancyGuard_init();
+        __UUPSUpgradeable_init();
+
         usdc = IUSDC(usdcAddress);
     }
 
     /**
-     * @dev Set the USDC contract address
-     * @param _usdcAddress The address of the USDC contract
+     * @dev Required override for UUPS proxy upgrade authorization
      */
-    function setUsdcAddress(address _usdcAddress) external onlyOwner {
-        require(_usdcAddress != address(0), "Invalid USDC address");
-        usdc = IUSDC(_usdcAddress);
-    }
+    function _authorizeUpgrade(
+        address newImplementation
+    ) internal override onlyOwner {}
 
     /**
      * @dev Validate ownership percentages add up to 10000 (100%)
@@ -196,7 +213,7 @@ contract DatasetToken is ERC1155, Ownable, ReentrancyGuard {
             "Already purchased this dataset"
         );
         uint256 currentPrice = bondingCurve.getCurrentPrice(tokenId);
-        
+
         // Check USDC allowance
         require(
             usdc.allowance(msg.sender, address(this)) >= currentPrice,
@@ -216,8 +233,8 @@ contract DatasetToken is ERC1155, Ownable, ReentrancyGuard {
             address owner = metadata.owners[i].owner;
 
             // Calculate owner's share of the payment
-            uint256 ownerShare = (currentPrice * metadata.owners[i].percentage) /
-                10000;
+            uint256 ownerShare = (currentPrice *
+                metadata.owners[i].percentage) / 10000;
 
             // Transfer USDC to owner
             require(

diff --git a/test/DatasetBondingCurve.t.sol b/test/DatasetBondingCurve.t.sol
@@ -5,10 +5,13 @@ import {Test, console2} from "forge-std/Test.sol";
 import {DatasetToken} from "../src/DeployDataset.sol";
 import {DatasetBondingCurve} from "../src/DatasetBondingCurve.sol";
 import {MockUSDC} from "./mocks/MockUSDC.sol";
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 
 contract DatasetBondingCurveTest is Test {
     DatasetBondingCurve public bondingCurve;
+    DatasetBondingCurve public bondingCurveImpl;
     DatasetToken public datasetToken;
+    DatasetToken public datasetTokenImpl;
     MockUSDC public mockUsdc;
     address public owner;
     address public user1;
@@ -30,15 +33,35 @@ contract DatasetBondingCurveTest is Test {
         // Deploy mock USDC
         mockUsdc = new MockUSDC();
 
-        // Deploy DatasetToken
-        datasetToken = new DatasetToken("ipfs://", owner, address(mockUsdc));
+        // Deploy implementation contracts
+        datasetTokenImpl = new DatasetToken();
+        bondingCurveImpl = new DatasetBondingCurve();
 
-        // Deploy BondingCurve
-        bondingCurve = new DatasetBondingCurve(
+        // Deploy and initialize DatasetToken proxy
+        bytes memory datasetTokenData = abi.encodeWithSelector(
+            DatasetToken.initialize.selector,
+            "ipfs://",
+            owner,
+            address(mockUsdc)
+        );
+        ERC1967Proxy datasetTokenProxy = new ERC1967Proxy(
+            address(datasetTokenImpl),
+            datasetTokenData
+        );
+        datasetToken = DatasetToken(address(datasetTokenProxy));
+
+        // Deploy and initialize BondingCurve proxy
+        bytes memory bondingCurveData = abi.encodeWithSelector(
+            DatasetBondingCurve.initialize.selector,
             address(datasetToken),
             address(mockUsdc),
             owner
         );
+        ERC1967Proxy bondingCurveProxy = new ERC1967Proxy(
+            address(bondingCurveImpl),
+            bondingCurveData
+        );
+        bondingCurve = DatasetBondingCurve(address(bondingCurveProxy));
 
         // Set bonding curve in dataset token
         datasetToken.setBondingCurve(address(bondingCurve));
@@ -56,6 +79,17 @@ contract DatasetBondingCurveTest is Test {
         mockUsdc.mint(user2, 100_000_000); // 100 USDC
     }
 
+    function test_UpgradeBondingCurve() public {
+        // Deploy new implementation
+        DatasetBondingCurve newImpl = new DatasetBondingCurve();
+
+        // Upgrade to new implementation
+        bondingCurve.upgradeToAndCall(address(newImpl), "");
+
+        // Verify upgrade
+        assertEq(bondingCurve.owner(), owner);
+    }
+
     function test_IndependentBondingCurves() public {
         vm.startPrank(owner);