From e9b64d187bf297ad85e422b2da839e2fe408dd62 Mon Sep 17 00:00:00 2001
From: Carter Carlson <pegahcarter@gmail.com>
Date: Tue, 10 Dec 2024 19:39:07 +0100
Subject: [PATCH] feat: `setSendReceiveLib()` (#27)

* feat: ProxyAdmin => FraxProxyAdmin

* feat: setLibs
---
 contracts/FraxProxyAdmin.sol                  |  8 ++
 scripts/BaseL0Script.sol                      |  3 +-
 .../DeployFraxOFTProtocol.s.sol               | 93 +++++++++++++++----
 3 files changed, 86 insertions(+), 18 deletions(-)
 create mode 100644 contracts/FraxProxyAdmin.sol

diff --git a/contracts/FraxProxyAdmin.sol b/contracts/FraxProxyAdmin.sol
new file mode 100644
index 0000000..d4b75eb
--- /dev/null
+++ b/contracts/FraxProxyAdmin.sol
@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.22;
+
+import { ProxyAdmin } from "@fraxfinance/layerzero-v2-upgradeable/messagelib/contracts/upgradeable/proxy/ProxyAdmin.sol";
+
+contract FraxProxyAdmin is ProxyAdmin {
+    constructor(address _admin) ProxyAdmin(_admin) {}
+}
\ No newline at end of file
diff --git a/scripts/BaseL0Script.sol b/scripts/BaseL0Script.sol
index aadfa4d..9dc1acc 100644
--- a/scripts/BaseL0Script.sol
+++ b/scripts/BaseL0Script.sol
@@ -8,6 +8,7 @@ import { console } from "frax-std/FraxTest.sol";
 
 import { SerializedTx, SafeTxUtil } from "scripts/SafeBatchSerialize.sol";
 import { FraxOFTUpgradeable } from "contracts/FraxOFTUpgradeable.sol";
+import { FraxProxyAdmin } from "contracts/FraxProxyAdmin.sol";
 import { ImplementationMock } from "contracts/mocks/ImplementationMock.sol";
 
 import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
@@ -97,7 +98,7 @@ contract BaseL0Script is Script {
     string public json;
 
     function version() public virtual pure returns (uint256, uint256, uint256) {
-        return (1, 2, 0);
+        return (1, 2, 1);
     }
 
     modifier broadcastAs(uint256 privateKey) {
diff --git a/scripts/DeployFraxOFTProtocol/DeployFraxOFTProtocol.s.sol b/scripts/DeployFraxOFTProtocol/DeployFraxOFTProtocol.s.sol
index 2490115..3e1310d 100644
--- a/scripts/DeployFraxOFTProtocol/DeployFraxOFTProtocol.s.sol
+++ b/scripts/DeployFraxOFTProtocol/DeployFraxOFTProtocol.s.sol
@@ -17,7 +17,7 @@ contract DeployFraxOFTProtocol is BaseL0Script {
     using Strings for uint256;
 
     function version() public virtual override pure returns (uint256, uint256, uint256) {
-        return (1, 2, 0);
+        return (1, 2, 1);
     }
 
     function setUp() public virtual override {
@@ -41,15 +41,6 @@ contract DeployFraxOFTProtocol is BaseL0Script {
         setupProxyDestinations();
     }
 
-    function setupLegacyDestinations() public virtual {
-        for (uint256 i=0; i<legacyConfigs.length; i++) {
-            setupDestination({
-                _connectedConfig: legacyConfigs[i],
-                _connectedOfts: legacyOfts
-            });
-        }
-    }
-
     function setupProxyDestinations() public virtual {
         for (uint256 i=0; i<proxyConfigs.length; i++) {
             // skip if destination == source
@@ -61,6 +52,10 @@ contract DeployFraxOFTProtocol is BaseL0Script {
         }
     }
 
+    function setupLegacyDestinations() public pure {
+        // DEPRECATED
+    }
+
     function setupDestination(
         L0Config memory _connectedConfig,
         address[] memory _connectedOfts
@@ -88,12 +83,19 @@ contract DeployFraxOFTProtocol is BaseL0Script {
         setupEvms();
         setupNonEvms();
 
+        /// @dev configures legacy configs as well
         setDVNs({
             _connectedConfig: broadcastConfig,
             _connectedOfts: proxyOfts,
             _configs: allConfigs
         });
 
+        setLibs({
+            _connectedConfig: broadcastConfig,
+            _connectedOfts: proxyOfts,
+            _configs: allConfigs
+        });
+
         setPriviledgedRoles();
     }
 
@@ -103,12 +105,6 @@ contract DeployFraxOFTProtocol is BaseL0Script {
             _configs: evmConfigs
         });
 
-        /// @dev legacy, non-upgradeable OFTs
-        setEvmPeers({
-            _connectedOfts: proxyOfts,
-            _peerOfts: legacyOfts,
-            _configs: legacyConfigs
-        });
         /// @dev Upgradeable OFTs maintaining the same address cross-chain.
         setEvmPeers({
             _connectedOfts: proxyOfts,
@@ -160,7 +156,7 @@ contract DeployFraxOFTProtocol is BaseL0Script {
     function deployFraxOFTUpgradeablesAndProxies() broadcastAs(oftDeployerPK) public virtual {
 
         // Proxy admin (0x223a681fc5c5522c85C96157c0efA18cd6c5405c)
-        proxyAdmin = address(new ProxyAdmin(vm.addr(configDeployerPK)));
+        proxyAdmin = address(new FraxProxyAdmin(vm.addr(configDeployerPK)));
 
         // Implementation mock (0x8f1B9c1fd67136D525E14D96Efb3887a33f16250)
         implementationMock = address(new ImplementationMock());
@@ -503,4 +499,67 @@ contract DeployFraxOFTProtocol is BaseL0Script {
             );
         }
     }
+
+    function setLibs(
+        L0Config memory _connectedConfig,
+        address[] memory _connectedOfts,
+        L0Config[] memory _configs 
+    ) public virtual {
+        // for each oft
+        for (uint256 o=0; o < _connectedOfts.length; o++) {
+            // for each destination
+            for (uint256 c=0; c < _configs.length; c++) {
+                setLib({
+                    _connectedConfig: _connectedConfig,
+                    _connectedOft: _connectedOfts[o],
+                    _config: _configs[c]
+                });
+            }
+        }
+    }
+
+    function setLib(
+        L0Config memory _connectedConfig,
+        address _connectedOft,
+        L0Config memory _config
+    ) public virtual {
+        // skip if the connected and target are the same
+        if (_connectedConfig.eid == _config.eid) return;
+
+        // set sendLib to default
+        bytes memory data = abi.encodeCall(
+            IMessageLibManager.setSendLibrary,
+            (
+                _connectedOft, uint32(_config.eid), _connectedConfig.sendLib302
+            )
+        );
+        (bool success,) = _connectedConfig.endpoint.call(data);
+        require(success, "Unable to call setSendLibrary");
+        serializedTxs.push(
+            SerializedTx({
+                name: "setSendLibrary",
+                to: _connectedConfig.endpoint,
+                value: 0,
+                data: data
+            })
+        );
+
+        // set receiveLib to default
+        data = abi.encodeCall(
+            IMessageLibManager.setReceiveLibrary,
+            (
+                _connectedOft, uint32(_config.eid), _connectedConfig.receiveLib302, 0
+            )
+        );
+        (success,) = _connectedConfig.endpoint.call(data);
+        require(success, "Unable to call setReceiveLibrary");
+        serializedTxs.push(
+            SerializedTx({
+                name: "setReceiveLibrary",
+                to: _connectedConfig.endpoint,
+                value: 0,
+                data: data
+            })
+        );
+    }
 }