Skip to content

Commit aec5f7d

Browse files
xdaniortegaxdaniortega
committed
feat: Improved error handling with custom errors.
1 parent bbf156c commit aec5f7d

File tree

3 files changed

+80
-63
lines changed

3 files changed

+80
-63
lines changed

‎src/TransactionManager.sol‎

Lines changed: 52 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import "./interfaces/ILLMOracle.sol";
1616
* - Slashing mechanism for false challenges
1717
*/
1818
contract TransactionManager {
19-
// Events
2019
event ProposalSubmitted(bytes32 indexed proposalId, string transaction, address indexed submitter);
2120
event ProposalOptimisticallyApproved(bytes32 indexed proposalId);
2221
event ProposalChallenged(bytes32 indexed proposalId, address indexed challenger);
@@ -27,7 +26,28 @@ contract TransactionManager {
2726
event ChallengeResolved(bytes32 indexed proposalId, bool approved, uint256 yesVotes, uint256 noVotes);
2827
event ValidatorSlashed(bytes32 indexed proposalId, address indexed challenger, uint256 amount, bool wasHonest);
2928

30-
// State variables
29+
error InvalidValidatorFactory();
30+
error InvalidLLMOracle();
31+
error EmptyTransaction();
32+
error ProposalAlreadyExists();
33+
error ChallengePeriodExpired();
34+
error NotEnoughValidators();
35+
error NotAValidator();
36+
error AlreadyVoted();
37+
error InvalidSignature();
38+
error ChallengePeriodNotEnded();
39+
error UseResolveChallengeForVotingProposals();
40+
error InvalidProposalStateForFinalization();
41+
error VotingPeriodNotEnded();
42+
error InvalidProposalState();
43+
error ProposalNotFound();
44+
error NotASelectedValidator();
45+
error AlreadySigned();
46+
error VotingPeriodExpired();
47+
error NotInVotingPeriod();
48+
error InvalidSignatureLength();
49+
error InvalidSignatureV();
50+
3151
ValidatorFactory public validatorFactory;
3252
ILLMOracle public llmOracle;
3353

@@ -39,14 +59,12 @@ contract TransactionManager {
3959
mapping(bytes32 => address[]) public proposalVoters;
4060
uint256 public proposalCount;
4161

42-
// Constants
4362
uint256 public constant CHALLENGE_PERIOD = 10; // blocks - as per requirement
4463
uint256 public constant VOTING_PERIOD = 30; // blocks for voting after challenge
4564
uint256 public constant REQUIRED_SIGNATURES = 3; // 3 out of 5 validators
4665
uint256 public constant VALIDATOR_SET_SIZE = 5; // top 5 validators for consensus
4766
uint256 public constant SLASH_PERCENTAGE = 10; // 10% slash for false challenges
4867

49-
// Proposal states
5068
enum ProposalState {
5169
Proposed, // Just submitted
5270
OptimisticApproved, // Enough signatures, optimistically approved
@@ -56,7 +74,6 @@ contract TransactionManager {
5674
Reverted // Proposal was invalid/rejected
5775
}
5876

59-
// Structs
6077
struct Proposal {
6178
bytes32 proposalId;
6279
string transaction;
@@ -75,8 +92,8 @@ contract TransactionManager {
7592
}
7693

7794
constructor(address _validatorFactory, address _llmOracle) {
78-
require(_validatorFactory != address(0), "Invalid validator factory");
79-
require(_llmOracle != address(0), "Invalid LLM oracle");
95+
if (_validatorFactory == address(0)) revert InvalidValidatorFactory();
96+
if (_llmOracle == address(0)) revert InvalidLLMOracle();
8097

8198
validatorFactory = ValidatorFactory(_validatorFactory);
8299
llmOracle = ILLMOracle(_llmOracle);
@@ -88,14 +105,14 @@ contract TransactionManager {
88105
* @return proposalId Unique identifier for the proposal
89106
*/
90107
function submitProposal(string calldata transaction) external returns (bytes32 proposalId) {
91-
require(bytes(transaction).length > 0, "Empty transaction");
108+
if (bytes(transaction).length == 0) revert EmptyTransaction();
92109

93110
proposalId = keccak256(abi.encodePacked(transaction, block.timestamp, msg.sender));
94-
require(proposals[proposalId].proposalId == bytes32(0), "Proposal already exists");
111+
if (proposals[proposalId].proposalId != bytes32(0)) revert ProposalAlreadyExists();
95112

96113
// Get top validators for this proposal
97114
address[] memory topValidators = _getTopValidators();
98-
require(topValidators.length >= REQUIRED_SIGNATURES, "Not enough validators");
115+
if (topValidators.length < REQUIRED_SIGNATURES) revert NotEnoughValidators();
99116

100117
// Perform LLM validation using external oracle
101118
bool llmResult = llmOracle.validateTransaction(transaction);
@@ -132,18 +149,18 @@ contract TransactionManager {
132149
*/
133150
function signProposal(bytes32 proposalId, bytes calldata signature) external {
134151
Proposal storage proposal = proposals[proposalId];
135-
require(proposal.proposalId != bytes32(0), "Proposal not found");
136-
require(proposal.state == ProposalState.Proposed, "Invalid proposal state");
137-
require(block.number <= proposal.challengeDeadline, "Challenge period expired");
138-
require(!hasValidatorSigned[proposalId][msg.sender], "Already signed");
152+
if (proposal.proposalId == bytes32(0)) revert ProposalNotFound();
153+
if (proposal.state != ProposalState.Proposed) revert InvalidProposalState();
154+
if (block.number > proposal.challengeDeadline) revert ChallengePeriodExpired();
155+
if (hasValidatorSigned[proposalId][msg.sender]) revert AlreadySigned();
139156

140157
// Verify that sender is one of the selected validators for this proposal
141-
require(_isSelectedValidator(proposalId, msg.sender), "Not a selected validator");
158+
if (!_isSelectedValidator(proposalId, msg.sender)) revert NotASelectedValidator();
142159

143160
// Verify the signature
144161
bytes32 messageHash = _getProposalHash(proposalId, proposal.transaction);
145162
address recoveredSigner = _recoverSigner(messageHash, signature);
146-
require(recoveredSigner == msg.sender, "Invalid signature");
163+
if (recoveredSigner != msg.sender) revert InvalidSignature();
147164

148165
// Record the signature
149166
hasValidatorSigned[proposalId][msg.sender] = true;
@@ -166,10 +183,10 @@ contract TransactionManager {
166183
*/
167184
function challengeProposal(bytes32 proposalId) external {
168185
Proposal storage proposal = proposals[proposalId];
169-
require(proposal.proposalId != bytes32(0), "Proposal not found");
170-
require(proposal.state == ProposalState.OptimisticApproved, "Cannot challenge");
171-
require(block.number <= proposal.challengeDeadline, "Challenge period expired");
172-
require(validatorFactory.isActiveValidator(msg.sender), "Not a validator");
186+
if (proposal.proposalId == bytes32(0)) revert ProposalNotFound();
187+
if (proposal.state != ProposalState.OptimisticApproved) revert InvalidProposalState();
188+
if (block.number > proposal.challengeDeadline) revert ChallengePeriodExpired();
189+
if (!validatorFactory.isActiveValidator(msg.sender)) revert NotAValidator();
173190

174191
proposal.state = ProposalState.Voting;
175192
proposal.challenger = msg.sender;
@@ -187,16 +204,16 @@ contract TransactionManager {
187204
*/
188205
function submitVote(bytes32 proposalId, bool support, bytes calldata signature) external {
189206
Proposal storage proposal = proposals[proposalId];
190-
require(proposal.proposalId != bytes32(0), "Proposal not found");
191-
require(proposal.state == ProposalState.Voting, "Not in voting period");
192-
require(block.number <= proposal.votingDeadline, "Voting period expired");
193-
require(validatorFactory.isActiveValidator(msg.sender), "Not a validator");
194-
require(!hasValidatorVoted[proposalId][msg.sender], "Already voted");
207+
if (proposal.proposalId == bytes32(0)) revert ProposalNotFound();
208+
if (proposal.state != ProposalState.Voting) revert InvalidProposalState();
209+
if (block.number > proposal.votingDeadline) revert VotingPeriodExpired();
210+
if (!validatorFactory.isActiveValidator(msg.sender)) revert NotAValidator();
211+
if (hasValidatorVoted[proposalId][msg.sender]) revert AlreadyVoted();
195212

196213
// Verify the vote signature
197214
bytes32 voteHash = _getVoteHash(proposalId, support);
198215
address recoveredSigner = _recoverSigner(voteHash, signature);
199-
require(recoveredSigner == msg.sender, "Invalid vote signature");
216+
if (recoveredSigner != msg.sender) revert InvalidSignature();
200217

201218
// Record the vote
202219
hasValidatorVoted[proposalId][msg.sender] = true;
@@ -218,9 +235,9 @@ contract TransactionManager {
218235
*/
219236
function resolveChallenge(bytes32 proposalId) external {
220237
Proposal storage proposal = proposals[proposalId];
221-
require(proposal.proposalId != bytes32(0), "Proposal not found");
222-
require(proposal.state == ProposalState.Voting, "Not in voting state");
223-
require(block.number > proposal.votingDeadline, "Voting period not ended");
238+
if (proposal.proposalId == bytes32(0)) revert ProposalNotFound();
239+
if (proposal.state != ProposalState.Voting) revert InvalidProposalState();
240+
if (block.number <= proposal.votingDeadline) revert VotingPeriodNotEnded();
224241

225242
uint256 totalVotes = proposal.yesVotes + proposal.noVotes;
226243
bool approved = false;
@@ -254,21 +271,21 @@ contract TransactionManager {
254271
*/
255272
function finalizeProposal(bytes32 proposalId) external {
256273
Proposal storage proposal = proposals[proposalId];
257-
require(proposal.proposalId != bytes32(0), "Proposal not found");
274+
if (proposal.proposalId == bytes32(0)) revert ProposalNotFound();
258275

259276
bool approved = false;
260277

261278
if (proposal.state == ProposalState.OptimisticApproved) {
262-
require(block.number > proposal.challengeDeadline, "Challenge period not ended");
279+
if (block.number <= proposal.challengeDeadline) revert ChallengePeriodNotEnded();
263280
// No challenge during period - approve
264281
approved = true;
265282
proposal.executed = true;
266283
proposal.state = ProposalState.Finalized;
267284
} else if (proposal.state == ProposalState.Voting) {
268285
// Should use resolveChallenge instead
269-
revert("Use resolveChallenge for voting proposals");
286+
revert UseResolveChallengeForVotingProposals();
270287
} else {
271-
revert("Invalid proposal state for finalization");
288+
revert InvalidProposalStateForFinalization();
272289
}
273290

274291
emit ProposalFinalized(proposalId, approved);
@@ -363,7 +380,7 @@ contract TransactionManager {
363380
* @return signer Recovered signer address
364381
*/
365382
function _recoverSigner(bytes32 messageHash, bytes memory signature) internal pure returns (address) {
366-
require(signature.length == 65, "Invalid signature length");
383+
if (signature.length != 65) revert InvalidSignatureLength();
367384

368385
bytes32 r;
369386
bytes32 s;
@@ -379,7 +396,7 @@ contract TransactionManager {
379396
v += 27;
380397
}
381398

382-
require(v == 27 || v == 28, "Invalid signature v value");
399+
if (v != 27 && v != 28) revert InvalidSignatureV();
383400

384401
return ecrecover(messageHash, v, r, s);
385402
}

‎src/ValidatorFactory.sol‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ import "./ValidatorLogic.sol";
1818
*/
1919
contract ValidatorFactory is ReentrancyGuard {
2020
using SafeERC20 for IERC20;
21-
// Events
2221
event ValidatorCreated(address indexed validator, address indexed proxy, uint256 stake);
2322
event ValidatorRemoved(address indexed validator);
2423
event ConsensusModuleUpdated(address indexed oldModule, address indexed newModule);
@@ -29,7 +28,18 @@ contract ValidatorFactory is ReentrancyGuard {
2928
event StakingPositionCreated(address indexed validator, uint256 positionId, uint256 amount);
3029
event StakingPositionClosed(address indexed validator, uint256 positionId, uint256 amount);
3130

32-
// State variables
31+
error SenderNotValidator();
32+
error InsufficientStakeAmount();
33+
error AlreadyValidator();
34+
error MaxValidatorsReached();
35+
error TransferFailed();
36+
error AmountExceedsTotalStake();
37+
error PositionNotFound();
38+
error PositionNotActive();
39+
error NotAValidator();
40+
error ValidatorProxyNotFound();
41+
error SlashAmountExceedsStake();
42+
3343
UpgradeableBeacon public beacon;
3444
IERC20 public stakingToken;
3545
uint256 public constant WITHDRAW_COOLDOWN_PERIOD = 1 days;
@@ -43,16 +53,6 @@ contract ValidatorFactory is ReentrancyGuard {
4353
mapping(address => bool) public isValidator;
4454
address[] public validators;
4555

46-
// Custom errors
47-
error SenderNotValidator();
48-
error InsufficientStakeAmount();
49-
error AlreadyValidator();
50-
error MaxValidatorsReached();
51-
error TransferFailed();
52-
error AmountExceedsTotalStake();
53-
error PositionNotFound();
54-
error PositionNotActive();
55-
5656
modifier onlyValidator() {
5757
if (!isValidator[msg.sender]) {
5858
revert SenderNotValidator();
@@ -254,14 +254,14 @@ contract ValidatorFactory is ReentrancyGuard {
254254
function slashValidator(address validator, uint256 slashAmount, string memory reason) external {
255255
// Only allow TransactionManager to slash (in a real implementation, this would be controlled)
256256
// For now, allow any caller for testing purposes
257-
require(isValidator[validator], "Not a validator");
257+
if (!isValidator[validator]) revert NotAValidator();
258258

259259
address proxy = validatorToProxy[validator];
260-
require(proxy != address(0), "Validator proxy not found");
260+
if (proxy == address(0)) revert ValidatorProxyNotFound();
261261

262262
ValidatorLogic validatorLogic = ValidatorLogic(proxy);
263263
uint256 currentStake = validatorLogic.getStakeAmount();
264-
require(currentStake >= slashAmount, "Slash amount exceeds stake");
264+
if (currentStake < slashAmount) revert SlashAmountExceedsStake();
265265

266266
// Perform the slashing by calling unstake on the ValidatorLogic
267267
validatorLogic.unstake(slashAmount);

‎test/TransactionManager.t.sol‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ contract TransactionManagerTest is Test {
187187

188188
function test_RevertWhen_SubmitProposal_EmptyTransaction() public {
189189
vm.prank(proposer);
190-
vm.expectRevert("Empty transaction");
190+
vm.expectRevert(TransactionManager.EmptyTransaction.selector);
191191
transactionManager.submitProposal("");
192192
}
193193

@@ -197,7 +197,7 @@ contract TransactionManagerTest is Test {
197197
transactionManager.submitProposal(TEST_TRANSACTION);
198198

199199
// Try to submit same proposal again
200-
vm.expectRevert("Proposal already exists");
200+
vm.expectRevert(TransactionManager.ProposalAlreadyExists.selector);
201201
transactionManager.submitProposal(TEST_TRANSACTION);
202202

203203
vm.stopPrank();
@@ -259,7 +259,7 @@ contract TransactionManagerTest is Test {
259259
bytes memory invalidSignature = abi.encodePacked(bytes32(0), bytes32(0), uint8(27));
260260

261261
vm.prank(validator1);
262-
vm.expectRevert("Invalid signature");
262+
vm.expectRevert(TransactionManager.InvalidSignature.selector);
263263
transactionManager.signProposal(proposalId, invalidSignature);
264264
}
265265

@@ -281,7 +281,7 @@ contract TransactionManagerTest is Test {
281281
vm.startPrank(validator1);
282282
transactionManager.signProposal(proposalId, signature);
283283

284-
vm.expectRevert("Already signed");
284+
vm.expectRevert(TransactionManager.AlreadySigned.selector);
285285
transactionManager.signProposal(proposalId, signature);
286286
vm.stopPrank();
287287
}
@@ -358,7 +358,7 @@ contract TransactionManagerTest is Test {
358358
// Don't get optimistic approval
359359

360360
vm.prank(validator1);
361-
vm.expectRevert("Cannot challenge");
361+
vm.expectRevert(TransactionManager.InvalidProposalState.selector);
362362
transactionManager.challengeProposal(proposalId);
363363
}
364364

@@ -371,7 +371,7 @@ contract TransactionManagerTest is Test {
371371
vm.roll(block.number + 11);
372372

373373
vm.prank(validator1);
374-
vm.expectRevert("Challenge period expired");
374+
vm.expectRevert(TransactionManager.ChallengePeriodExpired.selector);
375375
transactionManager.challengeProposal(proposalId);
376376
}
377377

@@ -440,7 +440,7 @@ contract TransactionManagerTest is Test {
440440
bytes memory voteSignature = abi.encodePacked(r, s, v);
441441

442442
vm.prank(validator1);
443-
vm.expectRevert("Not in voting period");
443+
vm.expectRevert(TransactionManager.InvalidProposalState.selector);
444444
transactionManager.submitVote(proposalId, support, voteSignature);
445445
}
446446

@@ -465,7 +465,7 @@ contract TransactionManagerTest is Test {
465465
bytes memory voteSignature = abi.encodePacked(r, s, v);
466466

467467
vm.prank(validator2);
468-
vm.expectRevert("Already voted");
468+
vm.expectRevert(TransactionManager.AlreadyVoted.selector);
469469
transactionManager.submitVote(proposalId, false, voteSignature);
470470
}
471471

@@ -601,7 +601,7 @@ contract TransactionManagerTest is Test {
601601
_signProposalWithValidators(proposalId, TEST_TRANSACTION, 3);
602602

603603
// Try to finalize before challenge period ends
604-
vm.expectRevert("Challenge period not ended");
604+
vm.expectRevert(TransactionManager.ChallengePeriodNotEnded.selector);
605605
transactionManager.finalizeProposal(proposalId);
606606
}
607607

@@ -615,7 +615,7 @@ contract TransactionManagerTest is Test {
615615

616616
vm.roll(block.number + 11);
617617

618-
vm.expectRevert("Use resolveChallenge for voting proposals");
618+
vm.expectRevert(TransactionManager.UseResolveChallengeForVotingProposals.selector);
619619
transactionManager.finalizeProposal(proposalId);
620620
}
621621

@@ -738,7 +738,7 @@ contract TransactionManagerTest is Test {
738738
transactionManager.challengeProposal(proposalId);
739739

740740
// Try to resolve before voting period ends
741-
vm.expectRevert("Voting period not ended");
741+
vm.expectRevert(TransactionManager.VotingPeriodNotEnded.selector);
742742
transactionManager.resolveChallenge(proposalId);
743743
}
744744

@@ -747,7 +747,7 @@ contract TransactionManagerTest is Test {
747747
bytes32 proposalId = transactionManager.submitProposal(TEST_TRANSACTION);
748748

749749
// Try to finalize before optimistic approval
750-
vm.expectRevert("Invalid proposal state for finalization");
750+
vm.expectRevert(TransactionManager.InvalidProposalStateForFinalization.selector);
751751
transactionManager.finalizeProposal(proposalId);
752752
}
753753

@@ -800,7 +800,7 @@ contract TransactionManagerTest is Test {
800800
assertFalse(transactionManager.isInVotingPeriod(proposalId));
801801

802802
// Try to vote after period ends
803-
vm.expectRevert("Voting period expired");
803+
vm.expectRevert(TransactionManager.VotingPeriodExpired.selector);
804804
_submitVote(proposalId, validator2, true);
805805
}
806806

0 commit comments

Comments
 (0)