Skip to content

N
nebula
Commit ab4a40cc authored by Maurelian's avatar Maurelian
Browse files
Options

safe-tools: Refactor to move free functions into SafeTestLib

parent 3787e706
Hide whitespace changes
Inline Side-by-side
Showing with 116 additions and 111 deletions
packages/contracts-bedrock/test/LivenessGuard.t.sol
View file @ ab4a40cc
......@@ -137,7 +137,7 @@ contract LivenessGuard_OwnerManagement_Test is LivenessGuard_TestInit {
safeInstance.execTransaction({
to: address(safeInstance.safe),
value: 0,
data: abi.encodeWithSelector(OwnerManager.removeOwner.selector, SENTINEL_OWNERS, ownerToRemove, 1)
data: abi.encodeWithSelector(OwnerManager.removeOwner.selector, SafeTestLib.SENTINEL_OWNERS, ownerToRemove, 1)
});
assertFalse(safeInstance.safe.isOwner(ownerToRemove));
......
packages/contracts-bedrock/test/LivenessModule.t.sol
View file @ ab4a40cc
......@@ -42,7 +42,7 @@ contract LivenessModule_TestInit is Test, SafeTestTools {
vm.warp(initTime);
// Create a Safe with 10 owners
(, uint256[] memory keys) = makeAddrsAndKeys(10);
(, uint256[] memory keys) = SafeTestLib.makeAddrsAndKeys(10);
safeInstance = _setupSafe(keys, 8);
livenessGuard = new LivenessGuard(safeInstance.safe);
......
packages/contracts-bedrock/test/SafeSigners.t.sol
View file @ ab4a40cc
......@@ -39,10 +39,11 @@ contract SafeSigners_Test is Test, SafeTestTools {
// Limit the number of signatures to 25
uint256 numSigs = bound(_numSigs, 1, 25);
(, uint256[] memory keys) = makeAddrsAndKeys(numSigs);
(, uint256[] memory keys) = SafeTestLib.makeAddrsAndKeys(numSigs);
for (uint256 i = 0; i < keys.length; i++) {
if (sigType(keys[i]) == SigTypes.Contract) {
keys[i] = encodeSmartContractWalletAsPK(decodeSmartContractWalletAsAddress(keys[i]));
keys[i] =
SafeTestLib.encodeSmartContractWalletAsPK(SafeTestLib.decodeSmartContractWalletAsAddress(keys[i]));
}
}
......@@ -66,15 +67,15 @@ contract SafeSigners_Test is Test, SafeTestTools {
v += 4;
signatures = bytes.concat(signatures, abi.encodePacked(r, s, v));
} else if (sigType(pks[i]) == SigTypes.ApprovedHash) {
vm.prank(getAddr(pks[i]));
vm.prank(SafeTestLib.getAddr(pks[i]));
safeInstance.safe.approveHash(digest);
v = 1;
// s is not checked on approved hash signatures, so we can leave it as zero.
r = bytes32(uint256(uint160(getAddr(pks[i]))));
r = bytes32(uint256(uint160(SafeTestLib.getAddr(pks[i]))));
signatures = bytes.concat(signatures, abi.encodePacked(r, s, v));
} else if (sigType(pks[i]) == SigTypes.Contract) {
contractSigs++;
address addr = decodeSmartContractWalletAsAddress(pks[i]);
address addr = SafeTestLib.decodeSmartContractWalletAsAddress(pks[i]);
r = bytes32(uint256(uint160(addr)));
vm.mockCall(
addr, abi.encodeWithSignature("isValidSignature(bytes,bytes)"), abi.encode(EIP1271_MAGIC_VALUE)
......
packages/contracts-bedrock/test/safe-tools/SafeTestTools.sol
View file @ ab4a40cc
......@@ -11,104 +11,7 @@ import "./CompatibilityFallbackHandler_1_3_0.sol";
// Tools to simplify testing Safe contracts
// Author: Colin Nielsen (https://github.com/colinnielsen/safe-tools)
address constant VM_ADDR = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
bytes12 constant ADDR_MASK = 0xffffffffffffffffffffffff;
/// @dev The address of the first owner in the linked list of owners
address constant SENTINEL_OWNERS = address(0x1);
/// @dev Get the address from a private key
function getAddr(uint256 pk) pure returns (address) {
return Vm(VM_ADDR).addr(pk);
}
/// @dev Get arrays of addresses and private keys. The arrays are sorted by address, and the addresses are labelled
function makeAddrsAndKeys(uint256 num) returns (address[] memory addrs, uint256[] memory keys) {
keys = new uint256[](num);
addrs = new address[](num);
for (uint256 i; i < num; i++) {
uint256 key = uint256(keccak256(abi.encodePacked(i)));
keys[i] = key;
}
for (uint256 i; i < num; i++) {
addrs[i] = Vm(VM_ADDR).addr(keys[i]);
Vm(VM_ADDR).label(getAddr(keys[i]), string.concat("SAFETEST: Signer ", string(abi.encodePacked(bytes32(i)))));
}
}
/// @dev Encode a smart contract wallet as a private key
function encodeSmartContractWalletAsPK(address addr) pure returns (uint256 encodedPK) {
assembly {
let addr_b32 := addr
encodedPK := or(addr, ADDR_MASK)
}
}
/// @dev Decode a smart contract wallet as an address from a private key
function decodeSmartContractWalletAsAddress(uint256 pk) pure returns (address decodedAddr) {
assembly {
let addr := shl(96, pk)
decodedAddr := shr(96, addr)
}
}
/// @dev Checks if a private key is an encoded smart contract address
function isSmartContractPK(uint256 pk) pure returns (bool isEncoded) {
assembly {
isEncoded := eq(shr(160, pk), shr(160, ADDR_MASK))
}
}
library Sort {
/// @dev Sorts an array of addresses in place
function sort(address[] memory arr) public pure returns (address[] memory) {
LibSort.sort(arr);
return arr;
}
}
/// @dev Sorts an array of private keys by the computed address
/// If the private key is a smart contract wallet, it will be decoded and sorted by the address
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++) {
uint256 pk = _pks[i];
address signer = getAddr(pk);
if (isSmartContractPK(pk)) {
signer = decodeSmartContractWalletAsAddress(pk);
}
addresses[i] = signer;
accounts[i][0] = bytes32(abi.encode(signer));
accounts[i][1] = bytes32(pk);
}
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;
}
// With expanded and improved functionality by OP Labs
/// @dev A minimal wrapper around the OwnerManager contract. This contract is meant to be initialized with
/// the same owners as a Safe instance, and then used to simulate the resulting owners list
......@@ -146,7 +49,108 @@ struct SafeInstance {
DeployedSafe safe;
}
library Sort {
/// @dev Sorts an array of addresses in place
function sort(address[] memory arr) public pure returns (address[] memory) {
LibSort.sort(arr);
return arr;
}
}
library SafeTestLib {
/// @dev The address of foundry's VM contract
address constant VM_ADDR = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
/// @dev The address of the first owner in the linked list of owners
address constant SENTINEL_OWNERS = address(0x1);
/// @dev Get the address from a private key
function getAddr(uint256 pk) internal pure returns (address) {
return Vm(VM_ADDR).addr(pk);
}
/// @dev Get arrays of addresses and private keys. The arrays are sorted by address, and the addresses are labelled
function makeAddrsAndKeys(uint256 num) internal returns (address[] memory addrs, uint256[] memory keys) {
keys = new uint256[](num);
addrs = new address[](num);
for (uint256 i; i < num; i++) {
uint256 key = uint256(keccak256(abi.encodePacked(i)));
keys[i] = key;
}
for (uint256 i; i < num; i++) {
addrs[i] = Vm(VM_ADDR).addr(keys[i]);
Vm(VM_ADDR).label(
getAddr(keys[i]), string.concat("SAFETEST: Signer ", string(abi.encodePacked(bytes32(i))))
);
}
}
bytes12 constant ADDR_MASK = 0xffffffffffffffffffffffff;
/// @dev Encode a smart contract wallet as a private key
function encodeSmartContractWalletAsPK(address addr) internal pure returns (uint256 encodedPK) {
assembly {
let addr_b32 := addr
encodedPK := or(addr, ADDR_MASK)
}
}
/// @dev Decode a smart contract wallet as an address from a private key
function decodeSmartContractWalletAsAddress(uint256 pk) internal pure returns (address decodedAddr) {
assembly {
let addr := shl(96, pk)
decodedAddr := shr(96, addr)
}
}
/// @dev Checks if a private key is an encoded smart contract address
function isSmartContractPK(uint256 pk) internal pure returns (bool isEncoded) {
assembly {
isEncoded := eq(shr(160, pk), shr(160, ADDR_MASK))
}
}
/// @dev Sorts an array of private keys by the computed address
/// If the private key is a smart contract wallet, it will be decoded and sorted by the address
function sortPKsByComputedAddress(uint256[] memory _pks) internal 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++) {
uint256 pk = _pks[i];
address signer = SafeTestLib.getAddr(pk);
if (isSmartContractPK(pk)) {
signer = decodeSmartContractWalletAsAddress(pk);
}
addresses[i] = signer;
accounts[i][0] = bytes32(abi.encode(signer));
accounts[i][1] = bytes32(pk);
}
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;
}
/// @dev A wrapper for the full execTransaction method, if no signatures are provided it will
/// generate them for all owners.
function execTransaction(
......@@ -442,14 +446,14 @@ contract SafeTestTools {
public
returns (SafeInstance memory)
{
uint256[] memory sortedPKs = sortPKsByComputedAddress(ownerPKs);
uint256[] memory sortedPKs = SafeTestLib.sortPKsByComputedAddress(ownerPKs);
address[] memory owners = new address[](sortedPKs.length);
for (uint256 i; i < sortedPKs.length; i++) {
if (isSmartContractPK(sortedPKs[i])) {
owners[i] = decodeSmartContractWalletAsAddress(sortedPKs[i]);
if (SafeTestLib.isSmartContractPK(sortedPKs[i])) {
owners[i] = SafeTestLib.decodeSmartContractWalletAsAddress(sortedPKs[i]);
} else {
owners[i] = getAddr(sortedPKs[i]);
owners[i] = SafeTestLib.getAddr(sortedPKs[i]);
}
}
// store the initialization parameters
......@@ -482,7 +486,7 @@ contract SafeTestTools {
});
instances.push(instance0);
Vm(VM_ADDR).deal(address(safe0), initialBalance);
Vm(SafeTestLib.VM_ADDR).deal(address(safe0), initialBalance);
return instance0;
}
......@@ -531,7 +535,7 @@ contract SafeTestTools {
}
function _setupSafe() public returns (SafeInstance memory) {
(, uint256[] memory defaultPKs) = makeAddrsAndKeys(3);
(, uint256[] memory defaultPKs) = SafeTestLib.makeAddrsAndKeys(3);
return _setupSafe(
defaultPKs,
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment