Add safe address and fee handling to AaveWithdrawWorkflow, update tes…

…ts and addresses.json accordingly.
KintoXYZ · Dec 24, 2024 · a095fea · a095fea
1 parent 1a4bd8e
commit a095fea
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 10 deletions.
diff --git a/script/migrations/145-upgrade_access_protocol.s.sol b/script/migrations/145-upgrade_access_protocol.s.sol
@@ -48,8 +48,9 @@ contract DeployScript is Script, MigrationHelper {
         registry.allowWorkflow(address(aaveRepayWorkflow));
 
         registry.disallowWorkflow(_getChainDeployment("AaveWithdrawWorkflow"));
-        AaveWithdrawWorkflow aaveWithdrawWorkflow =
-            new AaveWithdrawWorkflow(getAavePoolProvider(), _getChainDeployment("Bridger"));
+        AaveWithdrawWorkflow aaveWithdrawWorkflow = new AaveWithdrawWorkflow(
+            getAavePoolProvider(), _getChainDeployment("Bridger"), getMamoriSafeByChainId(block.chainid)
+        );
         saveContractAddress("AaveWithdrawWorkflow", address(aaveWithdrawWorkflow));
         registry.allowWorkflow(address(aaveWithdrawWorkflow));
 

diff --git a/src/access/workflows/AaveWithdrawWorkflow.sol b/src/access/workflows/AaveWithdrawWorkflow.sol
@@ -23,16 +23,21 @@ contract AaveWithdrawWorkflow {
     IPoolAddressesProvider public immutable poolAddressProvider;
     /// @notice Address of the Bridger contract
     IBridger public immutable bridger;
+    /// @notice Address of the Safe contract
+    address public immutable safe;
+    /// @notice Fee charged upon withdrawal. 10bps.
+    uint256 public constant FEE = 1e15;
 
     /* ============ Constructor ============ */
 
     /**
      * @notice Initializes the contract with Aave's pool address provider
      * @param poolAddressProvider_ The address of Aave's pool address provider
      */
-    constructor(address poolAddressProvider_, address bridger_) {
+    constructor(address poolAddressProvider_, address bridger_, address safe_) {
         poolAddressProvider = IPoolAddressesProvider(poolAddressProvider_);
         bridger = IBridger(bridger_);
+        safe = safe_;
     }
 
     /* ============ External Functions ============ */
@@ -42,7 +47,7 @@ contract AaveWithdrawWorkflow {
      * @param asset The address of the asset to withdraw
      * @param amount The amount to withdraw (use type(uint256).max for max available)
      */
-    function withdraw(address asset, uint256 amount) public {
+    function withdraw(address asset, uint256 amount) public returns (uint256) {
         address pool = poolAddressProvider.getPool();
 
         // If amount is max uint256, withdraw all available
@@ -52,6 +57,10 @@ contract AaveWithdrawWorkflow {
 
         // Withdraw from Aave
         IAavePool(pool).withdraw(asset, amount, address(this));
+        // Send the fee to the Safe
+        uint256 fee = amount * FEE / 1e18;
+        IERC20(asset).transfer(safe, fee);
+        return amount - fee;
     }
 
     function withdrawAndBridge(
@@ -60,7 +69,7 @@ contract AaveWithdrawWorkflow {
         address kintoWallet,
         IBridger.BridgeData calldata bridgeData
     ) external payable returns (uint256 amountOut) {
-        withdraw(asset, amount);
+        amount = withdraw(asset, amount);
 
         // Approve max allowance to save on gas for future transfers
         if (IERC20(asset).allowance(address(this), address(bridger)) < amount) {

diff --git a/test/artifacts/42161/addresses.json b/test/artifacts/42161/addresses.json
@@ -29,7 +29,7 @@
   "Viewer": "0x8888886e1d7c1468d7300cF08db89FFE68F29830",
   "Viewer-impl": "0x80338A3f75614491c8DC383fFaA663b9a27CD05d",
   "WethWorkflow": "0x7F7c594eE170a62d7e7615972831038Cf7d4Fc1A",
-  "WithdrawWorkflow": "0xbc22c860C1ED7330271eeF19FB47Eb08548f1723",
+  "WithdrawWorkflow": "0x794E1908A1D41760B8E2b798134c9856E24dCe65",
   "AaveLendWorkflow": "0xB47Ed636c8296729E81463109FEbf833CeEa71fb",
   "AaveRepayWorkflow": "0x24f71379C39b515Ff5182F4b0cc298793EC5998c",
   "AaveWithdrawWorkflow": "0xef4D6687372172c4af1802C208Ab40673b014309",

diff --git a/test/fork/workflows/AaveWithdrawWorkflow.t.sol b/test/fork/workflows/AaveWithdrawWorkflow.t.sol
@@ -34,6 +34,8 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader,
     IAccessPoint internal accessPoint;
     AaveWithdrawWorkflow internal aaveWithdrawWorkflow;
     IAavePool internal aavePool;
+    address internal safe = address(0x5afe);
+    uint256 constant FEE = 1e15;
 
     function setUp() public override {
         super.setUp();
@@ -50,7 +52,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader,
         accessPoint = accessRegistry.deployFor(address(alice0));
         vm.label(address(accessPoint), "accessPoint");
 
-        aaveWithdrawWorkflow = new AaveWithdrawWorkflow(ARB_AAVE_POOL_PROVIDER, address(bridger));
+        aaveWithdrawWorkflow = new AaveWithdrawWorkflow(ARB_AAVE_POOL_PROVIDER, address(bridger), safe);
         vm.label(address(aaveWithdrawWorkflow), "aaveWithdrawWorkflow");
 
         vm.prank(accessRegistry.owner());
@@ -89,7 +91,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader,
         // Assert balances changed correctly
         assertEq(
             IERC20(assetToWithdraw).balanceOf(address(accessPoint)),
-            initialAccessPointBalance + amountToWithdraw,
+            initialAccessPointBalance + amountToWithdraw - amountToWithdraw * FEE / 1e18,
             "Invalid USDC balance"
         );
         assertEq(
@@ -126,7 +128,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader,
         // Assert balances changed correctly
         assertEq(
             IERC20(assetToWithdraw).balanceOf(address(accessPoint)),
-            initialAccessPointBalance + amountToSupply,
+            initialAccessPointBalance + amountToSupply - amountToSupply * FEE / 1e18,
             "Invalid USDC balance"
         );
         assertEq(IERC20(aToken).balanceOf(address(accessPoint)), 0, "Invalid aToken balance");
@@ -175,7 +177,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader,
         assertEq(IERC20(assetToWithdraw).balanceOf(address(bridger)), initialBridgerBalance, "Invalid bridger balance");
         assertEq(
             IERC20(assetToWithdraw).balanceOf(address(bridgeData.vault)),
-            initialVaultBalance + amountToWithdraw,
+            initialVaultBalance + amountToWithdraw - amountToWithdraw * FEE / 1e18,
             "Invalid vault balance"
         );
     }