From 4ecba8b3a8b2cd9bf38e63811f0afddb1b7223d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Axel=20Wikstr=C3=B6m?= Date: Tue, 20 Feb 2024 17:14:21 +0100 Subject: [PATCH] new multiclam function with test --- src/royalty/ABClaim.sol | 93 +++++++++++++++++++++----------------- test/royalty/ABClaim.t.sol | 55 +++++++++++++++++++++- 2 files changed, 105 insertions(+), 43 deletions(-) diff --git a/src/royalty/ABClaim.sol b/src/royalty/ABClaim.sol index 3dfff18..008f0e7 100644 --- a/src/royalty/ABClaim.sol +++ b/src/royalty/ABClaim.sol @@ -291,7 +291,7 @@ contract ABClaim is Initializable, AccessControlUpgradeable { // abKycModule.beforeRoyaltyClaim(_user, _signature); } - function _claimMultiDrop( +function _claimMultiDrop( uint256[] calldata _dropIds, uint256[][] calldata _tokenIds, address _user, @@ -299,59 +299,68 @@ contract ABClaim is Initializable, AccessControlUpgradeable { ) internal { // Enforce KYC Signature validity through ABKYCModule contract _beforeClaim(_user, _signature); - - // Check parameters validity - uint256 dLength = _dropIds.length; - if (dLength != _tokenIds.length) revert ABErrors.INVALID_PARAMETER(); + // Check parameters validity + if (_dropIds.length != _tokenIds.length) revert ABErrors.INVALID_PARAMETER(); uint256 totalClaimable; - for (uint256 i; i < dLength;) { - uint256 dropId = _dropIds[i]; - + for (uint256 i; i < _dropIds.length;) { // Get drop data - ABDataTypes.DropData memory data = dropData[dropId]; + ABDataTypes.DropData memory data = dropData[_dropIds[i]]; + // get the claimableAmount depening on L1 or Base + uint256 dropClaimable = data.isL1 ? + handleL1Drop(_dropIds[i], _tokenIds[i], _user) : + handleBaseDrop(_dropIds[i], _tokenIds[i], _user, data.nft); + + //emit event with dropClaimable from the handler + emit ABEvents.RoyaltyClaimed(_dropIds[i], _tokenIds[i], dropClaimable, _user); + totalClaimable += dropClaimable; - // Calculate royalties per token - uint256 royaltiesPerToken = totalDepositedPerDrop[dropId] / data.supply; + unchecked { + ++i; + } + } + //transfer the totalClaimable + USDC.transfer(_user, totalClaimable); + } + + function handleL1Drop(uint256 dropId, uint256[] calldata tokenIds, address user) internal returns (uint256) { + uint256 totalClaimable; + uint256 royaltiesPerToken = totalDepositedPerDrop[dropId] / dropData[dropId].supply; - // Check if the drop is on Layer 1 or on Base - if (data.isL1) { - for (uint256 j; j < _tokenIds[i].length;) { - // Enforce token ownership - if (ownerOf[dropId][_tokenIds[i][j]] != _user) revert ABErrors.NOT_TOKEN_OWNER(); + for (uint256 j; j < tokenIds.length;) { + if (ownerOf[dropId][tokenIds[j]] != user) revert ABErrors.NOT_TOKEN_OWNER(); - // Calculate claimable amount - uint256 _claimable = royaltiesPerToken - claimedAmount[dropId][_tokenIds[i][j]]; - totalClaimable += _claimable; - claimedAmount[dropId][_tokenIds[i][j]] += _claimable; + uint256 claimable = royaltiesPerToken - claimedAmount[dropId][tokenIds[j]]; + totalClaimable += claimable; + claimedAmount[dropId][tokenIds[j]] += claimable; - unchecked { - ++j; - } - } - } else { - for (uint256 j; j < _tokenIds[i].length;) { - // Enforce token ownership - if (IERC721AB(data.nft).ownerOf(_tokenIds[i][j]) != _user) revert ABErrors.NOT_TOKEN_OWNER(); - - // Calculate claimable amount - uint256 _claimable = royaltiesPerToken - claimedAmount[dropId][_tokenIds[i][j]]; - totalClaimable += _claimable; - claimedAmount[dropId][_tokenIds[i][j]] += _claimable; - - unchecked { - ++j; - } - } + unchecked { + ++j; } - emit ABEvents.RoyaltyClaimed(_dropIds[i], _tokenIds[i], totalClaimable, _user); + } + + return totalClaimable; + } + + function handleBaseDrop(uint256 dropId, uint256[] calldata tokenIds, address user, address nftAddress) internal returns (uint256) { + uint256 totalClaimable; + uint256 royaltiesPerToken = totalDepositedPerDrop[dropId] / dropData[dropId].supply; + IERC721AB nft = IERC721AB(nftAddress); + + for (uint256 j; j < tokenIds.length;) { + if (nft.ownerOf(tokenIds[j]) != user) revert ABErrors.NOT_TOKEN_OWNER(); + + uint256 claimable = royaltiesPerToken - claimedAmount[dropId][tokenIds[j]]; + totalClaimable += claimable; + claimedAmount[dropId][tokenIds[j]] += claimable; + unchecked { - ++i; + ++j; } } - // Transfer total claimable amount to the shareholder - USDC.transfer(_user, totalClaimable); + + return totalClaimable; } function _claimSingleDrop(uint256 _dropId, uint256[] calldata _tokenIds, address _user, bytes calldata _signature) diff --git a/test/royalty/ABClaim.t.sol b/test/royalty/ABClaim.t.sol index a98f3a7..a17266f 100644 --- a/test/royalty/ABClaim.t.sol +++ b/test/royalty/ABClaim.t.sol @@ -8,7 +8,7 @@ import {ABClaim} from "src/royalty/ABClaim.sol"; import {ABErrors} from "src/libraries/ABErrors.sol"; import {MockNFT} from "test/_mocks/MockNFT.sol"; import {MockToken} from "test/_mocks/MockToken.sol"; - +import "../../lib/openzeppelin-contracts/contracts/utils/Strings.sol"; import {TransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol"; import {ProxyAdmin} from "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol"; import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; @@ -940,6 +940,59 @@ contract ABClaimTest is Test { assertEq(abClaim.getClaimableAmount(dropIds, uTokenIds), 0); } + function test_claim_multiScaleBaseDrop(address _sender, address u1, uint256 scaleAmount) public { + vm.assume(_sender != address(0)); + vm.assume(u1 != address(0)); + vm.assume(scaleAmount < 50); + + abClaim.grantRole(DEFAULT_ADMIN_ROLE_HASH, _sender); + + uint256[] memory dropIds = new uint256[](scaleAmount); + uint256[] memory amounts = new uint256[](scaleAmount); + uint256 amount = 6000; + MockNFT[] memory nfts = new MockNFT[](scaleAmount); + + for (uint256 i; i < scaleAmount; i++) { + amounts[i] = amount; + } + + mockUSD.mint(_sender, amount * scaleAmount); + + vm.startPrank(_sender); + mockUSD.approve(address(abClaim), amount * scaleAmount); + abClaim.depositRoyalty(dropIds, amounts); + vm.stopPrank(); + + + for (uint256 i; i < dropIds.length; i++ ) { + dropIds[i] = i; + string memory nftName = string(abi.encodePacked("NFT", Strings.toString(i))); + string memory nftSymbol = string(abi.encodePacked("AB", Strings.toString(i))); + nfts[i] = new MockNFT(nftName, nftSymbol); + abClaim.setDropData(dropIds[i], address(nfts[i]), false, 3); + } + uint256[][] memory u1TokenIdsParent = new uint256[][](scaleAmount); + + for (uint256 i; i < nfts.length; i++) { + nfts[i].mint(u1, 3); + + uint256[] memory tokenIds = new uint256[](3); + + for (uint256 j; j < 3; j++) { + tokenIds[j] = j; + } + u1TokenIdsParent[i] = tokenIds; + } + + bytes memory sigU1 = _generateValidSignature(kycSignerPkey, u1, 0); + + vm.prank(u1); + abClaim.claim(dropIds, u1TokenIdsParent, sigU1); + assertEq(mockUSD.balanceOf(u1), amount * scaleAmount); + assertEq(abClaim.getClaimableAmount(dropIds, u1TokenIdsParent), 0); + + } + function _generateValidSignature(uint256 _signerPkey, address _user, uint256 _nonce) internal pure