Skip to content

Commit

Permalink
Merge pull request #26 from axieinfinity/feature/wallet
Browse files Browse the repository at this point in the history
feat: allow forge script to handle loading `env` private key or `trezor` address
  • Loading branch information
TuDo1403 authored Dec 28, 2023
2 parents 08eff70 + 7c696b0 commit d524af3
Show file tree
Hide file tree
Showing 4 changed files with 190 additions and 30 deletions.
8 changes: 1 addition & 7 deletions run.sh
Original file line number Diff line number Diff line change
@@ -1,17 +1,11 @@
verify_arg=""
extra_argument=""
op_command="op run --env-file="./.env" --"

for arg in "$@"; do
case $arg in
--trezor)
op_command=""
extra_argument+=trezor@
;;
--broadcast)
op_command="op run --env-file="./.env" --"
# verify_arg="--verify --verifier sourcify --verifier-url https://sourcify.roninchain.com/server/"
;;
--log)
set -- "${@/#--log/}"
extra_argument+=log@
Expand All @@ -24,4 +18,4 @@ done
extra_argument="${extra_argument%%@}"

calldata=$(cast calldata 'run()')
${op_command} forge script ${verify_arg} --legacy ${@} --sig 'run(bytes,string)' ${calldata} "${extra_argument}"
forge script ${verify_arg} ${@} -g 200 --sig 'run(bytes,string)' ${calldata} "${extra_argument}"
25 changes: 9 additions & 16 deletions script/BaseGeneralConfig.sol
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.19;

import { Vm } from "../lib/forge-std/src/Vm.sol";
import { Vm, VmSafe } from "../lib/forge-std/src/Vm.sol";
import { StdStyle } from "../lib/forge-std/src/StdStyle.sol";
import { console2 as console } from "../lib/forge-std/src/console2.sol";
import { LibString } from "../lib/solady/src/utils/LibString.sol";
import { WalletConfig } from "./configs/WalletConfig.sol";
import { RuntimeConfig } from "./configs/RuntimeConfig.sol";
import { MigrationConfig } from "./configs/MigrationConfig.sol";
Expand All @@ -17,11 +16,8 @@ import { LibSharedAddress } from "./libraries/LibSharedAddress.sol";

contract BaseGeneralConfig is RuntimeConfig, WalletConfig, ContractConfig, NetworkConfig, MigrationConfig {
using StdStyle for string;
using LibString for string;
using EnumerableSet for EnumerableSet.AddressSet;

Vm internal constant vm = Vm(LibSharedAddress.VM);

fallback() external {
if (msg.sig == ISharedParameter.sharedArguments.selector) {
bytes memory returnData = getRawSharedArguments();
Expand Down Expand Up @@ -91,19 +87,12 @@ contract BaseGeneralConfig is RuntimeConfig, WalletConfig, ContractConfig, Netwo
}

function _setUpDefaultSender() private {
// by default we will read private key from .env
_envPk = vm.envUint(getPrivateKeyEnvLabel(getCurrentNetwork()));
_envSender = vm.rememberKey(_envPk);

label(block.chainid, _envSender, "ENVSender");
console.log("GeneralConfig:", vm.getLabel(_envSender));

_setUpSender();
}

function getSender() public view virtual override returns (address payable sender) {
sender = _option.trezor ? payable(_trezorSender) : payable(_envSender);
require(sender != address(0), "GeneralConfig: Sender is address(0x0)");
require(sender != address(0x0), "GeneralConfig: Sender is address(0x0)");
}

function setAddress(TNetwork network, TContract contractType, address contractAddr) public virtual {
Expand All @@ -127,10 +116,14 @@ contract BaseGeneralConfig is RuntimeConfig, WalletConfig, ContractConfig, Netwo

function _handleRuntimeConfig() internal virtual override {
if (_option.trezor) {
string memory str = vm.envString(deployerEnvLabel());
_trezorSender = vm.parseAddress(str.replace(trezorPrefix(), ""));
_loadTrezorAccount();
label(block.chainid, _trezorSender, "TrezorSender");
console.log("GeneralConfig:", vm.getLabel(_trezorSender));
console.log("GeneralConfig:", vm.getLabel(_trezorSender), "Enabled!");
} else {
string memory envLabel = getPrivateKeyEnvLabel(getCurrentNetwork());
_loadENVAccount(envLabel);
label(block.chainid, _envSender, "ENVSender");
console.log("GeneralConfig:", vm.getLabel(_envSender), "Enabled!");
}
}
}
158 changes: 152 additions & 6 deletions script/configs/WalletConfig.sol
Original file line number Diff line number Diff line change
@@ -1,24 +1,170 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.19;

import { CommonBase } from "../../lib/forge-std/src/Base.sol";
import { LibString } from "../../lib/solady/src/utils/LibString.sol";
import { IWalletConfig } from "../interfaces/configs/IWalletConfig.sol";

abstract contract WalletConfig is IWalletConfig {
uint256 internal _envPk;
abstract contract WalletConfig is CommonBase, IWalletConfig {
using LibString for string;

string internal _envLabel;
address internal _envSender;
address internal _trezorSender;
WalletOption internal _walletOption;

function getSender() public view virtual returns (address payable sender);

function getSenderPk() public view virtual returns (uint256) {
return _envPk;
}

function trezorPrefix() public view virtual returns (string memory) {
return "trezor://";
}

function deployerEnvLabel() public view virtual returns (string memory) {
return "DEPLOYER";
}

function ethSignMessage(address by, string memory message, WalletOption walletOption)
public
returns (bytes memory sig)
{
sig =
walletOption == WalletOption.Env ? envEthSignMessage(by, message, _envLabel) : trezorEthSignMessage(by, message);
}

function ethSignMessage(string memory message) public returns (bytes memory sig) {
sig = _walletOption == WalletOption.Env
? envEthSignMessage(_envSender, message, _envLabel)
: trezorEthSignMessage(_trezorSender, message);
}

function envEthSignMessage(address by, string memory message, string memory envLabel)
public
returns (bytes memory sig)
{
string[] memory commandInput = new string[](8);
commandInput[0] = "cast";
commandInput[1] = "wallet";
commandInput[2] = "sign";
commandInput[3] = "--from";
commandInput[4] = vm.toString(by);
commandInput[5] = "--private-key";
commandInput[6] = LibString.toHexString(_loadENVPrivateKey(envLabel));
commandInput[7] = message;

sig = vm.ffi(commandInput);
}

function trezorEthSignMessage(address by, string memory message) public returns (bytes memory sig) {
string[] memory commandInput = new string[](7);
commandInput[0] = "cast";
commandInput[1] = "wallet";
commandInput[2] = "sign";
commandInput[3] = "--from";
commandInput[4] = vm.toString(by);
commandInput[5] = "--trezor";
commandInput[6] = message;

sig = vm.ffi(commandInput);
}

function signTypedDataV4(address by, string memory filePath, WalletOption walletOption)
public
returns (bytes memory sig)
{
sig = walletOption == WalletOption.Env
? envSignTypedDataV4(by, filePath, _envLabel)
: trezorSignTypedDataV4(by, filePath);
}

function signTypedDataV4(string memory filePath) public returns (bytes memory sig) {
sig = _walletOption == WalletOption.Env
? envSignTypedDataV4(_envSender, filePath, _envLabel)
: trezorSignTypedDataV4(_trezorSender, filePath);
}

function envSignTypedDataV4(address by, string memory filePath, string memory envLabel)
public
returns (bytes memory sig)
{
string[] memory commandInput = new string[](10);
commandInput[0] = "cast";
commandInput[1] = "wallet";
commandInput[2] = "sign";
commandInput[3] = "--from";
commandInput[4] = vm.toString(by);
commandInput[5] = "--private-key";
commandInput[6] = LibString.toHexString(_loadENVPrivateKey(envLabel));
commandInput[7] = "--data";
commandInput[8] = "--from-file";
commandInput[9] = filePath;

sig = vm.ffi(commandInput);
}

function trezorSignTypedDataV4(address by, string memory filePath) public returns (bytes memory sig) {
string[] memory commandInput = new string[](9);
commandInput[0] = "cast";
commandInput[1] = "wallet";
commandInput[2] = "sign";
commandInput[3] = "--from";
commandInput[4] = vm.toString(by);
commandInput[5] = "--trezor";
commandInput[6] = "--data";
commandInput[7] = "--from-file";
commandInput[8] = filePath;

sig = vm.ffi(commandInput);
}

function _loadTrezorAccount() internal {
if (tx.origin != DEFAULT_SENDER) {
_trezorSender = tx.origin;
} else {
try vm.envString(deployerEnvLabel()) returns (string memory str) {
_trezorSender = vm.parseAddress(str.replace(trezorPrefix(), ""));
} catch {
revert(
string.concat(
"\nGeneralConfig: Error finding trezor address!\n- Please override default sender with `--sender {your_trezor_account}` tag \n- Or make `.env` file and create field `",
deployerEnvLabel(),
"=",
trezorPrefix(),
"{your_trezor_account}`"
)
);
}
}
_walletOption = WalletOption.Trezor;
}

function _loadENVAccount(string memory envLabel) internal {
_envLabel = envLabel;
_walletOption = WalletOption.Env;
_envSender = vm.rememberKey(_loadENVPrivateKey(envLabel));
}

function _loadENVPrivateKey(string memory envLabel) private returns (uint256) {
try vm.envUint(envLabel) returns (uint256 pk) {
return pk;
} catch {
string[] memory commandInput = new string[](3);

try vm.envString(envLabel) returns (string memory data) {
commandInput[2] = data;
} catch {
revert(
string.concat(
"\nGeneralConfig: Error finding env address!\n- Please make `.env` file and create field `",
envLabel,
"=",
"{op_secret_reference_or_your_private_key}`"
)
);
}
commandInput[0] = "op";
commandInput[1] = "read";

return vm.parseUint(vm.toString(vm.ffi(commandInput)));
}
}
}
29 changes: 28 additions & 1 deletion script/interfaces/configs/IWalletConfig.sol
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,38 @@
pragma solidity ^0.8.19;

interface IWalletConfig {
function getSenderPk() external view returns (uint256);
enum WalletOption {
Env,
Trezor
}

function getSender() external view returns (address payable sender);

function trezorPrefix() external view returns (string memory);

function deployerEnvLabel() external view returns (string memory);

function ethSignMessage(address by, string memory message, WalletOption walletOption)
external
returns (bytes memory sig);

function ethSignMessage(string memory message) external returns (bytes memory sig);

function envEthSignMessage(address by, string memory message, string memory envLabel)
external
returns (bytes memory sig);

function envSignTypedDataV4(address by, string memory filePath, string memory envLabel)
external
returns (bytes memory sig);

function trezorEthSignMessage(address by, string memory message) external returns (bytes memory sig);

function trezorSignTypedDataV4(address by, string memory filePath) external returns (bytes memory sig);

function signTypedDataV4(address by, string memory filePath, WalletOption walletOption)
external
returns (bytes memory sig);

function signTypedDataV4(string memory filePath) external returns (bytes memory sig);
}

0 comments on commit d524af3

Please sign in to comment.