From 981fc870badc8e5db0e3c303e113b5015a17c3a2 Mon Sep 17 00:00:00 2001 From: Anh Hoang Nguyen Date: Wed, 18 Dec 2024 13:29:32 +0700 Subject: [PATCH 1/4] feat: add support for managing ERC-20 tokens --- src/Wallet.sol | 36 +++++++++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/src/Wallet.sol b/src/Wallet.sol index 7fc7cca..995823d 100644 --- a/src/Wallet.sol +++ b/src/Wallet.sol @@ -11,6 +11,7 @@ contract Wallet { mapping(address => Transaction[]) public transactions; mapping(address => mapping(address => uint256)) token_balance; + mapping(address => bool) public supportedTokens; struct Transaction { uint256 amount; @@ -37,20 +38,38 @@ contract Wallet { function createWorldId() external onlyOwner {} - function transfer(address _recipient, uint256 _amount) external onlyVerified(msg.sender) { + function transfer( + address _recipient, + uint256 _amount + ) external onlyVerified(msg.sender) { require(_recipient != address(0), "Zero address detected"); require(usdt.balanceOf(msg.sender) >= _amount, "Insufficient balance"); require(_amount > 0, "Transfer amount must be greater than zero"); - require(usdt.transferFrom(msg.sender, _recipient, _amount), "Transfer failed"); + require( + usdt.transferFrom(msg.sender, _recipient, _amount), + "Transfer failed" + ); recordTransactionHistory(msg.sender, _amount, address(usdt)); } + function addSupportedToken(address _token) external onlyOwner { + require(_token != address(0), "Invalid token address"); + supportedTokens[_token] = true; + } + + function removeSupportedToken(address _token) external onlyOwner { + require(supportedTokens[_token], "Token not supported"); + supportedTokens[_token] = false; + } + ////////////////////////////////////////////// // View Functions // //////////////////////////////////////////// - function getTransactionHistory(address _user) external view returns (Transaction[] memory) { + function getTransactionHistory( + address _user + ) external view returns (Transaction[] memory) { return transactions[_user]; } @@ -58,8 +77,15 @@ contract Wallet { // Private Function // ////////////////////////////////////////////// - function recordTransactionHistory(address _user, uint256 _amount, address _token) private { - Transaction memory newTransaction = Transaction({amount: _amount, token: _token}); + function recordTransactionHistory( + address _user, + uint256 _amount, + address _token + ) private { + Transaction memory newTransaction = Transaction({ + amount: _amount, + token: _token + }); transactions[_user].push(newTransaction); } From fce4e6ddf25d71a1ad190ecba5f34ad66565bd53 Mon Sep 17 00:00:00 2001 From: Anh Hoang Nguyen Date: Wed, 18 Dec 2024 15:42:33 +0700 Subject: [PATCH 2/4] feat: transfer erc20 token --- src/Wallet.sol | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/Wallet.sol b/src/Wallet.sol index 995823d..6fc7aa1 100644 --- a/src/Wallet.sol +++ b/src/Wallet.sol @@ -29,28 +29,36 @@ contract Wallet { _; } + modifier onlySupportedToken(address _token) { + require(supportedTokens[_token], "Token not supported"); + _; + } + constructor(address _owner, address _worldID, address _usdt) { require(msg.sender != address(0), "zero address found"); owner = _owner; worldID = IWorldID(_worldID); usdt = IERC20(_usdt); + supportedTokens[_usdt] = true; // Add USDT as a default supported token } function createWorldId() external onlyOwner {} function transfer( address _recipient, + address _token, uint256 _amount - ) external onlyVerified(msg.sender) { + ) external onlyVerified(msg.sender) onlySupportedToken(_token) { require(_recipient != address(0), "Zero address detected"); - require(usdt.balanceOf(msg.sender) >= _amount, "Insufficient balance"); + IERC20 token = IERC20(_token); + require(token.balanceOf(msg.sender) >= _amount, "Insufficient balance"); require(_amount > 0, "Transfer amount must be greater than zero"); require( - usdt.transferFrom(msg.sender, _recipient, _amount), + token.transferFrom(msg.sender, _recipient, _amount), "Transfer failed" ); - recordTransactionHistory(msg.sender, _amount, address(usdt)); + recordTransactionHistory(msg.sender, _amount, _token); } function addSupportedToken(address _token) external onlyOwner { From 659f97021fd11c87971dfc31ab3ea9005decd5c2 Mon Sep 17 00:00:00 2001 From: Anh Hoang Nguyen Date: Wed, 18 Dec 2024 15:47:53 +0700 Subject: [PATCH 3/4] chore: update test for transfer function --- test/Transfer.t.sol | 14 +++++++------- test/Wallet.t.sol | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/test/Transfer.t.sol b/test/Transfer.t.sol index 2a19af8..a5f5083 100644 --- a/test/Transfer.t.sol +++ b/test/Transfer.t.sol @@ -121,7 +121,7 @@ contract USDTTransferTest is Test { // Perform the transfer vm.prank(user1); - transferContract.transfer(user2, transferAmount); + transferContract.transfer(user2, address(mockUSDT), transferAmount); // Assert balances after transfer assertEq(mockUSDT.balanceOf(user1), user1BalanceBefore - transferAmount); @@ -134,7 +134,7 @@ contract USDTTransferTest is Test { vm.prank(user1); vm.expectRevert("Insufficient balance"); - transferContract.transfer(user2, transferAmount); + transferContract.transfer(user2, address(mockUSDT), transferAmount); } // Test transfer function - Unverified user @@ -146,7 +146,7 @@ contract USDTTransferTest is Test { vm.prank(user1); vm.expectRevert("User not verified"); - transferContract.transfer(user2, transferAmount); + transferContract.transfer(user2, address(mockUSDT), transferAmount); } // Test transfer function - Transfer to unverified recipient @@ -157,7 +157,7 @@ contract USDTTransferTest is Test { vm.expectRevert("User not verified"); - transferContract.transfer(user2, transferAmount); + transferContract.transfer(user2, address(mockUSDT), transferAmount); } // Test transfer function - Multiple consecutive transfers @@ -171,7 +171,7 @@ contract USDTTransferTest is Test { // First transfer vm.prank(user1); - transferContract.transfer(user2, transferAmount1); + transferContract.transfer(user2, address(mockUSDT), transferAmount1); // Assert balances after first transfer assertEq(mockUSDT.balanceOf(user1), user1BalanceBefore - transferAmount1); @@ -182,7 +182,7 @@ contract USDTTransferTest is Test { user2BalanceBefore = mockUSDT.balanceOf(user2); vm.prank(user1); - transferContract.transfer(user2, transferAmount2); + transferContract.transfer(user2, address(mockUSDT), transferAmount2); assertEq(mockUSDT.balanceOf(user1), user1BalanceBefore - transferAmount2); assertEq(mockUSDT.balanceOf(user2), user2BalanceBefore + transferAmount2); @@ -195,6 +195,6 @@ contract USDTTransferTest is Test { vm.prank(user1); vm.expectRevert("Transfer amount must be greater than zero"); - transferContract.transfer(user2, transferAmount); + transferContract.transfer(user2, address(mockUSDT), transferAmount); } } diff --git a/test/Wallet.t.sol b/test/Wallet.t.sol index 7e6f61a..ad5c623 100644 --- a/test/Wallet.t.sol +++ b/test/Wallet.t.sol @@ -48,7 +48,7 @@ contract WalletTest is Test { /// @notice Test recording a single transaction function testRecordSingleTransaction() public { vm.prank(user1); - wallet.transfer(user2, 100); + wallet.transfer(user2, address(usdt), 100); vm.prank(user1); Wallet.Transaction[] memory history = wallet.getTransactionHistory(user1); @@ -60,8 +60,8 @@ contract WalletTest is Test { /// @notice Test recording multiple transactions function testRecordMultipleTransactions() public { vm.startPrank(user1); - wallet.transfer(user2, 100); - wallet.transfer(user2, 200); + wallet.transfer(user2, address(usdt), 100); + wallet.transfer(user2, address(usdt), 200); Wallet.Transaction[] memory history = wallet.getTransactionHistory(user1); vm.stopPrank(); @@ -76,13 +76,13 @@ contract WalletTest is Test { /// @notice Test recording transactions for different users function testRecordTransactionsForDifferentUsers() public { vm.prank(user1); - wallet.transfer(user2, 100); + wallet.transfer(user2, address(usdt), 100); vm.prank(user1); - wallet.transfer(user2, 50); + wallet.transfer(user2, address(usdt), 50); vm.prank(user2); - wallet.transfer(user1, 50); + wallet.transfer(user1, address(usdt), 50); vm.prank(user1); Wallet.Transaction[] memory user1History = wallet.getTransactionHistory(user1); @@ -103,7 +103,7 @@ contract WalletTest is Test { vm.startPrank(user1); usdt.approve(address(wallet), largeAmount); - wallet.transfer(user2, largeAmount); + wallet.transfer(user2, address(usdt), largeAmount); Wallet.Transaction[] memory history = wallet.getTransactionHistory(user1); vm.stopPrank(); From 3c590c5fa475a7dcdc99e6edbd286af36d40b0ad Mon Sep 17 00:00:00 2001 From: Anh Hoang Nguyen Date: Wed, 18 Dec 2024 16:01:13 +0700 Subject: [PATCH 4/4] test: add test for new functions --- src/Wallet.sol | 30 +++++++++--------------------- test/Wallet.t.sol | 44 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 21 deletions(-) diff --git a/src/Wallet.sol b/src/Wallet.sol index 6fc7aa1..5238665 100644 --- a/src/Wallet.sol +++ b/src/Wallet.sol @@ -44,19 +44,16 @@ contract Wallet { function createWorldId() external onlyOwner {} - function transfer( - address _recipient, - address _token, - uint256 _amount - ) external onlyVerified(msg.sender) onlySupportedToken(_token) { + function transfer(address _recipient, address _token, uint256 _amount) + external + onlyVerified(msg.sender) + onlySupportedToken(_token) + { require(_recipient != address(0), "Zero address detected"); IERC20 token = IERC20(_token); require(token.balanceOf(msg.sender) >= _amount, "Insufficient balance"); require(_amount > 0, "Transfer amount must be greater than zero"); - require( - token.transferFrom(msg.sender, _recipient, _amount), - "Transfer failed" - ); + require(token.transferFrom(msg.sender, _recipient, _amount), "Transfer failed"); recordTransactionHistory(msg.sender, _amount, _token); } @@ -75,9 +72,7 @@ contract Wallet { // View Functions // //////////////////////////////////////////// - function getTransactionHistory( - address _user - ) external view returns (Transaction[] memory) { + function getTransactionHistory(address _user) external view returns (Transaction[] memory) { return transactions[_user]; } @@ -85,15 +80,8 @@ contract Wallet { // Private Function // ////////////////////////////////////////////// - function recordTransactionHistory( - address _user, - uint256 _amount, - address _token - ) private { - Transaction memory newTransaction = Transaction({ - amount: _amount, - token: _token - }); + function recordTransactionHistory(address _user, uint256 _amount, address _token) private { + Transaction memory newTransaction = Transaction({amount: _amount, token: _token}); transactions[_user].push(newTransaction); } diff --git a/test/Wallet.t.sol b/test/Wallet.t.sol index ad5c623..c9f0f09 100644 --- a/test/Wallet.t.sol +++ b/test/Wallet.t.sol @@ -16,7 +16,9 @@ contract WalletTest is Test { address public owner; address public user1; address public user2; + address public nonOwner; MockERC20 public usdt; + MockERC20 public anotherToken; MockWorldID public worldID; /// @notice Set up the test environment before each test @@ -24,7 +26,10 @@ contract WalletTest is Test { owner = address(this); user1 = address(0x1); user2 = address(0x2); + nonOwner = address(0x999); usdt = new MockERC20("USDT", "USDT"); + anotherToken = new MockERC20("Another Token", "ATKN"); + worldID = new MockWorldID(); factory = new WalletFactory(); @@ -112,4 +117,43 @@ contract WalletTest is Test { assertEq(history[0].amount, largeAmount, "Recorded amount does not match"); assertEq(history[0].token, address(usdt), "Recorded token address does not match"); } + + function testAddSupportedTokenByOwner() public { + vm.prank(owner); // Set the caller as the owner + wallet.addSupportedToken(address(anotherToken)); + + assertTrue(wallet.supportedTokens(address(anotherToken))); + } + + function testAddSupportedTokenByNonOwnerReverts() public { + vm.prank(nonOwner); + vm.expectRevert("not owner"); + wallet.addSupportedToken(address(anotherToken)); + } + + function testRemoveSupportedTokenByOwner() public { + vm.startPrank(owner); // Start a transaction as the owner + wallet.addSupportedToken(address(anotherToken)); + assertTrue(wallet.supportedTokens(address(anotherToken))); + + wallet.removeSupportedToken(address(anotherToken)); + assertFalse(wallet.supportedTokens(address(anotherToken))); + vm.stopPrank(); + } + + function testRemoveSupportedTokenByNonOwnerReverts() public { + vm.startPrank(owner); + wallet.addSupportedToken(address(anotherToken)); + vm.stopPrank(); + + vm.prank(nonOwner); + vm.expectRevert("not owner"); + wallet.removeSupportedToken(address(anotherToken)); + } + + function testRemoveNonSupportedTokenReverts() public { + vm.prank(owner); + vm.expectRevert("Token not supported"); + wallet.removeSupportedToken(address(anotherToken)); + } }