Make Multicall context-aware

This commit is contained in:
ernestognw
2023-12-07 12:37:52 -06:00
parent 4eb67a408c
commit 9ce0340466
7 changed files with 111 additions and 13 deletions

View File

@ -13,6 +13,10 @@ import {Context} from "../utils/Context.sol";
* specification adding the address size in bytes (20) to the calldata size. An example of an unexpected
* behavior could be an unintended fallback (or another function) invocation while trying to invoke the `receive`
* function only accessible if `msg.data.length == 0`.
*
* WARNING: The usage of `delegatecall` in this contract is dangerous and may result in context corruption.
* Any forwarded request to this contract triggering a `delegatecall` to itself will result in an invalid {_msgSender}
* recovery.
*/
abstract contract ERC2771Context is Context {
/// @custom:oz-upgrades-unsafe-allow state-variable-immutable
@ -48,13 +52,11 @@ abstract contract ERC2771Context is Context {
* a call is not performed by the trusted forwarder or the calldata length is less than
* 20 bytes (an address length).
*/
function _msgSender() internal view virtual override returns (address sender) {
if (isTrustedForwarder(msg.sender) && msg.data.length >= 20) {
// The assembly code is more direct than the Solidity version using `abi.decode`.
/// @solidity memory-safe-assembly
assembly {
sender := shr(96, calldataload(sub(calldatasize(), 20)))
}
function _msgSender() internal view virtual override returns (address) {
uint256 calldataLength = msg.data.length;
uint256 contextSuffixLength = _contextSuffixLength();
if (isTrustedForwarder(msg.sender) && calldataLength >= contextSuffixLength) {
return address(bytes20(msg.data[calldataLength - contextSuffixLength:]));
} else {
return super._msgSender();
}
@ -66,10 +68,19 @@ abstract contract ERC2771Context is Context {
* 20 bytes (an address length).
*/
function _msgData() internal view virtual override returns (bytes calldata) {
if (isTrustedForwarder(msg.sender) && msg.data.length >= 20) {
return msg.data[:msg.data.length - 20];
uint256 calldataLength = msg.data.length;
uint256 contextSuffixLength = _contextSuffixLength();
if (isTrustedForwarder(msg.sender) && calldataLength >= contextSuffixLength) {
return msg.data[:calldataLength - contextSuffixLength];
} else {
return super._msgData();
}
}
/**
* @dev ERC-2771 specifies the context as being a single address (20 bytes).
*/
function _contextSuffixLength() internal view virtual override returns (uint256) {
return 20;
}
}