Skip to content

Commit 7f411d2

Browse files
guibescosguibescos
committed
Extra checks in load checked
1 parent 36bae1a commit 7f411d2

16 files changed

+169
-89
lines changed

‎program/rust/src/deserialize.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ use bytemuck::{
55
try_from_bytes_mut,
66
Pod,
77
};
8+
use solana_program::pubkey::Pubkey;
9+
use solana_program::rent::Rent;
810

911
use crate::c_oracle_header::{
1012
AccountHeader,
@@ -46,6 +48,7 @@ pub fn load_mut<T: Pod>(data: &mut [u8]) -> Result<&mut T, OracleError> {
4648
}
4749

4850
/// Get the data stored in `account` as a value of type `T`
51+
/// /// Warning: DON'T USE THIS UNCHECKED FUNCTION
4952
pub fn load_account_as<'a, T: Pod>(account: &'a AccountInfo) -> Result<Ref<'a, T>, ProgramError> {
5053
let data = account.try_borrow_data()?;
5154

@@ -69,7 +72,22 @@ pub fn load_account_as_mut<'a, T: Pod>(
6972
pub fn load_checked<'a, T: PythAccount>(
7073
account: &'a AccountInfo,
7174
version: u32,
75+
program_id: &Pubkey,
7276
) -> Result<RefMut<'a, T>, ProgramError> {
77+
// Check size and rent
78+
pyth_assert(
79+
account.data_len() >= T::minimum_size()
80+
&& Rent::default().is_exempt(account.lamports(), account.data_len()),
81+
OracleError::SmallSizeOrNotEnoughRent.into(),
82+
)?;
83+
84+
// Check owner
85+
pyth_assert(
86+
account.owner == program_id,
87+
OracleError::SmallSizeOrNotEnoughRent.into(),
88+
)?;
89+
90+
7391
{
7492
let account_header = load_account_as::<AccountHeader>(account)?;
7593
pyth_assert(

‎program/rust/src/error.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,10 @@ pub enum OracleError {
3232
InstructionDataTooShort = 610,
3333
#[error("InstructionDataSliceMisaligned")]
3434
InstructionDataSliceMisaligned = 611,
35+
#[error("SmallSizeOrNotEnoughRent")]
36+
SmallSizeOrNotEnoughRent = 612,
37+
#[error("AccountNotOwnerByOracle")]
38+
AccountNotOwnerByOracle = 613,
3539
}
3640

3741
impl From<OracleError> for ProgramError {

‎program/rust/src/rust_oracle.rs‎

Lines changed: 35 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -109,49 +109,49 @@ pub fn resize_price_account(
109109
accounts: &[AccountInfo],
110110
_instruction_data: &[u8],
111111
) -> ProgramResult {
112-
let [funding_account_info, price_account_info, system_program] = match accounts {
112+
let [funding_account_info, price_account, system_program] = match accounts {
113113
[x, y, z] => Ok([x, y, z]),
114114
_ => Err(ProgramError::InvalidArgument),
115115
}?;
116116

117117
check_valid_funding_account(funding_account_info)?;
118-
check_valid_signable_account(program_id, price_account_info, size_of::<PriceAccount>())?;
118+
check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
119119
pyth_assert(
120120
check_id(system_program.key),
121121
OracleError::InvalidSystemAccount.into(),
122122
)?;
123123
// Check that it is a valid initialized price account
124124
{
125-
load_checked::<PriceAccount>(price_account_info, PC_VERSION)?;
125+
load_checked::<PriceAccount>(price_account, PC_VERSION, program_id)?;
126126
}
127-
let account_len = price_account_info.try_data_len()?;
127+
let account_len = price_account.try_data_len()?;
128128
match account_len {
129129
PRICE_T_SIZE => {
130130
// Ensure account is still rent exempt after resizing
131131
let rent: Rent = Default::default();
132132
let lamports_needed: u64 = rent
133133
.minimum_balance(size_of::<PriceAccountWrapper>())
134-
.saturating_sub(price_account_info.lamports());
134+
.saturating_sub(price_account.lamports());
135135
if lamports_needed > 0 {
136136
send_lamports(
137137
funding_account_info,
138-
price_account_info,
138+
price_account,
139139
system_program,
140140
lamports_needed,
141141
)?;
142142
}
143143
// We do not need to zero allocate because we won't access the data in the same
144144
// instruction
145-
price_account_info.realloc(size_of::<PriceAccountWrapper>(), false)?;
145+
price_account.realloc(size_of::<PriceAccountWrapper>(), false)?;
146146

147147
// Check that everything is ok
148148
check_valid_signable_account(
149149
program_id,
150-
price_account_info,
150+
price_account,
151151
size_of::<PriceAccountWrapper>(),
152152
)?;
153153
let mut price_account =
154-
load_checked::<PriceAccountWrapper>(price_account_info, PC_VERSION)?;
154+
load_checked::<PriceAccountWrapper>(price_account, PC_VERSION, program_id)?;
155155
// Initialize Time Machine
156156
price_account.initialize_time_machine()?;
157157
Ok(())
@@ -206,7 +206,7 @@ pub fn add_mapping(
206206
check_valid_fresh_account(next_mapping)?;
207207

208208
let hdr = load::<CommandHeader>(instruction_data)?;
209-
let mut cur_mapping = load_checked::<MappingAccount>(cur_mapping, hdr.version)?;
209+
let mut cur_mapping = load_checked::<MappingAccount>(cur_mapping, hdr.version, program_id)?;
210210
pyth_assert(
211211
cur_mapping.num_ == PC_MAP_TABLE_SIZE && pubkey_is_zero(&cur_mapping.next_),
212212
ProgramError::InvalidArgument,
@@ -244,7 +244,8 @@ pub fn upd_price(
244244
let latest_aggregate_price: PriceInfo;
245245
{
246246
// Verify that symbol account is initialized
247-
let price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
247+
let price_data =
248+
load_checked::<PriceAccount>(price_account, cmd_args.header.version, program_id)?;
248249

249250
// Verify that publisher is authorized
250251
while publisher_index < price_data.num_ as usize {
@@ -306,7 +307,7 @@ pub fn upd_price(
306307

307308
{
308309
let mut price_data =
309-
load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
310+
load_checked::<PriceAccount>(price_account, cmd_args.header.version, program_id)?;
310311
let publisher_price = &mut price_data.comp_[publisher_index].latest_;
311312
publisher_price.price_ = cmd_args.price;
312313
publisher_price.conf_ = cmd_args.confidence;
@@ -359,7 +360,7 @@ pub fn add_price(
359360
check_valid_fresh_account(price_account)?;
360361

361362
let mut product_data =
362-
load_checked::<ProductAccount>(product_account, cmd_args.header.version)?;
363+
load_checked::<ProductAccount>(product_account, cmd_args.header.version, program_id)?;
363364

364365
let mut price_data =
365366
initialize_pyth_account_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
@@ -395,8 +396,9 @@ pub fn del_price(
395396

396397
{
397398
let cmd_args = load::<CommandHeader>(instruction_data)?;
398-
let mut product_data = load_checked::<ProductAccount>(product_account, cmd_args.version)?;
399-
let price_data = load_checked::<PriceAccount>(price_account, cmd_args.version)?;
399+
let mut product_data =
400+
load_checked::<ProductAccount>(product_account, cmd_args.version, program_id)?;
401+
let price_data = load_checked::<PriceAccount>(price_account, cmd_args.version, program_id)?;
400402
pyth_assert(
401403
pubkey_equal(&product_data.px_acc_, &price_account.key.to_bytes()),
402404
ProgramError::InvalidArgument,
@@ -437,7 +439,8 @@ pub fn init_price(
437439
check_valid_funding_account(funding_account)?;
438440
check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
439441

440-
let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
442+
let mut price_data =
443+
load_checked::<PriceAccount>(price_account, cmd_args.header.version, program_id)?;
441444
pyth_assert(
442445
price_data.ptype_ == cmd_args.price_type,
443446
ProgramError::InvalidArgument,
@@ -507,7 +510,8 @@ pub fn add_publisher(
507510
check_valid_funding_account(funding_account)?;
508511
check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
509512

510-
let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
513+
let mut price_data =
514+
load_checked::<PriceAccount>(price_account, cmd_args.header.version, program_id)?;
511515

512516
if price_data.num_ >= PC_COMP_SIZE {
513517
return Err(ProgramError::InvalidArgument);
@@ -560,7 +564,8 @@ pub fn del_publisher(
560564
check_valid_funding_account(funding_account)?;
561565
check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
562566

563-
let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
567+
let mut price_data =
568+
load_checked::<PriceAccount>(price_account, cmd_args.header.version, program_id)?;
564569

565570
for i in 0..(price_data.num_ as usize) {
566571
if pubkey_equal(&cmd_args.publisher, bytes_of(&price_data.comp_[i].pub_)) {
@@ -603,7 +608,8 @@ pub fn add_product(
603608
check_valid_fresh_account(new_product_account)?;
604609

605610
let hdr = load::<CommandHeader>(instruction_data)?;
606-
let mut mapping_data = load_checked::<MappingAccount>(tail_mapping_account, hdr.version)?;
611+
let mut mapping_data =
612+
load_checked::<MappingAccount>(tail_mapping_account, hdr.version, program_id)?;
607613
// The mapping account must have free space to add the product account
608614
pyth_assert(
609615
mapping_data.num_ < PC_MAP_TABLE_SIZE,
@@ -643,7 +649,8 @@ pub fn upd_product(
643649
let hdr = load::<CommandHeader>(instruction_data)?;
644650
{
645651
// Validate that product_account contains the appropriate account header
646-
let mut _product_data = load_checked::<ProductAccount>(product_account, hdr.version)?;
652+
let mut _product_data =
653+
load_checked::<ProductAccount>(product_account, hdr.version, program_id)?;
647654
}
648655

649656
pyth_assert(
@@ -679,7 +686,8 @@ pub fn upd_product(
679686
);
680687
}
681688

682-
let mut product_data = load_checked::<ProductAccount>(product_account, hdr.version)?;
689+
let mut product_data =
690+
load_checked::<ProductAccount>(product_account, hdr.version, program_id)?;
683691
product_data.size_ = try_convert(size_of::<ProductAccount>() + new_data.len())?;
684692

685693
Ok(())
@@ -705,7 +713,8 @@ pub fn set_min_pub(
705713
check_valid_funding_account(funding_account)?;
706714
check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
707715

708-
let mut price_account_data = load_checked::<PriceAccount>(price_account, cmd.header.version)?;
716+
let mut price_account_data =
717+
load_checked::<PriceAccount>(price_account, cmd.header.version, program_id)?;
709718
price_account_data.min_pub_ = cmd.minimum_publishers;
710719

711720
Ok(())
@@ -734,8 +743,10 @@ pub fn del_product(
734743

735744
{
736745
let cmd_args = load::<CommandHeader>(instruction_data)?;
737-
let mut mapping_data = load_checked::<MappingAccount>(mapping_account, cmd_args.version)?;
738-
let product_data = load_checked::<ProductAccount>(product_account, cmd_args.version)?;
746+
let mut mapping_data =
747+
load_checked::<MappingAccount>(mapping_account, cmd_args.version, program_id)?;
748+
let product_data =
749+
load_checked::<ProductAccount>(product_account, cmd_args.version, program_id)?;
739750

740751
// This assertion is just to make the subtractions below simpler
741752
pyth_assert(mapping_data.num_ >= 1, ProgramError::InvalidArgument)?;

‎program/rust/src/tests/test_add_mapping.rs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ use crate::c_oracle_header::{
66
};
77
use crate::deserialize::{
88
initialize_pyth_account_checked,
9-
load_account_as_mut,
109
load_checked,
1110
};
1211
use crate::instruction::{
@@ -44,7 +43,7 @@ fn test_add_mapping() {
4443

4544
{
4645
let mut cur_mapping_data =
47-
load_checked::<MappingAccount>(&cur_mapping, PC_VERSION).unwrap();
46+
load_checked::<MappingAccount>(&cur_mapping, PC_VERSION, &program_id).unwrap();
4847
cur_mapping_data.num_ = PC_MAP_TABLE_SIZE;
4948
}
5049

@@ -60,9 +59,10 @@ fn test_add_mapping() {
6059
.is_ok());
6160

6261
{
63-
let next_mapping_data = load_checked::<MappingAccount>(&next_mapping, PC_VERSION).unwrap();
62+
let next_mapping_data =
63+
load_checked::<MappingAccount>(&next_mapping, PC_VERSION, &program_id).unwrap();
6464
let mut cur_mapping_data =
65-
load_checked::<MappingAccount>(&cur_mapping, PC_VERSION).unwrap();
65+
load_checked::<MappingAccount>(&cur_mapping, PC_VERSION, &program_id).unwrap();
6666

6767
assert!(pubkey_equal(
6868
&cur_mapping_data.next_,
@@ -90,7 +90,7 @@ fn test_add_mapping() {
9090

9191
{
9292
let mut cur_mapping_data =
93-
load_checked::<MappingAccount>(&cur_mapping, PC_VERSION).unwrap();
93+
load_checked::<MappingAccount>(&cur_mapping, PC_VERSION, &program_id).unwrap();
9494
assert!(pubkey_is_zero(&cur_mapping_data.next_));
9595
cur_mapping_data.num_ = PC_MAP_TABLE_SIZE;
9696
cur_mapping_data.magic_ = 0;
@@ -110,7 +110,8 @@ fn test_add_mapping() {
110110
);
111111

112112
{
113-
let mut cur_mapping_data = load_account_as_mut::<MappingAccount>(&cur_mapping).unwrap();
113+
let mut cur_mapping_data =
114+
load_checked::<MappingAccount>(&cur_mapping, PC_VERSION, &program_id).unwrap();
114115
cur_mapping_data.magic_ = PC_MAGIC;
115116
}
116117

‎program/rust/src/tests/test_add_price.rs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,10 @@ fn test_add_price() {
7878
.is_ok());
7979

8080
{
81-
let price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
82-
let product_data = load_checked::<ProductAccount>(&product_account, PC_VERSION).unwrap();
81+
let price_data =
82+
load_checked::<PriceAccount>(&price_account, PC_VERSION, &program_id).unwrap();
83+
let product_data =
84+
load_checked::<ProductAccount>(&product_account, PC_VERSION, &program_id).unwrap();
8385
assert_eq!(price_data.expo_, 1);
8486
assert_eq!(price_data.ptype_, 1);
8587
assert!(pubkey_equal(
@@ -105,8 +107,10 @@ fn test_add_price() {
105107
.is_ok());
106108

107109
{
108-
let price_data_2 = load_checked::<PriceAccount>(&price_account_2, PC_VERSION).unwrap();
109-
let product_data = load_checked::<ProductAccount>(&product_account, PC_VERSION).unwrap();
110+
let price_data_2 =
111+
load_checked::<PriceAccount>(&price_account_2, PC_VERSION, &program_id).unwrap();
112+
let product_data =
113+
load_checked::<ProductAccount>(&product_account, PC_VERSION, &program_id).unwrap();
110114
assert_eq!(price_data_2.expo_, 1);
111115
assert_eq!(price_data_2.ptype_, 1);
112116
assert!(pubkey_equal(

‎program/rust/src/tests/test_add_product.rs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ use crate::c_oracle_header::{
1212
};
1313
use crate::deserialize::{
1414
initialize_pyth_account_checked,
15-
load_account_as,
1615
load_checked,
1716
};
1817
use crate::error::OracleError;
@@ -66,8 +65,10 @@ fn test_add_product() {
6665
.is_ok());
6766

6867
{
69-
let product_data = load_account_as::<ProductAccount>(&product_account).unwrap();
70-
let mapping_data = load_checked::<MappingAccount>(&mapping_account, PC_VERSION).unwrap();
68+
let product_data =
69+
load_checked::<ProductAccount>(&product_account, PC_VERSION, &program_id).unwrap();
70+
let mapping_data =
71+
load_checked::<MappingAccount>(&mapping_account, PC_VERSION, &program_id).unwrap();
7172

7273
assert_eq!(product_data.magic_, PC_MAGIC);
7374
assert_eq!(product_data.ver_, PC_VERSION);
@@ -92,7 +93,8 @@ fn test_add_product() {
9293
)
9394
.is_ok());
9495
{
95-
let mapping_data = load_checked::<MappingAccount>(&mapping_account, PC_VERSION).unwrap();
96+
let mapping_data =
97+
load_checked::<MappingAccount>(&mapping_account, PC_VERSION, &program_id).unwrap();
9698
assert_eq!(mapping_data.num_, 2);
9799
assert_eq!(mapping_data.size_, (MappingAccount::INITIAL_SIZE + 2 * 32));
98100
assert!(pubkey_equal(
@@ -145,7 +147,8 @@ fn test_add_product() {
145147
instruction_data
146148
)
147149
.is_ok());
148-
let mapping_data = load_checked::<MappingAccount>(&mapping_account, PC_VERSION).unwrap();
150+
let mapping_data =
151+
load_checked::<MappingAccount>(&mapping_account, PC_VERSION, &program_id).unwrap();
149152
assert_eq!(
150153
mapping_data.size_,
151154
MappingAccount::INITIAL_SIZE + (i + 1) * 32
@@ -168,6 +171,7 @@ fn test_add_product() {
168171
Err(ProgramError::InvalidArgument)
169172
);
170173

171-
let mapping_data = load_checked::<MappingAccount>(&mapping_account, PC_VERSION).unwrap();
174+
let mapping_data =
175+
load_checked::<MappingAccount>(&mapping_account, PC_VERSION, &program_id).unwrap();
172176
assert_eq!(mapping_data.num_, PC_MAP_TABLE_SIZE);
173177
}

‎program/rust/src/tests/test_add_publisher.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,8 @@ fn test_add_publisher() {
7171
.is_ok());
7272

7373
{
74-
let price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
74+
let price_data =
75+
load_checked::<PriceAccount>(&price_account, PC_VERSION, &program_id).unwrap();
7576
assert_eq!(price_data.num_, 1);
7677
assert_eq!(
7778
price_data.size_,
@@ -120,7 +121,8 @@ fn test_add_publisher() {
120121

121122

122123
{
123-
let price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
124+
let price_data =
125+
load_checked::<PriceAccount>(&price_account, PC_VERSION, &program_id).unwrap();
124126
assert_eq!(price_data.num_, i + 1);
125127
assert!(pubkey_equal(
126128
&price_data.comp_[i as usize].pub_,

0 commit comments

Comments
 (0)