From afd4805421b24409036536d6fbc349d9df39078e Mon Sep 17 00:00:00 2001 From: Steve Marx Date: Fri, 6 Nov 2020 22:03:07 -0500 Subject: [PATCH] rewrite ZeroEx in Yul (#23) rewrite ZeroEx in Yul --- contracts/zero-ex/CHANGELOG.json | 9 ++ contracts/zero-ex/contracts/src/IZeroEx.sol | 10 --- contracts/zero-ex/contracts/src/ZeroEx.sol | 90 +++++++++---------- .../ISimpleFunctionRegistryFeature.sol | 8 ++ .../SimpleFunctionRegistryFeature.sol | 13 +++ .../contracts/test/TestInitialMigration.sol | 4 +- .../features/simple_function_registry_test.ts | 2 +- contracts/zero-ex/test/full_migration_test.ts | 5 +- contracts/zero-ex/test/zero_ex_test.ts | 2 +- 9 files changed, 84 insertions(+), 59 deletions(-) diff --git a/contracts/zero-ex/CHANGELOG.json b/contracts/zero-ex/CHANGELOG.json index 9436988200..a348d2bf44 100644 --- a/contracts/zero-ex/CHANGELOG.json +++ b/contracts/zero-ex/CHANGELOG.json @@ -1,4 +1,13 @@ [ + { + "version": "0.9.0", + "changes": [ + { + "note": "Rewrite the ZeroEx contract in Yul", + "pr": 23 + } + ] + }, { "version": "0.8.0", "changes": [ diff --git a/contracts/zero-ex/contracts/src/IZeroEx.sol b/contracts/zero-ex/contracts/src/IZeroEx.sol index 8a84a0425a..beae4e5e2c 100644 --- a/contracts/zero-ex/contracts/src/IZeroEx.sol +++ b/contracts/zero-ex/contracts/src/IZeroEx.sol @@ -44,14 +44,4 @@ interface IZeroEx is /// @dev Fallback for just receiving ether. receive() external payable; - - // solhint-enable state-visibility - - /// @dev Get the implementation contract of a registered function. - /// @param selector The function selector. - /// @return impl The implementation contract address. - function getFunctionImplementation(bytes4 selector) - external - view - returns (address impl); } diff --git a/contracts/zero-ex/contracts/src/ZeroEx.sol b/contracts/zero-ex/contracts/src/ZeroEx.sol index a18d762a93..3ebc087f89 100644 --- a/contracts/zero-ex/contracts/src/ZeroEx.sol +++ b/contracts/zero-ex/contracts/src/ZeroEx.sol @@ -19,19 +19,12 @@ pragma solidity ^0.6.5; pragma experimental ABIEncoderV2; -import "@0x/contracts-utils/contracts/src/v06/LibBytesV06.sol"; -import "./migrations/LibBootstrap.sol"; import "./features/BootstrapFeature.sol"; import "./storage/LibProxyStorage.sol"; -import "./errors/LibProxyRichErrors.sol"; - /// @dev An extensible proxy contract that serves as a universal entry point for /// interacting with the 0x protocol. contract ZeroEx { - // solhint-disable separate-by-one-line-in-contract,indent,var-name-mixedcase - using LibBytesV06 for bytes; - /// @dev Construct this contract and register the `BootstrapFeature` feature. /// After constructing this contract, `bootstrap()` should be called /// by `bootstrap()` to seed the initial feature set. @@ -44,48 +37,55 @@ contract ZeroEx { address(bootstrap); } + // solhint-disable state-visibility /// @dev Forwards calls to the appropriate implementation contract. fallback() external payable { - bytes4 selector = msg.data.readBytes4(0); - address impl = getFunctionImplementation(selector); - if (impl == address(0)) { - _revertWithData(LibProxyRichErrors.NotImplementedError(selector)); + // This is used in assembly below as impls_slot. + mapping(bytes4 => address) storage impls = + LibProxyStorage.getStorage().impls; + + assembly { + let cdlen := calldatasize() + + // equivalent of receive() external payable {} + if iszero(cdlen) { + return(0, 0) + } + + // Store at 0x40, to leave 0x00-0x3F for slot calculation below. + calldatacopy(0x40, 0, cdlen) + let selector := and(mload(0x40), 0xFFFFFFFF00000000000000000000000000000000000000000000000000000000) + + // Slot for impls[selector] is keccak256(selector . impls_slot). + mstore(0, selector) + mstore(0x20, impls_slot) + let slot := keccak256(0, 0x40) + + let delegate := sload(slot) + if iszero(delegate) { + // Revert with: + // abi.encodeWithSelector( + // bytes4(keccak256("NotImplementedError(bytes4)")), + // selector) + mstore(0, 0x734e6e1c00000000000000000000000000000000000000000000000000000000) + mstore(4, selector) + revert(0, 0x24) + } + + let success := delegatecall( + gas(), + delegate, + 0x40, cdlen, + 0, 0 + ) + let rdlen := returndatasize() + returndatacopy(0, 0, rdlen) + if success { + return(0, rdlen) + } + revert(0, rdlen) } - - (bool success, bytes memory resultData) = impl.delegatecall(msg.data); - if (!success) { - _revertWithData(resultData); - } - _returnWithData(resultData); - } - - /// @dev Fallback for just receiving ether. - receive() external payable {} - - // solhint-enable state-visibility - - /// @dev Get the implementation contract of a registered function. - /// @param selector The function selector. - /// @return impl The implementation contract address. - function getFunctionImplementation(bytes4 selector) - public - view - returns (address impl) - { - return LibProxyStorage.getStorage().impls[selector]; - } - - /// @dev Revert with arbitrary bytes. - /// @param data Revert data. - function _revertWithData(bytes memory data) private pure { - assembly { revert(add(data, 32), mload(data)) } - } - - /// @dev Return with arbitrary bytes. - /// @param data Return data. - function _returnWithData(bytes memory data) private pure { - assembly { return(add(data, 32), mload(data)) } } } diff --git a/contracts/zero-ex/contracts/src/features/ISimpleFunctionRegistryFeature.sol b/contracts/zero-ex/contracts/src/features/ISimpleFunctionRegistryFeature.sol index 1e5ef61c25..583b97f9ad 100644 --- a/contracts/zero-ex/contracts/src/features/ISimpleFunctionRegistryFeature.sol +++ b/contracts/zero-ex/contracts/src/features/ISimpleFunctionRegistryFeature.sol @@ -57,4 +57,12 @@ interface ISimpleFunctionRegistryFeature { external view returns (address impl); + + /// @dev Get the implementation contract of a registered function. + /// @param selector The function selector. + /// @return impl The implementation contract address. + function getFunctionImplementation(bytes4 selector) + external + view + returns (address impl); } diff --git a/contracts/zero-ex/contracts/src/features/SimpleFunctionRegistryFeature.sol b/contracts/zero-ex/contracts/src/features/SimpleFunctionRegistryFeature.sol index ac98b64dc7..32a2bb7ab5 100644 --- a/contracts/zero-ex/contracts/src/features/SimpleFunctionRegistryFeature.sol +++ b/contracts/zero-ex/contracts/src/features/SimpleFunctionRegistryFeature.sol @@ -56,6 +56,7 @@ contract SimpleFunctionRegistryFeature is // Register getters. _extend(this.getRollbackLength.selector, _implementation); _extend(this.getRollbackEntryAtIndex.selector, _implementation); + _extend(this.getFunctionImplementation.selector, _implementation); return LibBootstrap.BOOTSTRAP_SUCCESS; } @@ -151,6 +152,18 @@ contract SimpleFunctionRegistryFeature is return LibSimpleFunctionRegistryStorage.getStorage().implHistory[selector][idx]; } + /// @dev Get the implementation contract of a registered function. + /// @param selector The function selector. + /// @return impl The implementation contract address. + function getFunctionImplementation(bytes4 selector) + external + override + view + returns (address impl) + { + return LibProxyStorage.getStorage().impls[selector]; + } + /// @dev Register or replace a function. /// @param selector The function selector. /// @param impl The implementation contract for the function. diff --git a/contracts/zero-ex/contracts/test/TestInitialMigration.sol b/contracts/zero-ex/contracts/test/TestInitialMigration.sol index f6337bc5ed..15e1e84279 100644 --- a/contracts/zero-ex/contracts/test/TestInitialMigration.sol +++ b/contracts/zero-ex/contracts/test/TestInitialMigration.sol @@ -22,6 +22,7 @@ pragma experimental ABIEncoderV2; import "../src/ZeroEx.sol"; import "../src/features/IBootstrapFeature.sol"; import "../src/migrations/InitialMigration.sol"; +import "../src/features/SimpleFunctionRegistryFeature.sol"; contract TestInitialMigration is @@ -44,7 +45,8 @@ contract TestInitialMigration is { success = InitialMigration.bootstrap(owner, features); // Snoop the bootstrap feature contract. - bootstrapFeature = ZeroEx(address(uint160(address(this)))) + bootstrapFeature = + SimpleFunctionRegistryFeature(address(uint160(address(this)))) .getFunctionImplementation(IBootstrapFeature.bootstrap.selector); } diff --git a/contracts/zero-ex/test/features/simple_function_registry_test.ts b/contracts/zero-ex/test/features/simple_function_registry_test.ts index 53d0ec99b0..00c03e8fa8 100644 --- a/contracts/zero-ex/test/features/simple_function_registry_test.ts +++ b/contracts/zero-ex/test/features/simple_function_registry_test.ts @@ -66,7 +66,7 @@ blockchainTests.resets('SimpleFunctionRegistry feature', env => { it('`rollback()` to zero impl succeeds for unregistered function', async () => { await registry.rollback(testFnSelector, NULL_ADDRESS).awaitTransactionSuccessAsync(); - const impl = await zeroEx.getFunctionImplementation(testFnSelector).callAsync(); + const impl = await registry.getFunctionImplementation(testFnSelector).callAsync(); expect(impl).to.eq(NULL_ADDRESS); }); diff --git a/contracts/zero-ex/test/full_migration_test.ts b/contracts/zero-ex/test/full_migration_test.ts index 0611e45420..993472e026 100644 --- a/contracts/zero-ex/test/full_migration_test.ts +++ b/contracts/zero-ex/test/full_migration_test.ts @@ -12,6 +12,7 @@ import { IMetaTransactionsFeatureContract, IOwnableFeatureContract, ISignatureValidatorFeatureContract, + ISimpleFunctionRegistryFeatureContract, ITokenSpenderFeatureContract, ITransformERC20FeatureContract, TestFullMigrationContract, @@ -25,6 +26,7 @@ blockchainTests.resets('Full migration', env => { let zeroEx: ZeroExContract; let features: FullFeatures; let migrator: TestFullMigrationContract; + let registry: ISimpleFunctionRegistryFeatureContract; const transformerDeployer = randomAddress(); before(async () => { @@ -47,6 +49,7 @@ blockchainTests.resets('Full migration', env => { await migrator .initializeZeroEx(owner, zeroEx.address, features, { transformerDeployer }) .awaitTransactionSuccessAsync(); + registry = new ISimpleFunctionRegistryFeatureContract(zeroEx.address, env.provider, env.txDefaults); }); it('ZeroEx has the correct owner', async () => { @@ -157,7 +160,7 @@ blockchainTests.resets('Full migration', env => { for (const fn of featureInfo.fns) { it(`${fn} is registered`, async () => { const selector = contract.getSelector(fn); - const impl = await zeroEx.getFunctionImplementation(selector).callAsync(); + const impl = await registry.getFunctionImplementation(selector).callAsync(); expect(impl).to.not.eq(NULL_ADDRESS); }); diff --git a/contracts/zero-ex/test/zero_ex_test.ts b/contracts/zero-ex/test/zero_ex_test.ts index 566504f157..99c72de78d 100644 --- a/contracts/zero-ex/test/zero_ex_test.ts +++ b/contracts/zero-ex/test/zero_ex_test.ts @@ -83,7 +83,7 @@ blockchainTests.resets('ZeroEx contract', env => { // registry.getSelector('extendSelf'), ]; const selectors = [...ownableSelectors, ...registrySelectors]; - const impls = await Promise.all(selectors.map(s => zeroEx.getFunctionImplementation(s).callAsync())); + const impls = await Promise.all(selectors.map(s => registry.getFunctionImplementation(s).callAsync())); for (let i = 0; i < impls.length; ++i) { const selector = selectors[i]; const impl = impls[i];