Skip to content

Commit ea00af3

Browse files
ayushbossayushboss
authored
Merge pull request #2928 from pyth-network/stylus-init-check
fix(stylus) - added initialization check
2 parents 562d0a8 + 9e194a1 commit ea00af3

File tree

4 files changed

+71
-48
lines changed

4 files changed

+71
-48
lines changed

‎target_chains/stylus/contracts/pyth-receiver/src/error.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ pub enum PythReceiverError {
2828
GovernanceMessageAlreadyExecuted,
2929
InvalidWormholeAddressToSet,
3030
WormholeUninitialized,
31+
AlreadyInitialized,
3132
}
3233

3334
impl core::fmt::Debug for PythReceiverError {
@@ -71,6 +72,7 @@ impl core::fmt::Debug for PythReceiverError {
7172
PythReceiverError::WormholeUninitialized => {
7273
write!(f, "Wormhole is uninitialized, please set the Wormhole address and initialize the contract first")
7374
}
75+
PythReceiverError::AlreadyInitialized => write!(f, "AlreadyInitialized"),
7476
}
7577
}
7678
}
@@ -118,6 +120,7 @@ impl core::fmt::Display for PythReceiverError {
118120
PythReceiverError::WormholeUninitialized => {
119121
write!(f, "Wormhole is uninitialized, please set the Wormhole address and initialize the contract first")
120122
}
123+
PythReceiverError::AlreadyInitialized => write!(f, "Contract is already initialized"),
121124
}
122125
}
123126
}
@@ -151,6 +154,7 @@ impl From<PythReceiverError> for Vec<u8> {
151154
PythReceiverError::GovernanceMessageAlreadyExecuted => 24,
152155
PythReceiverError::InvalidWormholeAddressToSet => 25,
153156
PythReceiverError::WormholeUninitialized => 26,
157+
PythReceiverError::AlreadyInitialized => 27,
154158
}]
155159
}
156160
}

‎target_chains/stylus/contracts/pyth-receiver/src/integration_tests.rs‎

Lines changed: 58 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -49,25 +49,11 @@ mod test {
4949
}
5050

5151
#[cfg(test)]
52-
fn pyth_wormhole_init(
52+
fn pyth_init(
5353
pyth_contract: &Contract<PythReceiver>,
5454
wormhole_contract: &Contract<WormholeContract>,
5555
alice: &Address,
56-
) {
57-
let guardians = current_guardians();
58-
let governance_contract =
59-
Address::from_slice(&GOVERNANCE_CONTRACT.to_be_bytes::<32>()[12..32]);
60-
wormhole_contract
61-
.sender(*alice)
62-
.initialize(
63-
guardians,
64-
4,
65-
CHAIN_ID,
66-
GOVERNANCE_CHAIN_ID,
67-
governance_contract,
68-
)
69-
.unwrap();
70-
56+
) -> Result<(), PythReceiverError> {
7157
let single_update_fee = SINGLE_UPDATE_FEE_IN_WEI;
7258
let valid_time_period = U256::from(3600u64);
7359

@@ -87,7 +73,33 @@ mod test {
8773
governance_chain_id,
8874
governance_emitter_address,
8975
governance_initial_sequence,
90-
);
76+
)?;
77+
78+
Ok(())
79+
}
80+
81+
#[cfg(test)]
82+
fn pyth_wormhole_init(
83+
pyth_contract: &Contract<PythReceiver>,
84+
wormhole_contract: &Contract<WormholeContract>,
85+
alice: &Address,
86+
) -> Result<(), PythReceiverError> {
87+
let guardians = current_guardians();
88+
let governance_contract =
89+
Address::from_slice(&GOVERNANCE_CONTRACT.to_be_bytes::<32>()[12..32]);
90+
wormhole_contract
91+
.sender(*alice)
92+
.initialize(
93+
guardians,
94+
4,
95+
CHAIN_ID,
96+
GOVERNANCE_CHAIN_ID,
97+
governance_contract,
98+
)
99+
.unwrap();
100+
101+
pyth_init(pyth_contract, wormhole_contract, alice)?;
102+
Ok(())
91103
}
92104

93105
#[motsu::test]
@@ -96,7 +108,7 @@ mod test {
96108
wormhole_contract: Contract<WormholeContract>,
97109
alice: Address,
98110
) {
99-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
111+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
100112

101113
let update_data = ban_usd_update();
102114
let update_fee = mock_get_update_fee(update_data.clone()).unwrap();
@@ -121,7 +133,7 @@ mod test {
121133
wormhole_contract: Contract<WormholeContract>,
122134
alice: Address,
123135
) {
124-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
136+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
125137

126138
alice.fund(U256::from(200));
127139

@@ -142,7 +154,7 @@ mod test {
142154
wormhole_contract: Contract<WormholeContract>,
143155
alice: Address,
144156
) {
145-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
157+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
146158

147159
let update_data1 = ban_usd_update();
148160
let update_fee1 = mock_get_update_fee(update_data1.clone()).unwrap();
@@ -175,7 +187,7 @@ mod test {
175187
wormhole_contract: Contract<WormholeContract>,
176188
alice: Address,
177189
) {
178-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
190+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
179191

180192
let price_result = pyth_contract
181193
.sender(alice)
@@ -194,7 +206,7 @@ mod test {
194206
alice: Address,
195207
) {
196208
MockClock::set_time(Duration::from_secs(1761573860)); // less than good_update2().timestamp + 1s
197-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
209+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
198210

199211
let random_id: [u8; 32] = [
200212
0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc,
@@ -219,7 +231,7 @@ mod test {
219231
alice: Address,
220232
) {
221233
MockClock::set_time(Duration::from_secs(1761573860));
222-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
234+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
223235

224236
let update_data = btc_usd_update();
225237
let update_fee = mock_get_update_fee(update_data.clone()).unwrap();
@@ -245,7 +257,7 @@ mod test {
245257
alice: Address,
246258
) {
247259
MockClock::set_time(Duration::from_secs(1761573860));
248-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
260+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
249261

250262
let update_data = btc_usd_update();
251263
let update_fee = mock_get_update_fee(update_data.clone()).unwrap();
@@ -274,7 +286,7 @@ mod test {
274286
wormhole_contract: Contract<WormholeContract>,
275287
alice: Address,
276288
) {
277-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
289+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
278290

279291
let update_data = multiple_updates_diff_vaa();
280292
let update_fee = mock_get_update_fee(update_data.clone()).unwrap();
@@ -311,7 +323,7 @@ mod test {
311323
wormhole_contract: Contract<WormholeContract>,
312324
alice: Address,
313325
) {
314-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
326+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
315327

316328
assert!(!pyth_contract
317329
.sender(alice)
@@ -338,7 +350,7 @@ mod test {
338350
wormhole_contract: Contract<WormholeContract>,
339351
alice: Address,
340352
) {
341-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
353+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
342354

343355
let price_result = pyth_contract
344356
.sender(alice)
@@ -357,7 +369,7 @@ mod test {
357369
wormhole_contract: Contract<WormholeContract>,
358370
alice: Address,
359371
) {
360-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
372+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
361373

362374
let update_data = ban_usd_update();
363375
let update_fee = mock_get_update_fee(update_data.clone()).unwrap();
@@ -384,7 +396,7 @@ mod test {
384396
wormhole_contract: Contract<WormholeContract>,
385397
alice: Address,
386398
) {
387-
pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
399+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
388400

389401
let update_data = multiple_updates_diff_vaa();
390402
let update_fee = mock_get_update_fee(update_data.clone()).unwrap();
@@ -417,4 +429,21 @@ mod test {
417429
multiple_updates_diff_vaa_results_full()[1]
418430
);
419431
}
432+
433+
#[motsu::test]
434+
fn test_double_initialization_reverts(
435+
pyth_contract: Contract<PythReceiver>,
436+
wormhole_contract: Contract<WormholeContract>,
437+
alice: Address,
438+
) {
439+
let _ = pyth_wormhole_init(&pyth_contract, &wormhole_contract, &alice);
440+
441+
let double_init = pyth_init(&pyth_contract, &wormhole_contract, &alice);
442+
443+
assert!(double_init.is_err());
444+
assert_eq!(
445+
double_init.unwrap_err(),
446+
PythReceiverError::AlreadyInitialized
447+
);
448+
}
420449
}

‎target_chains/stylus/contracts/pyth-receiver/src/lib.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ pub struct PythReceiver {
8080
pub governance_data_source_index: StorageUint<32, 1>,
8181
pub latest_price_info: StorageMap<FixedBytes<32>, PriceFeedStorage>,
8282
pub transaction_fee_in_wei: StorageU256,
83+
pub initialized: StorageBool,
8384
}
8485

8586
#[public]
@@ -94,7 +95,10 @@ impl PythReceiver {
9495
governance_emitter_chain_id: u16,
9596
governance_emitter_address: [u8; 32],
9697
governance_initial_sequence: u64,
97-
) {
98+
) -> Result<(), PythReceiverError> {
99+
if self.initialized.get() {
100+
return Err(PythReceiverError::AlreadyInitialized.into());
101+
}
98102
self.wormhole.set(wormhole);
99103
self.single_update_fee_in_wei.set(single_update_fee_in_wei);
100104
self.valid_time_period_seconds
@@ -123,6 +127,8 @@ impl PythReceiver {
123127

124128
self.is_valid_data_source.setter(data_source_key).set(true);
125129
}
130+
self.initialized.set(true);
131+
Ok(())
126132
}
127133

128134
pub fn price_feed_exists(&self, id: [u8; 32]) -> bool {

‎target_chains/stylus/contracts/pyth-receiver/src/pyth_governance_test.rs‎

Lines changed: 2 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -20,29 +20,13 @@ mod test {
2020
const GOVERNANCE_CONTRACT: U256 = U256::from_limbs([4, 0, 0, 0]);
2121

2222
const SINGLE_UPDATE_FEE_IN_WEI: U256 = U256::from_limbs([100, 0, 0, 0]);
23-
const TRANSACTION_FEE_IN_WEI: U256 = U256::from_limbs([32, 0, 0, 0]);
24-
25-
const TEST_SIGNER1: Address = Address::new([
26-
0xbe, 0xFA, 0x42, 0x9d, 0x57, 0xcD, 0x18, 0xb7, 0xF8, 0xA4, 0xd9, 0x1A, 0x2d, 0xa9, 0xAB,
27-
0x4A, 0xF0, 0x5d, 0x0F, 0xBe,
28-
]);
29-
const TEST_SIGNER2: Address = Address::new([
30-
0x4b, 0xa0, 0xC2, 0xdb, 0x9A, 0x26, 0x20, 0x8b, 0x3b, 0xB1, 0xa5, 0x0B, 0x01, 0xb1, 0x69,
31-
0x41, 0xc1, 0x0D, 0x76, 0xdb,
32-
]);
23+
3324
const GOVERNANCE_CHAIN_ID: u16 = 1;
3425
const GOVERNANCE_EMITTER: [u8; 32] = [
3526
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
3627
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
3728
0x00, 0x11,
3829
];
39-
const TEST_PYTH2_WORMHOLE_CHAIN_ID: u16 = 1;
40-
const TEST_PYTH2_WORMHOLE_EMITTER: [u8; 32] = [
41-
0x71, 0xf8, 0xdc, 0xb8, 0x63, 0xd1, 0x76, 0xe2, 0xc4, 0x20, 0xad, 0x66, 0x10, 0xcf, 0x68,
42-
0x73, 0x59, 0x61, 0x2b, 0x6f, 0xb3, 0x92, 0xe0, 0x64, 0x2b, 0x0c, 0xa6, 0xb1, 0xf1, 0x86,
43-
0xaa, 0x3b,
44-
];
45-
const TARGET_CHAIN_ID: u16 = 2;
4630

4731
#[cfg(test)]
4832
fn pyth_wormhole_init(
@@ -75,7 +59,7 @@ mod test {
7559
let governance_chain_id = 1u16;
7660
let governance_initial_sequence = 0u64;
7761

78-
pyth_contract.sender(*alice).initialize(
62+
let _ = pyth_contract.sender(*alice).initialize(
7963
wormhole_contract.address(),
8064
single_update_fee,
8165
valid_time_period,

0 commit comments

Comments
 (0)