PythGovernance.t.sol 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739
  1. // SPDX-License-Identifier: Apache 2
  2. // NOTE: These tests were migrated from target_chains/ethereum/contracts/test/pyth.js but exclude the Wormhole-specific tests,
  3. // which remain in the original JavaScript test file.
  4. pragma solidity ^0.8.0;
  5. import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
  6. import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
  7. import "forge-std/Test.sol";
  8. import "@pythnetwork/pyth-sdk-solidity/IPyth.sol";
  9. import "@pythnetwork/pyth-sdk-solidity/PythErrors.sol";
  10. import "@pythnetwork/pyth-sdk-solidity/PythStructs.sol";
  11. import "../contracts/pyth/PythInternalStructs.sol";
  12. import "../contracts/pyth/PythGovernanceInstructions.sol";
  13. import "../contracts/pyth/PythUpgradable.sol";
  14. import "../contracts/pyth/PythGetters.sol";
  15. import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
  16. import "../contracts/wormhole/interfaces/IWormhole.sol";
  17. import "../contracts/wormhole/Implementation.sol";
  18. import "../contracts/wormhole/Setup.sol";
  19. import "../contracts/wormhole/Wormhole.sol";
  20. import "../contracts/wormhole-receiver/WormholeReceiver.sol";
  21. import "../contracts/wormhole-receiver/ReceiverImplementation.sol";
  22. import "../contracts/wormhole-receiver/ReceiverSetup.sol";
  23. import "../contracts/wormhole-receiver/ReceiverGovernanceStructs.sol";
  24. import "../contracts/wormhole-receiver/ReceiverStructs.sol";
  25. import "../contracts/wormhole-receiver/ReceiverGovernance.sol";
  26. import "../contracts/libraries/external/BytesLib.sol";
  27. import "../contracts/pyth/mock/MockUpgradeableProxy.sol";
  28. import "./utils/WormholeTestUtils.t.sol";
  29. import "./utils/PythTestUtils.t.sol";
  30. import "./utils/RandTestUtils.t.sol";
  31. contract PythGovernanceTest is Test, PythTestUtils, PythGovernanceInstructions {
  32. using BytesLib for bytes;
  33. IPyth public pyth;
  34. address constant TEST_SIGNER1 = 0xbeFA429d57cD18b7F8A4d91A2da9AB4AF05d0FBe;
  35. address constant TEST_SIGNER2 = 0x4ba0C2db9A26208b3bB1a50B01b16941c10D76db;
  36. uint16 constant TEST_GOVERNANCE_CHAIN_ID = 1;
  37. bytes32 constant TEST_GOVERNANCE_EMITTER =
  38. 0x0000000000000000000000000000000000000000000000000000000000000011;
  39. uint16 constant TEST_PYTH2_WORMHOLE_CHAIN_ID = 1;
  40. bytes32 constant TEST_PYTH2_WORMHOLE_EMITTER =
  41. 0x71f8dcb863d176e2c420ad6610cf687359612b6fb392e0642b0ca6b1f186aa3b;
  42. uint16 constant TARGET_CHAIN_ID = 2;
  43. function setUp() public {
  44. pyth = IPyth(setUpPyth(new WormholeTestUtils(1)));
  45. }
  46. function testNoOwner() public {
  47. // Check that the ownership is renounced
  48. assertEq(OwnableUpgradeable(address(pyth)).owner(), address(0));
  49. }
  50. function testValidDataSources() public {
  51. assertTrue(
  52. PythGetters(address(pyth)).isValidDataSource(
  53. TEST_PYTH2_WORMHOLE_CHAIN_ID,
  54. TEST_PYTH2_WORMHOLE_EMITTER
  55. )
  56. );
  57. }
  58. function testSetFee() public {
  59. // Set fee to 5000 (5000 = 5 * 10^3)
  60. bytes memory setFeeMessage = abi.encodePacked(
  61. MAGIC,
  62. uint8(GovernanceModule.Target),
  63. uint8(GovernanceAction.SetFee),
  64. TARGET_CHAIN_ID,
  65. uint64(5), // value
  66. uint64(3) // exponent
  67. );
  68. bytes memory vaa = encodeAndSignMessage(
  69. setFeeMessage,
  70. TEST_GOVERNANCE_CHAIN_ID,
  71. TEST_GOVERNANCE_EMITTER,
  72. 1
  73. );
  74. uint oldFee = PythGetters(address(pyth)).singleUpdateFeeInWei();
  75. vm.expectEmit(true, true, true, true);
  76. emit FeeSet(oldFee, 5000);
  77. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  78. assertEq(PythGetters(address(pyth)).singleUpdateFeeInWei(), 5000);
  79. }
  80. function testSetValidPeriod() public {
  81. // Create governance VAA to set valid period to 0
  82. bytes memory data = abi.encodePacked(
  83. MAGIC,
  84. uint8(GovernanceModule.Target),
  85. uint8(GovernanceAction.SetValidPeriod),
  86. TARGET_CHAIN_ID, // Target chain ID
  87. uint64(0) // New valid period
  88. );
  89. bytes memory vaa = encodeAndSignMessage(
  90. data,
  91. TEST_GOVERNANCE_CHAIN_ID,
  92. TEST_GOVERNANCE_EMITTER,
  93. 1
  94. );
  95. uint oldValidPeriod = PythGetters(address(pyth))
  96. .validTimePeriodSeconds();
  97. vm.expectEmit(true, true, true, true);
  98. emit ValidPeriodSet(oldValidPeriod, 0);
  99. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  100. assertEq(PythGetters(address(pyth)).validTimePeriodSeconds(), 0);
  101. }
  102. function testInvalidGovernanceMessage() public {
  103. // Test with wrong magic number
  104. bytes memory data = abi.encodePacked(
  105. bytes4(0x12345678), // Wrong magic
  106. uint8(GovernanceModule.Target),
  107. uint8(GovernanceAction.SetValidPeriod),
  108. uint16(1), // Target chain ID
  109. uint64(0)
  110. );
  111. bytes memory vaa = encodeAndSignMessage(
  112. data,
  113. TEST_GOVERNANCE_CHAIN_ID,
  114. TEST_GOVERNANCE_EMITTER,
  115. 1
  116. );
  117. vm.expectRevert(PythErrors.InvalidGovernanceMessage.selector);
  118. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  119. }
  120. function testInvalidGovernanceTarget() public {
  121. // Test with wrong chain target
  122. bytes memory data = abi.encodePacked(
  123. MAGIC,
  124. uint8(GovernanceModule.Target),
  125. uint8(GovernanceAction.SetValidPeriod),
  126. uint16(3), // Different chain ID for testing invalid target
  127. uint64(0)
  128. );
  129. bytes memory vaa = encodeAndSignMessage(
  130. data,
  131. TEST_GOVERNANCE_CHAIN_ID,
  132. TEST_GOVERNANCE_EMITTER,
  133. 1
  134. );
  135. vm.expectRevert(PythErrors.InvalidGovernanceTarget.selector);
  136. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  137. }
  138. function testInvalidGovernanceDataSource() public {
  139. // Test with wrong emitter
  140. bytes memory data = abi.encodePacked(
  141. MAGIC,
  142. uint8(GovernanceModule.Target),
  143. uint8(GovernanceAction.SetValidPeriod),
  144. TARGET_CHAIN_ID,
  145. uint64(0)
  146. );
  147. bytes memory vaa = encodeAndSignMessage(
  148. data,
  149. TEST_GOVERNANCE_CHAIN_ID,
  150. bytes32(uint256(0x1111)), // Wrong emitter
  151. 1
  152. );
  153. vm.expectRevert(PythErrors.InvalidGovernanceDataSource.selector);
  154. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  155. }
  156. function testSetDataSources() public {
  157. // Create governance VAA to set new data sources
  158. bytes memory data = abi.encodePacked(
  159. MAGIC,
  160. uint8(GovernanceModule.Target),
  161. uint8(GovernanceAction.SetDataSources),
  162. TARGET_CHAIN_ID, // Target chain ID
  163. uint8(1), // Number of data sources
  164. uint16(1), // Chain ID
  165. bytes32(
  166. 0x0000000000000000000000000000000000000000000000000000000000001111
  167. ) // Emitter
  168. );
  169. bytes memory vaa = encodeAndSignMessage(
  170. data,
  171. TEST_GOVERNANCE_CHAIN_ID,
  172. TEST_GOVERNANCE_EMITTER,
  173. 1
  174. );
  175. PythInternalStructs.DataSource[] memory oldDataSources = PythGetters(
  176. address(pyth)
  177. ).validDataSources();
  178. PythInternalStructs.DataSource[]
  179. memory newDataSources = new PythInternalStructs.DataSource[](1);
  180. newDataSources[0] = PythInternalStructs.DataSource(
  181. 1,
  182. 0x0000000000000000000000000000000000000000000000000000000000001111
  183. );
  184. vm.expectEmit(true, true, true, true);
  185. emit DataSourcesSet(oldDataSources, newDataSources);
  186. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  187. // Verify old data source is no longer valid
  188. assertFalse(
  189. PythGetters(address(pyth)).isValidDataSource(
  190. TEST_PYTH2_WORMHOLE_CHAIN_ID,
  191. TEST_PYTH2_WORMHOLE_EMITTER
  192. )
  193. );
  194. // Verify new data source is valid
  195. assertTrue(
  196. PythGetters(address(pyth)).isValidDataSource(
  197. 1,
  198. 0x0000000000000000000000000000000000000000000000000000000000001111
  199. )
  200. );
  201. }
  202. function testSetWormholeAddress() public {
  203. // Deploy a new wormhole contract
  204. address newWormhole = new WormholeTestUtils(1)
  205. .getWormholeReceiverAddr();
  206. // Create governance VAA to set new wormhole address
  207. bytes memory data = abi.encodePacked(
  208. MAGIC,
  209. uint8(GovernanceModule.Target),
  210. uint8(GovernanceAction.SetWormholeAddress),
  211. TARGET_CHAIN_ID, // Target chain ID
  212. newWormhole // New wormhole address
  213. );
  214. bytes memory vaa = encodeAndSignMessage(
  215. data,
  216. TEST_GOVERNANCE_CHAIN_ID,
  217. TEST_GOVERNANCE_EMITTER,
  218. 1
  219. );
  220. address oldWormhole = address(PythGetters(address(pyth)).wormhole());
  221. vm.expectEmit(true, true, true, true);
  222. emit WormholeAddressSet(oldWormhole, newWormhole);
  223. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  224. assertEq(address(PythGetters(address(pyth)).wormhole()), newWormhole);
  225. }
  226. function testSetVerifierAddress() public {
  227. // Deploy a new verifier contract
  228. address newVerifier = new WormholeTestUtils(3)
  229. .getWormholeReceiverAddr();
  230. // Create governance VAA to set new verifier address
  231. bytes memory data = abi.encodePacked(
  232. MAGIC,
  233. uint8(GovernanceModule.Target),
  234. uint8(GovernanceAction.SetVerifierAddress),
  235. TARGET_CHAIN_ID, // Target chain ID
  236. newVerifier // New verifier address
  237. );
  238. bytes memory vaa = encodeAndSignMessage(
  239. data,
  240. TEST_GOVERNANCE_CHAIN_ID,
  241. TEST_GOVERNANCE_EMITTER,
  242. 1
  243. );
  244. address oldVerifier = address(PythGetters(address(pyth)).verifier());
  245. vm.expectEmit(true, true, true, true);
  246. emit VerifierAddressSet(oldVerifier, newVerifier);
  247. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  248. assertEq(address(PythGetters(address(pyth)).verifier()), newVerifier);
  249. }
  250. function testTransferGovernanceDataSource() public {
  251. uint16 newEmitterChain = 2;
  252. bytes32 newEmitterAddress = 0x0000000000000000000000000000000000000000000000000000000000001111;
  253. // Create claim VAA from new governance
  254. bytes memory claimData = abi.encodePacked(
  255. MAGIC,
  256. uint8(GovernanceModule.Target),
  257. uint8(GovernanceAction.RequestGovernanceDataSourceTransfer),
  258. TARGET_CHAIN_ID, // Target chain ID
  259. uint32(1) // New governance index
  260. );
  261. bytes memory claimVaa = encodeAndSignMessage(
  262. claimData,
  263. newEmitterChain,
  264. newEmitterAddress,
  265. 1
  266. );
  267. // Create authorize VAA from current governance
  268. bytes memory authData = abi.encodePacked(
  269. MAGIC,
  270. uint8(GovernanceModule.Target),
  271. uint8(GovernanceAction.AuthorizeGovernanceDataSourceTransfer),
  272. TARGET_CHAIN_ID, // Target chain ID
  273. claimVaa
  274. );
  275. bytes memory authVaa = encodeAndSignMessage(
  276. authData,
  277. TEST_GOVERNANCE_CHAIN_ID,
  278. TEST_GOVERNANCE_EMITTER,
  279. 1
  280. );
  281. PythInternalStructs.DataSource
  282. memory oldDataSource = PythInternalStructs.DataSource(
  283. TEST_GOVERNANCE_CHAIN_ID,
  284. TEST_GOVERNANCE_EMITTER
  285. );
  286. PythInternalStructs.DataSource
  287. memory newDataSource = PythInternalStructs.DataSource(
  288. newEmitterChain,
  289. newEmitterAddress
  290. );
  291. vm.expectEmit(true, true, true, true);
  292. emit GovernanceDataSourceSet(oldDataSource, newDataSource, 1);
  293. PythGovernance(address(pyth)).executeGovernanceInstruction(authVaa);
  294. // Verify old governance can't execute instructions anymore
  295. bytes memory invalidData = abi.encodePacked(
  296. MAGIC,
  297. uint8(GovernanceModule.Target),
  298. uint8(GovernanceAction.SetValidPeriod),
  299. uint16(1), // Wrong chain ID for testing invalid target
  300. uint64(0)
  301. );
  302. bytes memory invalidVaa = encodeAndSignMessage(
  303. invalidData,
  304. TEST_GOVERNANCE_CHAIN_ID,
  305. TEST_GOVERNANCE_EMITTER,
  306. 2
  307. );
  308. vm.expectRevert(PythErrors.InvalidGovernanceDataSource.selector);
  309. PythGovernance(address(pyth)).executeGovernanceInstruction(invalidVaa);
  310. }
  311. function testSequentialGovernanceMessages() public {
  312. // First governance message
  313. bytes memory data1 = abi.encodePacked(
  314. MAGIC,
  315. uint8(GovernanceModule.Target),
  316. uint8(GovernanceAction.SetValidPeriod),
  317. TARGET_CHAIN_ID, // Target chain ID
  318. uint64(10)
  319. );
  320. bytes memory vaa1 = encodeAndSignMessage(
  321. data1,
  322. TEST_GOVERNANCE_CHAIN_ID,
  323. TEST_GOVERNANCE_EMITTER,
  324. 1
  325. );
  326. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa1);
  327. // Second governance message
  328. bytes memory data2 = abi.encodePacked(
  329. MAGIC,
  330. uint8(GovernanceModule.Target),
  331. uint8(GovernanceAction.SetValidPeriod),
  332. TARGET_CHAIN_ID, // Target chain ID
  333. uint64(20)
  334. );
  335. bytes memory vaa2 = encodeAndSignMessage(
  336. data2,
  337. TEST_GOVERNANCE_CHAIN_ID,
  338. TEST_GOVERNANCE_EMITTER,
  339. 2
  340. );
  341. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa2);
  342. // Try to replay first message
  343. vm.expectRevert(PythErrors.OldGovernanceMessage.selector);
  344. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa1);
  345. // Try to replay second message
  346. vm.expectRevert(PythErrors.OldGovernanceMessage.selector);
  347. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa2);
  348. }
  349. function testUpgradeContractWithChainIdZeroIsInvalid() public {
  350. // Deploy a new PythUpgradable contract
  351. PythUpgradable newImplementation = new PythUpgradable();
  352. // Create governance VAA with chain ID 0 (unset)
  353. bytes memory data = abi.encodePacked(
  354. MAGIC,
  355. uint8(GovernanceModule.Target),
  356. uint8(GovernanceAction.UpgradeContract),
  357. uint16(0), // Chain ID 0 (unset)
  358. address(newImplementation) // New implementation address
  359. );
  360. bytes memory vaa = encodeAndSignMessage(
  361. data,
  362. TEST_GOVERNANCE_CHAIN_ID,
  363. TEST_GOVERNANCE_EMITTER,
  364. 1
  365. );
  366. // Should revert with InvalidGovernanceTarget
  367. vm.expectRevert(PythErrors.InvalidGovernanceTarget.selector);
  368. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  369. }
  370. // Helper function to get the second address from event data
  371. function getSecondAddressFromEventData(
  372. bytes memory data
  373. ) internal pure returns (address) {
  374. (, address secondAddr) = abi.decode(data, (address, address));
  375. return secondAddr;
  376. }
  377. function testUpgradeContractShouldWork() public {
  378. // Deploy a new PythUpgradable contract
  379. PythUpgradable newImplementation = new PythUpgradable();
  380. // Create governance VAA to upgrade the contract
  381. bytes memory data = abi.encodePacked(
  382. MAGIC,
  383. uint8(GovernanceModule.Target),
  384. uint8(GovernanceAction.UpgradeContract),
  385. TARGET_CHAIN_ID, // Valid target chain ID
  386. address(newImplementation) // New implementation address
  387. );
  388. bytes memory vaa = encodeAndSignMessage(
  389. data,
  390. TEST_GOVERNANCE_CHAIN_ID,
  391. TEST_GOVERNANCE_EMITTER,
  392. 1
  393. );
  394. // Create a custom event checker for ContractUpgraded event
  395. // Since we only care about the newImplementation parameter
  396. vm.recordLogs();
  397. // Execute the governance instruction
  398. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  399. // Get emitted logs and check the event parameters
  400. Vm.Log[] memory entries = vm.getRecordedLogs();
  401. bool foundUpgradeEvent = false;
  402. for (uint i = 0; i < entries.length; i++) {
  403. // The event signature for ContractUpgraded
  404. bytes32 eventSignature = keccak256(
  405. "ContractUpgraded(address,address)"
  406. );
  407. if (entries[i].topics[0] == eventSignature) {
  408. // This is a ContractUpgraded event
  409. // Get just the new implementation address using our helper
  410. address recordedNewImplementation = getSecondAddressFromEventData(
  411. entries[i].data
  412. );
  413. // Check newImplementation
  414. assertEq(recordedNewImplementation, address(newImplementation));
  415. foundUpgradeEvent = true;
  416. break;
  417. }
  418. }
  419. // Make sure we found the event
  420. assertTrue(foundUpgradeEvent, "ContractUpgraded event not found");
  421. // Verify the upgrade worked by checking the magic number
  422. assertEq(
  423. PythUpgradable(address(pyth)).pythUpgradableMagic(),
  424. 0x97a6f304
  425. );
  426. // Verify the implementation was upgraded to our new implementation
  427. // Access implementation using the ERC1967 storage slot
  428. address implAddr = address(
  429. uint160(
  430. uint256(
  431. vm.load(
  432. address(pyth),
  433. 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc // ERC1967 implementation slot
  434. )
  435. )
  436. )
  437. );
  438. assertEq(implAddr, address(newImplementation));
  439. }
  440. function testUpgradeContractToNonPythContractWontWork() public {
  441. // Deploy a mock upgradeable proxy that isn't a proper Pyth implementation
  442. MockUpgradeableProxy newImplementation = new MockUpgradeableProxy();
  443. // Create governance VAA to upgrade to an invalid implementation
  444. bytes memory data = abi.encodePacked(
  445. MAGIC,
  446. uint8(GovernanceModule.Target),
  447. uint8(GovernanceAction.UpgradeContract),
  448. TARGET_CHAIN_ID, // Valid target chain ID
  449. address(newImplementation) // Invalid implementation address
  450. );
  451. bytes memory vaa = encodeAndSignMessage(
  452. data,
  453. TEST_GOVERNANCE_CHAIN_ID,
  454. TEST_GOVERNANCE_EMITTER,
  455. 1
  456. );
  457. // Should revert with no specific error message because the mock implementation
  458. // doesn't have the pythUpgradableMagic method
  459. vm.expectRevert();
  460. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  461. }
  462. function testSetTransactionFee() public {
  463. // Set transaction fee to 1000 (1000 = 1 * 10^3)
  464. bytes memory setTransactionFeeMessage = abi.encodePacked(
  465. MAGIC,
  466. uint8(GovernanceModule.Target),
  467. uint8(GovernanceAction.SetTransactionFee),
  468. TARGET_CHAIN_ID,
  469. uint64(1), // value
  470. uint64(3) // exponent
  471. );
  472. bytes memory vaa = encodeAndSignMessage(
  473. setTransactionFeeMessage,
  474. TEST_GOVERNANCE_CHAIN_ID,
  475. TEST_GOVERNANCE_EMITTER,
  476. 1
  477. );
  478. uint oldFee = PythGetters(address(pyth)).transactionFeeInWei();
  479. vm.expectEmit(true, true, true, true);
  480. emit TransactionFeeSet(oldFee, 1000);
  481. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  482. assertEq(PythGetters(address(pyth)).transactionFeeInWei(), 1000);
  483. // Test that update fee includes transaction fee
  484. bytes[] memory updateData = new bytes[](0);
  485. assertEq(pyth.getUpdateFee(updateData), 1000);
  486. // Test that insufficient fee reverts
  487. vm.expectRevert(PythErrors.InsufficientFee.selector);
  488. pyth.updatePriceFeeds{value: 999}(updateData);
  489. // Test that sufficient fee works
  490. pyth.updatePriceFeeds{value: 1000}(updateData);
  491. }
  492. function testWithdrawFee() public {
  493. // First send some ETH to the contract
  494. bytes[] memory updateData = new bytes[](0);
  495. pyth.updatePriceFeeds{value: 1 ether}(updateData);
  496. assertEq(address(pyth).balance, 1 ether);
  497. address recipient = makeAddr("recipient");
  498. // Create governance VAA to withdraw fee
  499. bytes memory withdrawMessage = abi.encodePacked(
  500. MAGIC,
  501. uint8(GovernanceModule.Target),
  502. uint8(GovernanceAction.WithdrawFee),
  503. TARGET_CHAIN_ID,
  504. recipient,
  505. uint64(5), // value = 5
  506. uint64(17) // exponent = 17 (5 * 10^17 = 0.5 ether)
  507. );
  508. bytes memory vaa = encodeAndSignMessage(
  509. withdrawMessage,
  510. TEST_GOVERNANCE_CHAIN_ID,
  511. TEST_GOVERNANCE_EMITTER,
  512. 1
  513. );
  514. vm.expectEmit(true, true, true, true);
  515. emit FeeWithdrawn(recipient, 0.5 ether);
  516. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  517. assertEq(address(pyth).balance, 0.5 ether);
  518. assertEq(recipient.balance, 0.5 ether);
  519. }
  520. function testWithdrawFeeInsufficientBalance() public {
  521. // First send some ETH to the contract
  522. bytes[] memory updateData = new bytes[](0);
  523. pyth.updatePriceFeeds{value: 1 ether}(updateData);
  524. assertEq(address(pyth).balance, 1 ether);
  525. address recipient = makeAddr("recipient");
  526. // Create governance VAA to withdraw fee
  527. bytes memory withdrawMessage = abi.encodePacked(
  528. MAGIC,
  529. uint8(GovernanceModule.Target),
  530. uint8(GovernanceAction.WithdrawFee),
  531. TARGET_CHAIN_ID,
  532. recipient,
  533. uint64(2), // value = 2
  534. uint64(18) // exponent = 18 (2 * 10^18 = 2 ether, more than balance)
  535. );
  536. bytes memory vaa = encodeAndSignMessage(
  537. withdrawMessage,
  538. TEST_GOVERNANCE_CHAIN_ID,
  539. TEST_GOVERNANCE_EMITTER,
  540. 1
  541. );
  542. vm.expectRevert(PythErrors.InsufficientFee.selector);
  543. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  544. // Balances should remain unchanged
  545. assertEq(address(pyth).balance, 1 ether);
  546. assertEq(recipient.balance, 0);
  547. }
  548. function testWithdrawFeeInvalidGovernance() public {
  549. address recipient = makeAddr("recipient");
  550. // Create governance VAA with wrong emitter
  551. bytes memory withdrawMessage = abi.encodePacked(
  552. MAGIC,
  553. uint8(GovernanceModule.Target),
  554. uint8(GovernanceAction.WithdrawFee),
  555. TARGET_CHAIN_ID,
  556. recipient,
  557. uint64(5), // value = 5
  558. uint64(17) // exponent = 17 (5 * 10^17 = 0.5 ether)
  559. );
  560. bytes memory vaa = encodeAndSignMessage(
  561. withdrawMessage,
  562. TEST_GOVERNANCE_CHAIN_ID,
  563. bytes32(uint256(0x1111)), // Wrong emitter
  564. 1
  565. );
  566. vm.expectRevert(PythErrors.InvalidGovernanceDataSource.selector);
  567. PythGovernance(address(pyth)).executeGovernanceInstruction(vaa);
  568. }
  569. function encodeAndSignWormholeMessage(
  570. bytes memory data,
  571. uint16 emitterChainId,
  572. bytes32 emitterAddress,
  573. uint64 sequence,
  574. uint8 numGuardians
  575. ) internal view returns (bytes memory) {
  576. return
  577. generateVaa(
  578. uint32(block.timestamp),
  579. emitterChainId,
  580. emitterAddress,
  581. sequence,
  582. data,
  583. numGuardians,
  584. false
  585. );
  586. }
  587. function encodeAndSignMessage(
  588. bytes memory data,
  589. uint16 emitterChainId,
  590. bytes32 emitterAddress,
  591. uint64 sequence
  592. ) internal view returns (bytes memory) {
  593. return
  594. encodeAndSignWormholeMessage(
  595. data,
  596. emitterChainId,
  597. emitterAddress,
  598. sequence,
  599. 1 // Number of guardians
  600. );
  601. }
  602. // Events
  603. event ContractUpgraded(
  604. address oldImplementation,
  605. address newImplementation
  606. );
  607. event GovernanceDataSourceSet(
  608. PythInternalStructs.DataSource oldDataSource,
  609. PythInternalStructs.DataSource newDataSource,
  610. uint64 initialSequence
  611. );
  612. event DataSourcesSet(
  613. PythInternalStructs.DataSource[] oldDataSources,
  614. PythInternalStructs.DataSource[] newDataSources
  615. );
  616. event FeeSet(uint oldFee, uint newFee);
  617. event ValidPeriodSet(uint oldValidPeriod, uint newValidPeriod);
  618. event WormholeAddressSet(
  619. address oldWormholeAddress,
  620. address newWormholeAddress
  621. );
  622. event TransactionFeeSet(uint oldFee, uint newFee);
  623. event FeeWithdrawn(address recipient, uint256 fee);
  624. event VerifierAddressSet(
  625. address oldVerifierAddress,
  626. address newVerifierAddress
  627. );
  628. }