Refactor PufferProtocol to keep <24KB by eladiosch · Pull Request #119 · PufferFinance/puffer-contracts

eladiosch · 2025-07-08T17:52:00Z

Before this PR, due to the changes for Pectra and the Validator Tokens => Validation Time rework, the size of the PufferProtocol contract had reached 28.5 KB, way over the limit.

In this PR I've implemented a contract that offloads a big part of the logic of the PufferProtocol contract. The way to call this contract is via delegatecall so the storage is the PufferProtocol's. A new base contract has been created so both PufferProtocol and PufferProtocolLogic share constants, immutables, storage, errors, and other common stuff.

The contract size achieved is 18.2 KB, so there's more than enough space to continue with the Pectra implementation and add some more features.

The PR includes some changes to the tests but they're not related to the PufferProtocolLogic contract. They were adapted so the previous tests could work with the changes made regarding the signatures. The changes to the protocol and protocol logic contract have not altered the tests and they still pass

…o logic contract

bxmmm1 · 2025-07-10T06:07:04Z

mainnet-contracts/src/PufferProtocol.sol

        // will continue. If the `msg.sender` did a `VALIDATOR_TICKET.approve(spender, amount)` before calling this
        // And the spender is `msg.sender` the Permit call will revert, but the overall transaction will succeed
-        _callPermit(address(VALIDATOR_TICKET), permit);
+        _callPermit(address(_VALIDATOR_TICKET), permit);


can remove _callPermit and just inline the code here, should save some gas. Or remove it completely and just do the .approve flow since VT is deprecated.

Removed the Permit flow completely

bxmmm1 · 2025-07-10T06:17:20Z

mainnet-contracts/src/PufferProtocolLogic.sol

+            require(validatorSrc.node == msg.sender && validatorSrc.status == Status.ACTIVE, InvalidValidator());
+            srcPubkeys[i] = validatorSrc.pubKey;
+            validatorTarget = $.validators[moduleName][targetIndices[i]];
+            require(validatorTarget.node == msg.sender && validatorTarget.status == Status.ACTIVE, InvalidValidator());


Did we doument anywhere that scenario, where Source consolidates into a bad target and NoOp loses money?

No further work was done in relation to consolidation. We mentioned removing numBatches from Validator info, I still need to check if that's reasonable or breaks some other flow. We'll probably need to sync again on consolidation

bxmmm1 · 2025-07-10T08:28:57Z

mainnet-contracts/src/GuardianModule.sol

     * @inheritdoc IGuardianModule
     */
-    function validateWithdrawalRequest(bytes[] calldata eoaSignatures, bytes32 messageHash) external view {
+    function validateWithdrawalRequest(


This would require us to deploy a new GuardianModule, we could have this message logic in the protocol, and use

GUARDIAN_MODULE.validateGuardiansEOASignatures

That way, we wouldn't need to re-audit/redeploy

GuardianModule

We need to re-deploy the GuardianModule anyways since the struct StoppedValidatorsInfo has changes. It was commented in the Pectra PR => #115 (comment)

If we move that signature logic to the protocol, we might not need to, please check my big comment on the PR

bxmmm1 · 2025-07-10T12:00:01Z

You can use the fallback of PufferProtocol to delegateCall to another logic as long as you have the same interface in the logic contract; everything will work.

E.G.  ```  fallback() external payable {
address implementation = _getPufferProtocolStorage().pufferProtocolLogic;

    assembly {
        // Copy msg.data. We take full control of memory in this inline assembly
        // block because it will not return to Solidity code. We overwrite the
        // Solidity scratch pad at memory position 0.
        calldatacopy(0, 0, calldatasize())

        // Call the implementation.
        // out and outsize are 0 because we don't know the size yet.
        let result := delegatecall(gas(), implementation, 0, calldatasize(), 0, 0)

        // Copy the returned data.
        returndatacopy(0, 0, returndatasize())

        switch result
        // delegatecall returns 0 on error.
        case 0 {
            revert(0, returndatasize())
        }
        default {
            return(0, returndatasize())
        }
    }
}

```  And then in the Logic contract, instead of _ method can just have a normal `function depositValidationTime(EpochsValidatedSignature memory epochsValidatedSignature)` signature.  Then, from the PufferProtocol, you can remove every method that uses _delegateCall.  That way, the flow will be:  User -> Proxy -> PufferProtocol (if method signature is not found) trigger `fallback` -> it then tries to find the method on PufferProtocolLogic   In the last in line (PufferProtocolLogic) you can have fallback revert so that it returns an error if the method is not found in thorugh the contract calls chain.

Maybe it wouldn’t be bad to compose PufferProtocol out of multiple interfaces.
E.G.
IPufferProtoclEvents
IPufferProtocolErrors
IPufferProtocolViewFunctions

Also, since this PufferProtocolLogic is being called by delegateCall, it doesn’t need to have the constants in it as well. You can have the constants from (PufferProtocolBase) only in the PufferProtocol.
But in our case, we don’t need to reduce bytecode even more.  Since this PR requires us to redeploy GuardianModule, we should remove unused functions from there. Or delete it and replace it with a function _validateSignatures somewhere in the protocol, and then have the protocol call it with a message + signer’s address(backend), no need for array or fancy things, since we never used more than one signature.. 

codecov · 2025-07-10T12:02:36Z

Codecov Report

❌ Patch coverage is 75.19685% with 63 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mainnet-contracts/src/PufferProtocolLogic.sol	75.13%	38 Missing and 8 partials ⚠️
mainnet-contracts/src/PufferProtocol.sol	73.21%	12 Missing and 3 partials ⚠️
mainnet-contracts/src/PufferProtocolBase.sol	84.61%	0 Missing and 2 partials ⚠️

Files with missing lines	Coverage Δ
mainnet-contracts/src/ProtocolSignatureNonces.sol	`100.00% <ø> (ø)`
mainnet-contracts/src/PufferProtocolBase.sol	`84.61% <84.61%> (ø)`
mainnet-contracts/src/PufferProtocol.sol	`88.02% <73.21%> (+8.39%)`	⬆️
mainnet-contracts/src/PufferProtocolLogic.sol	`75.13% <75.13%> (ø)`

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

bxmmm1 · 2025-07-15T09:32:28Z

mainnet-contracts/src/PufferProtocolLogic.sol

+            InvalidETHAmount()
+        );
+
+        epochsValidatedSignature.functionSelector = _FUNCTION_SELECTOR_DEPOSIT_VALIDATION_TIME;


You can use IPufferProtocolLogic.depositValidationTime.selector directly.

bxmmm1 · 2025-07-15T09:33:06Z

mainnet-contracts/src/PufferProtocolBase.sol

+     */
+    uint256 internal constant _32_ETH_GWEI = 32 * 10 ** 9;
+
+    bytes32 internal constant _FUNCTION_SELECTOR_REGISTER_VALIDATOR_KEY =


You can get rid of all these internal constants and just inline IPufferProtocolLogic.whatever.selector

I had to create the constants because of the stack too deep. Maybe it isn't necessary anymore. I'll try

bxmmm1 · 2025-07-15T09:46:04Z

mainnet-contracts/src/PufferProtocol.sol


    receive() external payable { }

+    fallback() external payable {


Maybe add a simple explainer.
If a function selector is not found in this contract ...

bxmmm1 · 2025-07-15T09:55:30Z

mainnet-contracts/src/PufferProtocolLogic.sol

+            revert DeadlineExceeded();
+        }
+
+        _GUARDIAN_MODULE.validateBatchWithdrawals(validatorInfos, guardianEOASignatures, deadline);


I'd use _GUARDIAN_MODULE.validateGuardiansEOASignatures. Prepare the message here, in this contract. That way, we won't need to redeploy GuardianModule contract at all

bxmmm1 · 2025-07-15T09:55:51Z

mainnet-contracts/src/PufferProtocolLogic.sol

+            return;
+        }
+
+        _GUARDIAN_MODULE.validateTotalEpochsValidated({


You can do it here as well validateGuardiansEOASignatures

ksatyarth2 · 2025-07-15T13:54:10Z

mainnet-contracts/src/PufferProtocolLogic.sol

+import { EpochsValidatedSignature } from "./struct/Signatures.sol";
+import { InvalidAddress } from "./Errors.sol";
+
+contract PufferProtocolLogic is PufferProtocolBase, IPufferProtocolLogic {


natspec for this contract plz

Done (yesterday)

ksatyarth2 · 2025-07-15T14:19:33Z

mainnet-contracts/src/PufferProtocolLogic.sol

+     * @notice Check IPufferProtocol.withdrawValidationTime
+     * @dev This function should only be called by the PufferProtocol contract through a delegatecall
+     */
+    function withdrawValidationTime(uint96 amount, address recipient) external override {


the natspec says that fn can be called by delegatecall, but since this is a separate contract, can't be this fn directly called by someone? we are not checking anywhere that the caller should be PufferProtocol, or msg.sender == address(this) since this is delegate call. Am I missing something here?

Think of a Proxy pattern. If you interact with the implementation contract you're not accessing the proxy storage, right? This is the same. If you call the PufferLogicContract you'll be reading/writing the PufferLogicContract storage, not the PufferProtocol one, which is the important one. So no checks are needed.

However I just noticed due to your comment that since we're not checking the restricted in the original PufferProtocol anymore due to the last refactor (using just the fallback instead of the old functions), we need to add the check somehow. Probably will add it in the fallback function directly

ksatyarth2 · 2025-07-15T14:30:53Z

mainnet-contracts/src/PufferProtocolLogic.sol

+    /**
+     * @notice Check IPufferProtocol.withdrawValidationTime
+     * @dev This function should only be called by the PufferProtocol contract through a delegatecall
+     */


add amount>0 and recipient!=o validation check to avoid informational issue from auditors :D

Done (yesterday)

ksatyarth2 · 2025-07-15T14:39:05Z

mainnet-contracts/src/PufferProtocolLogic.sol

+            InvalidETHAmount()
+        );
+
+        epochsValidatedSignature.functionSelector = IPufferProtocolLogic.depositValidationTime.selector;


why are we modifying the memory param here?

We have that struct (EpochsValidatedSignature) that contains all info regarding the epochsvalidated. One of that items is the function selector, which we need to get the proper nonce (see ProtocolSignatureNonces contract), that are now separated by function selector, in order to avoid DOS by blocking a nonce used by another function. This is, we previously used just nonce by address, but that way a user could send a tx that would block a nonce so it couldn't be used for another flow.

Well this function receives that struct, but the function selector doesn't need to be set by the user (they could tamper it), so it's set in the function body. It's also set in other places where the struct is created (like registerValidatorKey or batchHandleWithdrawals)

ksatyarth2 · 2025-07-15T14:40:38Z

mainnet-contracts/src/PufferProtocolLogic.sol

+        if (block.timestamp > epochsValidatedSignature.deadline) {
+            revert DeadlineExceeded();
+        }


few places we are using require and other places revert - inconsistent

True, but that was also the case before this PR. IIRC we're not enforcing consistency with require/revert. I think this came up in a different PR

Yeah, we got no consensus on this. I think that this uses slightly less gas. Just use whatever you feel is more readable to human.
On another note, this is 3 LOC when auditing the contract, while require would likely be 1 :D
I don't like inline `if(something) revert error() without the {}. I need to re-wire my brain to read it properly.

When I was optimizing I looked it up and it's equivalent. I usually prefer the require approach. And the LOC thing is valid, I'll see where it makes sense to change it

…tract

bxmmm1 · 2025-07-16T07:14:48Z

mainnet-contracts/src/PufferProtocolLogic.sol

+        payable
+        override
+    {
+        if (srcIndices.length == 0) {


Maybe we could replace some of these if revert with require for less LOC and readability, wyt?

Agreed. I just wrote that in the other comment :D

bxmmm1 · 2025-07-16T07:16:10Z

mainnet-contracts/src/PufferProtocolLogic.sol

+        if (block.timestamp > deadline) {
+            revert DeadlineExceeded();
+        }


This can probably be a modifier, it shouldn't add too much bytecode, it will be a little more readable. I've seen it 3 times in this contract

mainnet-contracts/src/PufferProtocol.sol

mainnet-contracts/src/PufferProtocolLogic.sol

mainnet-contracts/src/PufferProtocol.sol

Co-authored-by: Benjamin <[email protected]>

eladiosch added 7 commits July 3, 2025 12:14

Removed unused function

86afa4b

Created PufferLogic contract and adapted existing ones (WIP)

202c7f1

Adapted previous tests and improved signatures

0b4a2b8

Changed ProtocolLogic structure. Moved more logic to ext contract (WIP)

37fff25

Fixed and simplify delegatecalls

27dea10

Refactor delegatecalls and small fixes

50536f7

Refactor PufferConstants to PufferProtocolBase and moved more logic t…

7b53076

…o logic contract

eladiosch self-assigned this Jul 8, 2025

forge fmt

ec3e5fe

eladiosch requested review from bxmmm1, ksatyarth2 and tdenisenko July 8, 2025 17:55

bxmmm1 reviewed Jul 10, 2025

View reviewed changes

Removed flow to deposit VT with Permit

9d2251f

eladiosch and others added 4 commits July 10, 2025 17:05

Added tests to increase coverage

c104f39

forge fmt

48c5b53

Refactored interfaces and adapted tests (protocol WIP)

4e44801

Changed approach of the delegatecall to logic contract

eab8cfc

bxmmm1 reviewed Jul 15, 2025

View reviewed changes

eladiosch added 2 commits July 15, 2025 12:07

Removed selector constants

495ac9d

Added some natspec to the fallback function

bf0c713

ksatyarth2 reviewed Jul 15, 2025

View reviewed changes

eladiosch and others added 4 commits July 15, 2025 17:59

Moved logic from GuardianModule to Protocol to avoid re-deploying con…

a19d9a0

…tract

forge fmt

5501856

Added basic checks to withdrawValidationTime and added missing natspec

f3dafab

forge fmt

8d546b1

bxmmm1 reviewed Jul 16, 2025

View reviewed changes

mainnet-contracts/src/PufferProtocol.sol Outdated Show resolved Hide resolved

bxmmm1 reviewed Jul 16, 2025

View reviewed changes

mainnet-contracts/src/PufferProtocol.sol Outdated Show resolved Hide resolved

bxmmm1 reviewed Jul 16, 2025

View reviewed changes

mainnet-contracts/src/PufferProtocol.sol Outdated Show resolved Hide resolved

bxmmm1 reviewed Jul 16, 2025

View reviewed changes

mainnet-contracts/src/PufferProtocol.sol Outdated Show resolved Hide resolved

bxmmm1 reviewed Jul 16, 2025

View reviewed changes

mainnet-contracts/src/PufferProtocolLogic.sol Show resolved Hide resolved

bxmmm1 reviewed Jul 16, 2025

View reviewed changes

mainnet-contracts/src/PufferProtocol.sol Outdated Show resolved Hide resolved

eladiosch and others added 9 commits July 16, 2025 10:37

Update mainnet-contracts/src/PufferProtocol.sol

522f503

Co-authored-by: Benjamin <[email protected]>

Update mainnet-contracts/src/PufferProtocol.sol

e7fb64e

Co-authored-by: Benjamin <[email protected]>

Update mainnet-contracts/src/PufferProtocol.sol

76a3233

Co-authored-by: Benjamin <[email protected]>

Update mainnet-contracts/src/PufferProtocol.sol

c90f5da

Co-authored-by: Benjamin <[email protected]>

Update mainnet-contracts/src/PufferProtocolLogic.sol

e181e10

Co-authored-by: Benjamin <[email protected]>

Update mainnet-contracts/src/PufferProtocol.sol

fad967d

Co-authored-by: Benjamin <[email protected]>

Using modifier (with require) to check deadlines

b982b2e

Changed reverts to require and refactored to avoid stack too deep

3c16eb3

Added restricted for the Logic contract and improved natspec

3599a63

eladiosch merged commit 59f2899 into feat/pectra-upgrade Jul 28, 2025
3 of 4 checks passed

Comments

Conversation

eladiosch commented Jul 8, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bxmmm1 commented Jul 10, 2025

Uh oh!

codecov bot commented Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ksatyarth2 Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Jul 10, 2025 •

edited

Loading

ksatyarth2 Jul 15, 2025 •

edited

Loading