Skip to content

Commit

Permalink
Fix M2
Browse files Browse the repository at this point in the history
  • Loading branch information
crispymangoes committed Oct 27, 2023
1 parent 370c0b1 commit 4ee541d
Show file tree
Hide file tree
Showing 4 changed files with 294 additions and 2 deletions.
45 changes: 44 additions & 1 deletion src/base/ERC4626SharePriceOracle.sol
Expand Up @@ -87,6 +87,8 @@ contract ERC4626SharePriceOracle is AutomationCompatibleInterface {
error ERC4626SharePriceOracle__FuturePerformData();
error ERC4626SharePriceOracle__ContractKillSwitch();
error ERC4626SharePriceOracle__AlreadyInitialized();
error ERC4626SharePriceOracle__ParamHashDiffers();
error ERC4626SharePriceOracle__NoPendingUpkeepToHandle();

//============================== EVENTS ===============================

Expand Down Expand Up @@ -284,12 +286,53 @@ contract ERC4626SharePriceOracle is AutomationCompatibleInterface {
emit UpkeepRegistered(upkeepID, forwarder);
} else {
// Upkeep is pending.
bytes32 paramHash = keccak256(abi.encode(params));
bytes32 paramHash = keccak256(
abi.encode(
params.upkeepContract,
params.gasLimit,
params.adminAddress,
params.triggerType,
params.checkData,
params.offchainConfig
)
);
pendingUpkeepParamHash = paramHash;
emit UpkeepPending(paramHash);
}
}

/**
* @notice Finish setting forwarder address if `initialize` did not get an auto-approved upkeep.
*/
function handlePendingUpkeep(uint256 _upkeepId) external {
if (pendingUpkeepParamHash == bytes32(0) || automationForwarder != address(0))
revert ERC4626SharePriceOracle__NoPendingUpkeepToHandle();

IRegistry registry = IRegistry(automationRegistry);

IRegistry.UpkeepInfo memory upkeepInfo = registry.getUpkeep(_upkeepId);
// Build the param hash using upkeepInfo.
// The upkeep id has 16 bytes of entropy, that need to be shifted out(16*8=128).
// Then take the resulting number and only take the last byte of it to get the trigger type.
uint8 triggerType = uint8(_upkeepId >> 128);
bytes32 proposedParamHash = keccak256(
abi.encode(
upkeepInfo.target,
upkeepInfo.executeGas,
upkeepInfo.admin,
triggerType,
upkeepInfo.checkData,
upkeepInfo.offchainConfig
)
);
if (pendingUpkeepParamHash != proposedParamHash) revert ERC4626SharePriceOracle__ParamHashDiffers();

// Hashes match, so finish initialization.
address forwarder = registry.getForwarder(_upkeepId);
automationForwarder = forwarder;
emit UpkeepRegistered(_upkeepId, forwarder);
}

//============================== CHAINLINK AUTOMATION ===============================

/**
Expand Down
26 changes: 26 additions & 0 deletions src/interfaces/external/Chainlink/IRegistrar.sol
Expand Up @@ -15,5 +15,31 @@ interface IRegistrar {
uint96 amount;
}

enum AutoApproveType {
DISABLED,
ENABLED_SENDER_ALLOWLIST,
ENABLED_ALL
}

function registerUpkeep(RegistrationParams calldata requestParams) external returns (uint256 id);

function setTriggerConfig(
uint8 triggerType,
AutoApproveType autoApproveType,
uint32 autoApproveMaxAllowed
) external;

function owner() external view returns (address);

function approve(
string memory name,
address upkeepContract,
uint32 gasLimit,
address adminAddress,
uint8 triggerType,
bytes calldata checkData,
bytes memory triggerConfig,
bytes calldata offchainConfig,
bytes32 hash
) external;
}
15 changes: 15 additions & 0 deletions src/interfaces/external/Chainlink/IRegistry.sol
Expand Up @@ -2,5 +2,20 @@
pragma solidity 0.8.21;

interface IRegistry {
struct UpkeepInfo {
address target;
uint32 executeGas;
bytes checkData;
uint96 balance;
address admin;
uint64 maxValidBlocknumber;
uint32 lastPerformBlockNumber;
uint96 amountSpent;
bool paused;
bytes offchainConfig;
}

function getForwarder(uint256 upkeepID) external view returns (address forwarder);

function getUpkeep(uint256 id) external view returns (UpkeepInfo memory upkeepInfo);
}
210 changes: 209 additions & 1 deletion test/ERC4626SharePriceOracle/ERC4626SharePriceOracle.t.sol
Expand Up @@ -4,7 +4,7 @@ pragma solidity 0.8.21;
import { AaveATokenAdaptor } from "src/modules/adaptors/Aave/AaveATokenAdaptor.sol";
import { IPool } from "src/interfaces/external/IPool.sol";
import { MockDataFeed } from "src/mocks/MockDataFeed.sol";
import { ERC4626SharePriceOracle } from "src/base/ERC4626SharePriceOracle.sol";
import { ERC4626SharePriceOracle, IRegistry, IRegistrar } from "src/base/ERC4626SharePriceOracle.sol";

// Import Everything from Starter file.
import "test/resources/MainnetStarter.t.sol";
Expand Down Expand Up @@ -1151,6 +1151,214 @@ contract ERC4626SharePriceOracleTest is MainnetStarterTest, AdaptorHelperFunctio
sharePriceOracle.performUpkeep(performData);
}

function testCallingHandlePendingUpkeepWithNothingPending() external {
ERC4626 _target = ERC4626(address(cellar));
uint64 _heartbeat = 1 days;
uint64 _deviationTrigger = 0.0005e4;
uint64 _gracePeriod = 60 * 60; // 1 hr
uint16 _observationsToUse = 4; // TWAA duration is heartbeat * (observationsToUse - 1), so ~3 days.
address _automationRegistry = automationRegistryV2;
address _automationRegistrar = automationRegistrarV2;
address _automationAdmin = address(this);

// Setup share price oracle.
sharePriceOracle = new ERC4626SharePriceOracle(
_target,
_heartbeat,
_deviationTrigger,
_gracePeriod,
_observationsToUse,
_automationRegistry,
_automationRegistrar,
_automationAdmin,
address(LINK),
1e18,
0.01e4,
10e4
);

// Try calling `handlePendingUpkeep` before calling `initialize`.
vm.expectRevert(
bytes(
abi.encodeWithSelector(
ERC4626SharePriceOracle.ERC4626SharePriceOracle__NoPendingUpkeepToHandle.selector
)
)
);
sharePriceOracle.handlePendingUpkeep(0);

uint96 initialUpkeepFunds = 10e18;
deal(address(LINK), address(this), initialUpkeepFunds);
LINK.safeApprove(address(sharePriceOracle), initialUpkeepFunds);
sharePriceOracle.initialize(initialUpkeepFunds);

// Try calling `handlePendingUpkeep` after successfully calling `initialize`.
// Try calling `handlePendingUpkeep` before calling `initialize`.
vm.expectRevert(
bytes(
abi.encodeWithSelector(
ERC4626SharePriceOracle.ERC4626SharePriceOracle__NoPendingUpkeepToHandle.selector
)
)
);
sharePriceOracle.handlePendingUpkeep(0);
}

function testAuotmationDDOSAttackMakingNewUpkeepsPending() external {
IRegistrar registrarV2 = IRegistrar(automationRegistrarV2);
IRegistry registryV2 = IRegistry(automationRegistryV2);

ERC4626 _target = ERC4626(address(cellar));
uint64 _heartbeat = 1 days;
uint64 _deviationTrigger = 0.0005e4;
uint64 _gracePeriod = 60 * 60; // 1 hr
uint16 _observationsToUse = 4; // TWAA duration is heartbeat * (observationsToUse - 1), so ~3 days.
address _automationRegistry = automationRegistryV2;
address _automationRegistrar = automationRegistrarV2;
address _automationAdmin = address(this);

// Setup share price oracle.
sharePriceOracle = new ERC4626SharePriceOracle(
_target,
_heartbeat,
_deviationTrigger,
_gracePeriod,
_observationsToUse,
_automationRegistry,
_automationRegistrar,
_automationAdmin,
address(LINK),
1e18,
0.01e4,
10e4
);

IRegistrar.RegistrationParams memory params = IRegistrar.RegistrationParams({
name: "Share Price Oracle",
encryptedEmail: hex"",
upkeepContract: address(sharePriceOracle),
gasLimit: sharePriceOracle.UPKEEP_GAS_LIMIT(),
adminAddress: address(this),
triggerType: 0,
checkData: hex"",
triggerConfig: hex"",
offchainConfig: hex"",
amount: 10e18
});
LINK.safeApprove(automationRegistrarV2, type(uint256).max);
deal(address(LINK), address(this), type(uint128).max);

// Create 6 phony upkeeps we will try to use in `handlePendingUpkeep`
uint256[] memory phonyIds = new uint256[](6);
params.upkeepContract = address(this);
phonyIds[0] = registrarV2.registerUpkeep(params);

params.upkeepContract = address(sharePriceOracle);
params.gasLimit = 100_000;
phonyIds[1] = registrarV2.registerUpkeep(params);

params.gasLimit = sharePriceOracle.UPKEEP_GAS_LIMIT();
params.adminAddress = address(1);
phonyIds[2] = registrarV2.registerUpkeep(params);

params.adminAddress = address(this);
params.triggerType = 1;
phonyIds[3] = registrarV2.registerUpkeep(params);

params.triggerType = 0;
params.checkData = hex"01";
phonyIds[4] = registrarV2.registerUpkeep(params);

params.checkData = hex"";
params.offchainConfig = hex"01";
phonyIds[5] = registrarV2.registerUpkeep(params);

params.offchainConfig = hex"";

// Set auto-approval to false, so our share price oracle upkeep is made pending.
vm.startPrank(registrarV2.owner());
registrarV2.setTriggerConfig(0, IRegistrar.AutoApproveType.DISABLED, 500);
vm.stopPrank();

// Now try creating our share price oracle upkeep.
uint96 initialUpkeepFunds = 10e18;
deal(address(LINK), address(this), initialUpkeepFunds);
LINK.safeApprove(address(sharePriceOracle), initialUpkeepFunds);
sharePriceOracle.initialize(initialUpkeepFunds);

assertTrue(sharePriceOracle.pendingUpkeepParamHash() != bytes32(0), "Pending param hash should be set");

// Trying to call initialize again should fail.
vm.expectRevert(
bytes(abi.encodeWithSelector(ERC4626SharePriceOracle.ERC4626SharePriceOracle__AlreadyInitialized.selector))
);
sharePriceOracle.initialize(initialUpkeepFunds);

// Try using the phony upkeeps to call `handlePendingUpkeep`.
for (uint256 i; i < phonyIds.length; ++i) {
vm.expectRevert(
bytes(
abi.encodeWithSelector(ERC4626SharePriceOracle.ERC4626SharePriceOracle__ParamHashDiffers.selector)
)
);
sharePriceOracle.handlePendingUpkeep(phonyIds[i]);
}

// Approve pending upkeep.
bytes32 hash = keccak256(
abi.encode(
address(sharePriceOracle),
sharePriceOracle.UPKEEP_GAS_LIMIT(),
address(this),
0,
hex"",
hex"",
hex""
)
);
vm.startPrank(registrarV2.owner());

registrarV2.approve(
"Simple Aave Cellar V0.0 Share Price Oracle",
address(sharePriceOracle),
sharePriceOracle.UPKEEP_GAS_LIMIT(),
address(this),
0,
hex"",
hex"",
hex"",
hash
);
vm.stopPrank();

// This can be pulled from stack trace, if you intentionally add a failing require statement right after the approve call.
uint256 approvedUpkeepId = 28590879878797250925087462897928979949870272636475581605666189596973513302011;

// Call `handlePendingUpkeep` with it.
sharePriceOracle.handlePendingUpkeep(approvedUpkeepId);

assertTrue(
sharePriceOracle.automationForwarder() == registryV2.getForwarder(approvedUpkeepId),
"Automation forwarder should be set."
);

// Make sure we can not call `handlePendingUpkeep` again.
vm.expectRevert(
bytes(
abi.encodeWithSelector(
ERC4626SharePriceOracle.ERC4626SharePriceOracle__NoPendingUpkeepToHandle.selector
)
)
);
sharePriceOracle.handlePendingUpkeep(approvedUpkeepId);

// Make sure we can't call initialize again.
vm.expectRevert(
bytes(abi.encodeWithSelector(ERC4626SharePriceOracle.ERC4626SharePriceOracle__AlreadyInitialized.selector))
);
sharePriceOracle.initialize(initialUpkeepFunds);
}

function _passTimeAlterSharePriceAndUpkeepForWethCellar(uint256 timeToPass, uint256 sharePriceMultiplier) internal {
vm.warp(block.timestamp + timeToPass);
wethMockFeed.setMockUpdatedAt(block.timestamp);
Expand Down

0 comments on commit 4ee541d

Please sign in to comment.