From a1605533278bf36b3d2f980ecff86263c4950352 Mon Sep 17 00:00:00 2001
From: telome <>
Date: Wed, 20 Dec 2023 11:46:13 +0000
Subject: [PATCH] Use EIP1167 for urn

---
 src/LockstakeEngine.sol    | 25 ++++++++++++++++++++++---
 src/LockstakeUrn.sol       |  9 +++++++--
 test/LockstakeEngine.t.sol |  2 +-
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/src/LockstakeEngine.sol b/src/LockstakeEngine.sol
index 2cff3252..a21cf24e 100644
--- a/src/LockstakeEngine.sol
+++ b/src/LockstakeEngine.sol
@@ -93,6 +93,7 @@ contract LockstakeEngine is Multicall {
     MkrNgtLike          immutable public mkrNgt;
     GemLike             immutable public ngt;
     uint256             immutable public mkrNgtRate;
+    address             immutable public urnImplementation;
 
     // --- events ---   
 
@@ -150,6 +151,7 @@ contract LockstakeEngine is Multicall {
         ngt.approve(address(mkrNgt), type(uint256).max);
         gov.approve(address(mkrNgt), type(uint256).max);
         mkrNgtRate = mkrNgt.rate();
+        urnImplementation = address(new LockstakeUrn(address(vat), stkGov_));
 
         wards[msg.sender] = 1;
         emit Rely(msg.sender);
@@ -167,6 +169,16 @@ contract LockstakeEngine is Multicall {
         ok = urnOwners[urn] == usr || urnCan[urn][usr] == 1;
     }
 
+    // See the reference implementation in https://eips.ethereum.org/EIPS/eip-1167
+    function _initCode(bytes20 _target) internal pure returns (bytes memory code) {
+        code = new bytes(55);
+        assembly {
+            mstore(add(code, 0x20), 0x3d602d80600a3d3981f3363d3d373d3d3d363d73000000000000000000000000)
+            mstore(add(code, add(0x20,0x14)), _target)
+            mstore(add(code, add(0x20,0x28)), 0x5af43d82803e903d91602b57fd5bf30000000000000000000000000000000000)
+        }
+    }
+
     // --- administration ---
 
     function rely(address usr) external auth {
@@ -200,7 +212,7 @@ contract LockstakeEngine is Multicall {
 
     function getUrn(address owner, uint256 index) external view returns (address urn) {
         uint256 salt = uint256(keccak256(abi.encode(owner, index)));
-        bytes32 codeHash = keccak256(abi.encodePacked(type(LockstakeUrn).creationCode, abi.encode(vat, stkGov)));
+        bytes32 codeHash = keccak256(abi.encodePacked(_initCode(bytes20(urnImplementation))));
         urn = address(uint160(uint256(
             keccak256(
                 abi.encodePacked(bytes1(0xff), address(this), salt, codeHash)
@@ -215,8 +227,15 @@ contract LockstakeEngine is Multicall {
 
     function open(uint256 index) external returns (address urn) {
         require(index == usrAmts[msg.sender]++, "LockstakeEngine/urn-already-opened");
-        bytes32 salt = keccak256(abi.encode(msg.sender, index));
-        urn = address(new LockstakeUrn{salt: salt}(address(vat), address(stkGov)));
+
+        uint256 salt = uint256(keccak256(abi.encode(msg.sender, index)));
+        bytes memory initCode = _initCode(bytes20(urnImplementation));
+        assembly {
+            urn := create2(0, add(initCode, 0x20), mload(initCode), salt)
+            if iszero(extcodesize(urn)) { revert(0, 0) }
+        }
+        LockstakeUrn(urn).init();
+
         urnOwners[urn] = msg.sender;
         emit Open(msg.sender, urn);
     }
diff --git a/src/LockstakeUrn.sol b/src/LockstakeUrn.sol
index 1c1ea4d4..9f3eda76 100644
--- a/src/LockstakeUrn.sol
+++ b/src/LockstakeUrn.sol
@@ -38,6 +38,7 @@ contract LockstakeUrn {
 
     address immutable public engine;
     GemLike immutable public stkGov;
+    VatLike immutable public vat;
 
     // --- modifiers ---
 
@@ -46,12 +47,16 @@ contract LockstakeUrn {
         _;
     }
 
-    // --- constructor ---
+    // --- constructor & init ---
 
     constructor(address vat_, address stkGov_) {
         engine = msg.sender;
+        vat = VatLike(vat_);
         stkGov = GemLike(stkGov_);
-        VatLike(vat_).hope(msg.sender);
+    }
+
+    function init() external isEngine {
+        vat.hope(msg.sender);
         stkGov.approve(msg.sender, type(uint256).max);
     }
 
diff --git a/test/LockstakeEngine.t.sol b/test/LockstakeEngine.t.sol
index 5ec66e30..fe09a5c2 100644
--- a/test/LockstakeEngine.t.sol
+++ b/test/LockstakeEngine.t.sol
@@ -55,7 +55,7 @@ interface CalcLike {
     function file(bytes32, uint256) external;
 }
 
-contract AllocatorVaultTest is DssTest {
+contract LockstakeEngineTest is DssTest {
     using stdStorage for StdStorage;
 
     address             public pauseProxy;