Skip to content

Commit

Permalink
feat: scaffold out attach logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@colinnielsen
colinnielsen committed Nov 7, 2023
1 parent 085ac6b commit 092df5e
Show file tree
Hide file tree
Showing 2 changed files with 161 additions and 73 deletions.
160 changes: 87 additions & 73 deletions src/SafeTestTools.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,80 +2,11 @@
pragma solidity >=0.7.0 <0.9.0;

import "forge-std/Test.sol";
import "solady/utils/LibSort.sol";
import "safe-contracts/Safe.sol";
import "safe-contracts/proxies/SafeProxyFactory.sol";
import "safe-contracts/libraries/SignMessageLib.sol";
import "./CompatibilityFallbackHandler_1_4_1.sol";

address constant VM_ADDR = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
bytes12 constant ADDR_MASK = 0xffffffffffffffffffffffff;

function getAddr(uint256 pk) pure returns (address) {
return Vm(VM_ADDR).addr(pk);
}

function encodeSmartContractWalletAsPK(address addr) pure returns (uint256 encodedPK) {
assembly {
let addr_b32 := addr
encodedPK := or(addr, ADDR_MASK)
}
}

function decodeSmartContractWalletAsAddress(uint256 pk) pure returns (address decodedAddr) {
assembly {
let addr := shl(96, pk)
decodedAddr := shr(96, addr)
}
}

function isSmartContractPK(uint256 pk) pure returns (bool isEncoded) {
assembly {
isEncoded := eq(shr(160, pk), shr(160, ADDR_MASK))
}
}

library Sort {
function sort(address[] memory arr) public pure returns (address[] memory) {
LibSort.sort(arr);
return arr;
}
}

function sortPKsByComputedAddress(uint256[] memory _pks) pure returns (uint256[] memory) {
uint256[] memory sortedPKs = new uint256[](_pks.length);

address[] memory addresses = new address[](_pks.length);
bytes32[2][] memory accounts = new bytes32[2][](_pks.length);

for (uint256 i; i < _pks.length; i++) {
address signer = getAddr(_pks[i]);
addresses[i] = signer;
accounts[i][0] = bytes32(abi.encode(signer));
accounts[i][1] = bytes32(_pks[i]);
}

addresses = Sort.sort(addresses);

uint256 found;
for (uint256 j; j < addresses.length; j++) {
address signer = addresses[j];
uint256 pk;
for (uint256 k; k < accounts.length; k++) {
if (address(uint160(uint256(accounts[k][0]))) == signer) {
pk = uint256(accounts[k][1]);
found++;
}
}

sortedPKs[j] = pk;
}

if (found < _pks.length) {
revert("SAFETESTTOOLS: issue with private key sorting, please open a ticket on github");
}
return sortedPKs;
}
import "./Utils.sol";

// collapsed interface that includes comapatibilityfallback handler calls
abstract contract DeployedSafe is Safe, CompatibilityFallbackHandler {}
Expand All @@ -91,7 +22,13 @@ struct AdvancedSafeInitParams {
bytes initData;
}

enum InstanceType {
Live,
Test
}

struct SafeInstance {
InstanceType instanceType;
uint256 instanceId;
uint256[] ownerPKs;
address[] owners;
Expand All @@ -100,7 +37,7 @@ struct SafeInstance {
}

library SafeTestLib {
function execTransaction(
function _execTxWithLocalPKs(
SafeInstance memory instance,
address to,
uint256 value,
Expand Down Expand Up @@ -164,6 +101,66 @@ library SafeTestLib {
});
}

function _spoofSigWithStorageOverride(
SafeInstance memory instance,
address to,
uint256 value,
bytes memory data,
Enum.Operation operation,
uint256 safeTxGas,
uint256 baseGas,
uint256 gasPrice,
address gasToken,
address refundReceiver,
bytes memory signatures
) internal returns (bool) {
return true;
}

function execTransaction(
SafeInstance memory instance,
address to,
uint256 value,
bytes memory data,
Enum.Operation operation,
uint256 safeTxGas,
uint256 baseGas,
uint256 gasPrice,
address gasToken,
address refundReceiver,
bytes memory signatures
) internal returns (bool) {
if (instance.instanceType == InstanceType.Test) {
return _execTxWithLocalPKs(
instance,
to,
value,
data,
operation,
safeTxGas,
baseGas,
gasPrice,
gasToken,
refundReceiver,
signatures
);
} else {
return _spoofSigWithStorageOverride(
instance,
to,
value,
data,
operation,
safeTxGas,
baseGas,
gasPrice,
gasToken,
refundReceiver,
signatures
);
}
}

function execTransaction(
SafeInstance memory instance,
address to,
Expand Down Expand Up @@ -282,7 +279,7 @@ library SafeTestLib {
(v, r, s) = Vm(VM_ADDR).sign(pk, txDataHash);
}

function incrementNonce(SafeInstance memory instance) public returns (uint256 newNonce) {
function incrementNonce(SafeInstance memory instance) public returns (uint256) {
execTransaction(instance, address(0), 0, "", Enum.Operation.Call, 0, 0, 0, address(0), address(0), "");
return instance.safe.nonce();
}
Expand All @@ -306,6 +303,23 @@ contract SafeTestTools {
handler = new CompatibilityFallbackHandler();
}

function _attachToSafe(address _safe) public returns (SafeInstance memory) {
DeployedSafe safe = DeployedSafe(payable(_safe));

uint256[] memory ownerPKs = new uint256[](0);

SafeInstance memory instance0 = SafeInstance({
instanceType: InstanceType.Live,
instanceId: instances.length,
ownerPKs: ownerPKs,
owners: safe.getOwners(),
threshold: safe.getThreshold(),
safe: safe
});

instances.push(instance0);
}

function _setupSafe(
uint256[] memory ownerPKs,
uint256 threshold,
Expand Down Expand Up @@ -351,11 +365,11 @@ contract SafeTestTools {
);

SafeInstance memory instance0 = SafeInstance({
instanceType: InstanceType.Test,
instanceId: instances.length,
ownerPKs: sortedPKs,
owners: owners,
threshold: threshold,
// setup safe ecosystem, singleton, proxy factory, fallback handler, and create a new safe
safe: safe0
});
instances.push(instance0);
Expand Down
74 changes: 74 additions & 0 deletions src/Utils.sol
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
// SPDX-License-Identifier: MIT
pragma solidity >=0.7.0 <0.9.0;

import "forge-std/Test.sol";
import "solady/utils/LibSort.sol";

address constant VM_ADDR = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
bytes12 constant ADDR_MASK = 0xffffffffffffffffffffffff;

function getAddr(uint256 pk) pure returns (address) {
return Vm(VM_ADDR).addr(pk);
}

function encodeSmartContractWalletAsPK(address addr) pure returns (uint256 encodedPK) {
assembly {
let addr_b32 := addr
encodedPK := or(addr, ADDR_MASK)
}
}

function decodeSmartContractWalletAsAddress(uint256 pk) pure returns (address decodedAddr) {
assembly {
let addr := shl(96, pk)
decodedAddr := shr(96, addr)
}
}

function isSmartContractPK(uint256 pk) pure returns (bool isEncoded) {
assembly {
isEncoded := eq(shr(160, pk), shr(160, ADDR_MASK))
}
}

library Sort {
function sort(address[] memory arr) public pure returns (address[] memory) {
LibSort.sort(arr);
return arr;
}
}

function sortPKsByComputedAddress(uint256[] memory _pks) pure returns (uint256[] memory) {
uint256[] memory sortedPKs = new uint256[](_pks.length);

address[] memory addresses = new address[](_pks.length);
bytes32[2][] memory accounts = new bytes32[2][](_pks.length);

for (uint256 i; i < _pks.length; i++) {
address signer = getAddr(_pks[i]);
addresses[i] = signer;
accounts[i][0] = bytes32(abi.encode(signer));
accounts[i][1] = bytes32(_pks[i]);
}

addresses = Sort.sort(addresses);

uint256 found;
for (uint256 j; j < addresses.length; j++) {
address signer = addresses[j];
uint256 pk;
for (uint256 k; k < accounts.length; k++) {
if (address(uint160(uint256(accounts[k][0]))) == signer) {
pk = uint256(accounts[k][1]);
found++;
}
}

sortedPKs[j] = pk;
}

if (found < _pks.length) {
revert("SAFETESTTOOLS: issue with private key sorting, please open a ticket on github");
}
return sortedPKs;
}

0 comments on commit 092df5e

Please sign in to comment.