import Message from '@nomiclabs/ethereumjs-vm/dist/evm/message' ; import { FactoryOptions } from '@nomiclabs/hardhat-ethers/types' ; import { BaseContract , ContractFactory , ethers } from 'ethers' ; import { Interface } from 'ethers/lib/utils' ; import { ethers as hardhatEthers } from 'hardhat' ; import { Observable } from 'rxjs' ; import { distinct , filter , map , share , withLatestFrom } from 'rxjs/operators' ; import { EditableStorageLogic as EditableStorage } from '../logic/editable-storage-logic' ; import { ProgrammableFunctionLogic , SafeProgrammableContract } from '../logic/programmable-function-logic' ; import { ObservableVM } from '../observable-vm' ; import { Sandbox } from '../sandbox' ; import { ContractCall , FakeContract , MockContractFactory , ProgrammableContractFunction , ProgrammedReturnValue } from '../types' ; import { convertPojoToStruct , fromFancyAddress , impersonate , isPojo , toFancyAddress , toHexString } from '../utils' ; import { getStorageLayout } from '../utils/storage' ; export async function createFakeContract < Contract extends BaseContract > ( vm , address , contractInterface , provider ) { const fake = ( await initContract ( vm , address , contractInterface , provider ) ) as unknown as FakeContract < Contract > ; const contractFunctions = getContractFunctionsNameAndSighash ( contractInterface , Object . keys ( fake . functions ) ) ; contractFunctions . forEach ( ( [ sighash , name ] ) => { const { encoder , calls$ , results$ } = getFunctionEventData ( vm , contractInterface , fake . address , sighash ) ; const functionLogic = new SafeProgrammableContract ( name , calls$ , results$ , encoder ) ; fillProgrammableContractFunction ( fake [ name ] , functionLogic ) ; } ) ; return fake ; } function mockifyContractFactory < T extends ContractFactory > ( vm , contractName , factory ) { const realDeploy = factory . deploy ; factory . deploy = async ( ... args ) => { const mock = await realDeploy . apply ( factory , args ) ; const contractFunctions = getContractFunctionsNameAndSighash ( mock . interface , Object . keys ( mock . functions ) ) ; contractFunctions . forEach ( ( [ sighash , name ] ) => { const { encoder , calls$ , results$ } = getFunctionEventData ( vm , mock . interface , mock . address , sighash ) ; const functionLogic = new ProgrammableFunctionLogic ( name , calls$ , results$ , encoder ) ; fillProgrammableContractFunction ( mock [ name ] , functionLogic ) ; } ) ; const editableStorage = new EditableStorage ( await getStorageLayout ( contractName ) , vm . getManager ( ) , mock . address ) ; mock . setVariable = editableStorage . setVariable . bind ( editableStorage ) ; mock . wallet = await impersonate ( mock . address ) ; return mock ; } ; const realConnect = factory . connect ; factory . connect = ( ... args ) : MockContractFactory < T > => { const newFactory = realConnect . apply ( factory , args ) as MockContractFactory < T > ; return mockifyContractFactory ( vm , contractName , newFactory ) ; } ; return factory ; } export async function createMockContractFactory < T extends ContractFactory > ( vm , contractName , signerOrOptions ? ) { const factory = ( await hardhatEthers . getContractFactory ( contractName , signerOrOptions ) ) as unknown as MockContractFactory < T > ; return mockifyContractFactory ( vm , contractName , factory ) ; } async function initContract ( vm , address , contractInterface , provider ) { const contract = new ethers . Contract ( address , contractInterface , provider ) ; await vm . getManager ( ) . putContractCode ( toFancyAddress ( contract . address ) , Buffer . from ( '00' , 'hex' ) ) ; ( contract as any ) . wallet = await impersonate ( contract . address ) ; return contract ; } function getFunctionEventData ( vm , contractInterface , contractAddress , sighash ) { const encoder = getFunctionEncoder ( contractInterface , sighash ) ; const calls$ = parseAndFilterBeforeMessages ( vm . getBeforeMessages ( ) , contractInterface , contractAddress , sighash ) ; const results$ = vm . getAfterMessages ( ) . pipe ( withLatestFrom ( calls$ ) , distinct ( ( [ , call ] ) => call ) , map ( ( [ answer ] ) => answer ) ) ; return { encoder , calls$ , results$ } ; } function getFunctionEncoder ( contractInterface , sighash ) { if ( sighash === null ) { return ( values ) => values ; } else { return ( values ) => { const fnFragment = contractInterface . getFunction ( sighash ) ; try { return contractInterface . encodeFunctionResult ( fnFragment , [ values ] ) ; } catch { try { return contractInterface . encodeFunctionResult ( fnFragment , values ) ; } catch ( err ) { if ( isPojo ( values ) ) { return contractInterface . encodeFunctionResult ( fnFragment , convertPojoToStruct ( values , fnFragment ) ) ; } throw err ; } } } ; } } function parseAndFilterBeforeMessages ( messages$ , contractInterface , contractAddress , sighash ) { return messages$ . pipe ( filter ( ( message ) => { if ( sighash === null ) { return message . data . length === 0 ; } else { return toHexString ( message . data . slice ( 0 , 4 ) ) === sighash ; } } ) , filter ( ( message ) => message . to . toString ( ) . toLowerCase ( ) === contractAddress . toLowerCase ( ) ) , map ( ( message ) => parseMessage ( message , contractInterface , sighash ) ) , share ( ) ) ; } function fillProgrammableContractFunction ( fn , logic ) { fn . _watchable = logic ; fn . atCall = logic . atCall . bind ( logic ) ; fn . getCall = logic . getCall . bind ( logic ) ; fn . returns = logic . returns . bind ( logic ) ; fn . returnsAtCall = logic . returnsAtCall . bind ( logic ) ; fn . reverts = logic . reverts . bind ( logic ) ; fn . revertsAtCall = logic . revertsAtCall . bind ( logic ) ; fn . whenCalledWith = logic . whenCalledWith . bind ( logic ) ; fn . reset = logic . reset . bind ( logic ) ; } param contractInterface param names returns function getContractFunctionsNameAndSighash ( contractInterface , names ) { let functions = { } ; names . forEach ( ( name ) => { const sighash = contractInterface . getSighash ( name ) ; if ( ! functions [ sighash ] || ! name . includes ( '(' ) ) { functions [ sighash ] = name ; } } ) ; return [ ... Object . entries ( functions ) , [ null , 'fallback' ] ] ; } function parseMessage ( message , contractInterface , sighash ) { return { args : sighash === null ? toHexString ( message . data ) : getMessageArgs ( message . data , contractInterface , sighash ) , nonce : Sandbox . getNextNonce ( ) , target : fromFancyAddress ( message . delegatecall ? message . _codeAddress : message . to ) , } ; } function getMessageArgs ( messageData , contractInterface , sighash ) { try { return contractInterface . decodeFunctionData ( contractInterface . getFunction ( sighash ) . format ( ) , toHexString ( messageData ) ) as unknown [ ] ; } catch ( err ) { throw new Error ( `Failed to decode message data: ${ err }` ) ; } }