Skip to content

Commit

Permalink
adds referral code in the deposit
Browse files Browse the repository at this point in the history
  • Loading branch information
@WalidOfNow
WalidOfNow committed Jun 16, 2024
1 parent cefaf09 commit d21f0e6
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 29 deletions.
20 changes: 14 additions & 6 deletions mainnet-contracts/src/PufferL2Depositor.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ contract PufferL2Depositor is IPufferL2Depositor, AccessManaged {
* @inheritdoc IPufferL2Depositor
* @dev Restricted in this context is like `whenNotPaused` modifier from Pausable.sol
*/
function deposit(address token, address account, Permit calldata permitData)
function deposit(address token, address account, Permit calldata permitData, uint64 referralCode)
external
onlySupportedTokens(token)
restricted
Expand All @@ -62,17 +62,23 @@ contract PufferL2Depositor is IPufferL2Depositor, AccessManaged {

IERC20(token).safeTransferFrom(msg.sender, address(this), permitData.amount);

_deposit({ token: token, depositor: msg.sender, account: account, amount: permitData.amount });
_deposit({
token: token,
depositor: msg.sender,
account: account,
amount: permitData.amount,
referralCode: referralCode
});
}

/**
* @inheritdoc IPufferL2Depositor
* @dev Restricted in this context is like `whenNotPaused` modifier from Pausable.sol
*/
function depositETH(address account) external payable restricted {
function depositETH(address account, uint64 referralCode) external payable restricted {
IWETH(WETH).deposit{ value: msg.value }();

_deposit({ token: WETH, depositor: msg.sender, account: account, amount: msg.value });
_deposit({ token: WETH, depositor: msg.sender, account: account, amount: msg.value, referralCode: referralCode });
}

/**
Expand Down Expand Up @@ -111,14 +117,16 @@ contract PufferL2Depositor is IPufferL2Depositor, AccessManaged {
*/
function revertIfPaused() external restricted { }

function _deposit(address token, address depositor, address account, uint256 amount) internal {
function _deposit(address token, address depositor, address account, uint256 amount, uint64 referralCode)
internal
{
PufToken pufToken = PufToken(tokens[token]);

IERC20(token).safeIncreaseAllowance(address(pufToken), amount);

pufToken.deposit(depositor, account, amount);

emit DepositedToken(token, msg.sender, account, amount);
emit DepositedToken(token, msg.sender, account, amount, referralCode);
}

function _addNewToken(address token) internal {
Expand Down
10 changes: 7 additions & 3 deletions mainnet-contracts/src/interface/IPufferL2Depositor.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,11 @@ interface IPufferL2Depositor {
* @dev Signature "0x1a0b42192bd87f901af5da67a080a510d8c86ccd976904272a9b78c11e7fe085"
*/
event DepositedToken(
address indexed token, address indexed depositor, address indexed account, uint256 tokenAmount
address indexed token,
address indexed depositor,
address indexed account,
uint256 tokenAmount,
uint64 referralCode
);

event DepositCapUpdated(address indexed token, uint256 oldDepositCap, uint256 newDepositCap);
Expand All @@ -50,14 +54,14 @@ interface IPufferL2Depositor {
*
* @dev Restricted in this context is like `whenNotPaused` modifier from Pausable.sol
*/
function deposit(address token, address account, Permit calldata permitData) external;
function deposit(address token, address account, Permit calldata permitData, uint64 referralCode) external;

/**
* @notice Deposits naative ETH by wrapping it into WETH and then depositing to corresponding token contract
*
* @dev Restricted in this context is like `whenNotPaused` modifier from Pausable.sol
*/
function depositETH(address account) external payable;
function depositETH(address account, uint64 referralCode) external payable;

/**
* @notice Called by the Token contracts to check if the system is paused
Expand Down
42 changes: 22 additions & 20 deletions mainnet-contracts/test/unit/PufferL2Staking.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ contract PufferL2Staking is UnitTestHelper {

address bob = makeAddr("bob");

uint64 referralCode = 0;

function setUp() public override {
super.setUp();

Expand Down Expand Up @@ -122,7 +124,7 @@ contract PufferL2Staking is UnitTestHelper {
}

// Bad permit signature + approve
function test_depositFor_dai_approve(uint32 amount) public {
function test_depositFor_dai_approve(uint32 amount, uint64 refCode) public {
vm.assume(amount > 0);

// This is a bad permit signature
Expand All @@ -136,8 +138,8 @@ contract PufferL2Staking is UnitTestHelper {
dai.approve(address(depositor), amount);

vm.expectEmit(true, true, true, true);
emit IPufferL2Depositor.DepositedToken(address(dai), bob, bob, amount);
depositor.deposit(address(dai), bob, permit);
emit IPufferL2Depositor.DepositedToken(address(dai), bob, bob, amount, refCode);
depositor.deposit(address(dai), bob, permit, refCode);

PufToken pufToken = PufToken(depositor.tokens(address(dai)));
assertEq(pufToken.balanceOf(bob), amount, "bob got pufToken");
Expand All @@ -157,8 +159,8 @@ contract PufferL2Staking is UnitTestHelper {
sixDecimal.approve(address(depositor), amount);

vm.expectEmit(true, true, true, true);
emit IPufferL2Depositor.DepositedToken(address(sixDecimal), bob, bob, amount);
depositor.deposit(address(sixDecimal), bob, permit);
emit IPufferL2Depositor.DepositedToken(address(sixDecimal), bob, bob, amount, referralCode);
depositor.deposit(address(sixDecimal), bob, permit, referralCode);

PufToken pufToken = PufToken(depositor.tokens(address(sixDecimal)));

Expand All @@ -184,8 +186,8 @@ contract PufferL2Staking is UnitTestHelper {
twentyTwoDecimal.approve(address(depositor), amount);

vm.expectEmit(true, true, true, true);
emit IPufferL2Depositor.DepositedToken(address(twentyTwoDecimal), bob, bob, amount);
depositor.deposit(address(twentyTwoDecimal), bob, permit);
emit IPufferL2Depositor.DepositedToken(address(twentyTwoDecimal), bob, bob, amount, referralCode);
depositor.deposit(address(twentyTwoDecimal), bob, permit, referralCode);

PufToken pufToken = PufToken(depositor.tokens(address(twentyTwoDecimal)));

Expand All @@ -198,7 +200,7 @@ contract PufferL2Staking is UnitTestHelper {
}

// Good Permit signature signature
function test_depositFor_dai_permit(uint32 amount) public {
function test_depositFor_dai_permit(uint32 amount, uint64 refCode) public {
vm.assume(amount > 0);

// Good permit signature
Expand All @@ -211,15 +213,15 @@ contract PufferL2Staking is UnitTestHelper {
dai.approve(depositor.tokens(address(dai)), amount);

vm.expectEmit(true, true, true, true);
emit IPufferL2Depositor.DepositedToken(address(dai), bob, bob, amount);
depositor.deposit(address(dai), bob, permit);
emit IPufferL2Depositor.DepositedToken(address(dai), bob, bob, amount, refCode);
depositor.deposit(address(dai), bob, permit, refCode);

PufToken pufToken = PufToken(depositor.tokens(address(dai)));
assertEq(pufToken.balanceOf(bob), amount, "bob got pufToken");
}

// Weth doesn't have `permit` at all
function test_deposiFor_WETH(uint32 amount) public {
function test_deposiFor_WETH(uint32 amount, uint64 refCode) public {
vm.assume(amount > 0);

// WETH Doesn't have permit
Expand All @@ -235,24 +237,24 @@ contract PufferL2Staking is UnitTestHelper {

// weth.permit triggers weth.fallback() and it doesn't revert
vm.expectEmit(true, true, true, true);
emit IPufferL2Depositor.DepositedToken(address(weth), bob, bob, amount);
depositor.deposit(address(weth), bob, permit);
emit IPufferL2Depositor.DepositedToken(address(weth), bob, bob, amount, refCode);
depositor.deposit(address(weth), bob, permit, refCode);

PufToken pufToken = PufToken(depositor.tokens(address(weth)));
assertEq(pufToken.balanceOf(bob), amount, "bob got pufToken");
}

// ETH deposit & weth withdrawal
function test_depositFor_ETH_withdraw_weth(uint16 amount) public {
function test_depositFor_ETH_withdraw_weth(uint16 amount, uint64 refCode) public {
vm.assume(amount > 0);

vm.deal(bob, amount);

vm.startPrank(bob);

vm.expectEmit(true, true, true, true);
emit IPufferL2Depositor.DepositedToken(address(weth), bob, bob, amount);
depositor.depositETH{ value: amount }(bob);
emit IPufferL2Depositor.DepositedToken(address(weth), bob, bob, amount, refCode);
depositor.depositETH{ value: amount }(bob, refCode);

PufToken pufToken = PufToken(depositor.tokens(address(weth)));
assertEq(pufToken.balanceOf(bob), amount, "bob got pufToken");
Expand Down Expand Up @@ -329,7 +331,7 @@ contract PufferL2Staking is UnitTestHelper {

vm.startPrank(bob);
vm.expectRevert();
depositor.deposit(address(notSupportedToken), bob, permit);
depositor.deposit(address(notSupportedToken), bob, permit, referralCode);
}

// zero address token reverts
Expand Down Expand Up @@ -361,7 +363,7 @@ contract PufferL2Staking is UnitTestHelper {
vm.startPrank(bob);

vm.expectRevert(IPufStakingPool.InvalidAmount.selector);
depositor.depositETH{ value: 0 }(bob);
depositor.depositETH{ value: 0 }(bob, referralCode);
}

// Mock address 123 is not allowed to be migrator
Expand All @@ -377,14 +379,14 @@ contract PufferL2Staking is UnitTestHelper {
vm.deal(bob, 1 ether);
vm.startPrank(bob);
vm.expectRevert();
depositor.depositETH{ value: 1 ether }(address(0));
depositor.depositETH{ value: 1 ether }(address(0), referralCode);
}

// 0 deposit eth reverts
function testRevert_zero_deposit_ETH() public {
vm.startPrank(bob);
vm.expectRevert();
depositor.depositETH{ value: 0 }(bob);
depositor.depositETH{ value: 0 }(bob, referralCode);
}

// No deposit reverts
Expand Down

0 comments on commit d21f0e6

Please sign in to comment.