Repository: defi-wonderland/smock (ts)

Issue:
fix: support faking and mocking delegated calls
**Describe the bug**
Smock currently does not support any interaction with delegated calls.

**Reproduction steps**

1. Create a Proxy contract
1. Fake a fake using an interface of a contract with the function `foo`
1. Force the fake to return `true` when the function `foo` is called
1. Call the method `foo` of the fake through the proxy
1. Expect to receive `true`

**Expected behavior**
Be able to work with mock and fakes when using delegated calls.

As well, it would be great to have a way to assert that a call to a function was delegated from address 0x123.

Pull request:
fix: delegated calls support
**Description**
Fixes #131

* Adds support to delegated calls in both fakes and mocks
* Adds custom assertion `expect(myFake.myFunction).to.be.delegatedFrom(myProxy.address)`

Canary version: 0.0.0-9be6e836

Patch:
diff --git a/docs/source/fakes.rst b/docs/source/fakes.rst
index f0cc0de..4c973a4 100644
--- a/docs/source/fakes.rst
+++ b/docs/source/fakes.rst
@@ -371,3 +371,17 @@ Modifying the :code:`receive` function
 .. code-block:: typescript
 
   myFake.receive.returns();
+
+Delegated calls
+***************
+
+Calls to a contract function via delegated calls do behave the same as a regular call, so you can enforce a return value, assert the calls details, etc...
+In addition, you also have custom assertions for delegated calls.
+
+Assert delegated caller
+#######################
+
+.. code-block:: typescript
+
+  expect(myFake.myFunction).to.be.delegatedFrom(myProxy.address);
+
diff --git a/docs/source/mocks.rst b/docs/source/mocks.rst
index 4a6e39b..a397d08 100644
--- a/docs/source/mocks.rst
+++ b/docs/source/mocks.rst
@@ -97,7 +97,7 @@ Setting the value of a struct
 
 
 Setting the value of a mapping (won't affect other keys)
-#######################################################
+########################################################
 
 .. code-block:: typescript
 
diff --git a/src/chai-plugin/matchers.ts b/src/chai-plugin/matchers.ts
index 3b54866..ee7394e 100644
--- a/src/chai-plugin/matchers.ts
+++ b/src/chai-plugin/matchers.ts
@@ -137,4 +137,5 @@ export const matchers: Chai.ChaiPlugin = (chai: Chai.ChaiStatic, utils: Chai.Cha
   smockMethodWithWatchableContractArg('calledImmediatelyAfter', 'been called immediately after %1');
   smockMethod('calledWith', 'been called with arguments %*', '%D');
   smockMethod('calledOnceWith', 'been called exactly once with arguments %*', '%D');
+  smockMethod('delegatedFrom', 'been called via a delegated call by %*', '');
 };
diff --git a/src/chai-plugin/types.ts b/src/chai-plugin/types.ts
index 4556d64..f7e2bf4 100644
--- a/src/chai-plugin/types.ts
+++ b/src/chai-plugin/types.ts
@@ -54,6 +54,10 @@ declare global {
        * Returns true when called at exactly once with the provided arguments.
        */
       calledOnceWith(...args: any[]): Assertion;
+      /**
+       * Returns true when the function call was made via delegatecall by the given address.
+       */
+      delegatedFrom(delegatorAddress: string): Assertion;
     }
   }
 }
diff --git a/src/factories/smock-contract.ts b/src/factories/smock-contract.ts
index 0978cc1..75e3b64 100644
--- a/src/factories/smock-contract.ts
+++ b/src/factories/smock-contract.ts
@@ -153,7 +153,10 @@ function parseAndFilterBeforeMessages(
       }
     }),
     // Ensure the message is directed to this contract
-    filter((message) => message.to.toString().toLowerCase() === contractAddress.toLowerCase()),
+    filter((message) => {
+      const target = message.delegatecall ? message.codeAddress : message.to;
+      return target.toString().toLowerCase() === contractAddress.toLowerCase();
+    }),
     map((message) => parseMessage(message, contractInterface, sighash)),
     share()
   );
@@ -199,7 +202,8 @@ function parseMessage(message: Message, contractInterface: Interface, sighash: s
   return {
     args: sighash === null ? toHexString(message.data) : getMessageArgs(message.data, contractInterface, sighash),
     nonce: Sandbox.getNextNonce(),
-    target: fromFancyAddress(message.delegatecall ? message._codeAddress : message.to),
+    target: fromFancyAddress(message.delegatecall ? message.codeAddress : message.to),
+    delegatedFrom: message.delegatecall ? fromFancyAddress(message.to) : undefined,
   };
 }
 
diff --git a/src/logic/watchable-function-logic.ts b/src/logic/watchable-function-logic.ts
index 33f0591..57e2bf4 100644
--- a/src/logic/watchable-function-logic.ts
+++ b/src/logic/watchable-function-logic.ts
@@ -40,6 +40,10 @@ export class WatchableFunctionLogic {
     return this.getCalledOnce() && this.calledWith(...expectedCallArgs);
   }
 
+  delegatedFrom(delegatorAddress: string): boolean {
+    return !!this.callHistory.find((call) => call.delegatedFrom?.toLowerCase() === delegatorAddress.toLowerCase());
+  }
+
   calledBefore(anotherWatchableContract: WatchableFunctionLogic): boolean {
     return this.compareWatchableContractNonces(
       this,
diff --git a/src/types.ts b/src/types.ts
index 06c909c..ff264ec 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -34,6 +34,7 @@ export interface ContractCall {
   args: unknown[] | string;
   nonce: number;
   target: string;
+  delegatedFrom?: string;
 }
 
 export interface SetVariablesType {


Test patch:
diff --git a/test/contracts/programmable-function-logic/Delegator.sol b/test/contracts/programmable-function-logic/Delegator.sol
new file mode 100644
index 0000000..726c949
--- /dev/null
+++ b/test/contracts/programmable-function-logic/Delegator.sol
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+contract Delegator {
+  function delegateGetBoolean(address _addy) external returns (bool) {
+    // solhint-disable-next-line avoid-low-level-calls
+    (, bytes memory _data) = _addy.delegatecall(abi.encodeWithSignature('getBoolean()'));
+    return abi.decode(_data, (bool));
+  }
+}
diff --git a/test/unit/programmable-function-logic/type-handling.spec.ts b/test/unit/programmable-function-logic/type-handling.spec.ts
index 5e186ee..6b135fb 100644
--- a/test/unit/programmable-function-logic/type-handling.spec.ts
+++ b/test/unit/programmable-function-logic/type-handling.spec.ts
@@ -1,6 +1,6 @@
 import { FakeContract, smock } from '@src';
 import { convertStructToPojo } from '@src/utils';
-import { Returner } from '@typechained';
+import { Delegator__factory, Returner } from '@typechained';
 import chai, { expect } from 'chai';
 import { BigNumber, utils } from 'ethers';
 import { ethers } from 'hardhat';
@@ -54,6 +54,20 @@ describe('ProgrammableFunctionLogic: Type Handling', () => {
 
         expect(await fake.callStatic.getBytes32()).to.equal(expected);
       });
+
+      it('should be able to return a boolean through a delegatecall', async () => {
+        const delegatorFactory = (await ethers.getContractFactory('Delegator')) as Delegator__factory;
+        const delegator = await delegatorFactory.deploy();
+
+        const expected = true;
+        fake.getBoolean.returns(expected);
+
+        const result = await delegator.callStatic.delegateGetBoolean(fake.address);
+
+        expect(result).to.equal(expected);
+        expect(fake.getBoolean).to.be.calledOnce;
+        expect(fake.getBoolean).to.be.delegatedFrom(delegator.address);
+      });
     });
   });
 
diff --git a/test/unit/watchable-function-logic/misc.spec.ts b/test/unit/watchable-function-logic/misc.spec.ts
index 601fb18..1a54d15 100644
--- a/test/unit/watchable-function-logic/misc.spec.ts
+++ b/test/unit/watchable-function-logic/misc.spec.ts
@@ -30,4 +30,10 @@ describe('WatchableFunctionLogic: Miscellaneous', () => {
       fake.receiveBoolean.should.always.have.been.calledOnceWith(true);
     }).to.throw('always flag is not supported for method calledOnceWith');
   });
+
+  it('should throw when expecting a delegatedFrom that did not happen', async () => {
+    expect(() => {
+      fake.receiveBoolean.should.be.delegatedFrom(fake.address);
+    }).to.throw(`expected receiveBoolean to have been called via a delegated call by '${fake.address}'`);
+  });
 });