Move minimum size check to load_checked and initialize_checked (#271)

* Size moved to deserialization functions

* Stop using load_account_as in rust

* Comments

Author: guibescos

program/rust/src/deserialize.rs

@@ -13,6 +13,7 @@ use crate::c_oracle_header::{
 };
 use crate::error::OracleError;
 use crate::utils::{
+    check_valid_fresh_account,
     clear_account,
     pyth_assert,
 };
@@ -45,7 +46,8 @@ pub fn load_mut<T: Pod>(data: &mut [u8]) -> Result<&mut T, OracleError> {
     .map_err(|_| OracleError::InstructionDataSliceMisaligned)
 }
 
-/// Get the data stored in `account` as a value of type `T`
+/// Get the data stored in `account` as a value of type `T`.
+/// WARNING : Use `load_checked` to load initialized Pyth accounts
 pub fn load_account_as<'a, T: Pod>(account: &'a AccountInfo) -> Result<Ref<'a, T>, ProgramError> {
     let data = account.try_borrow_data()?;
 
@@ -56,6 +58,7 @@ pub fn load_account_as<'a, T: Pod>(account: &'a AccountInfo) -> Result<Ref<'a, T
 
 /// Mutably borrow the data in `account` as a value of type `T`.
 /// Any mutations to the returned value will be reflected in the account data.
+/// WARNING : Use `load_checked` to load initialized Pyth accounts
 pub fn load_account_as_mut<'a, T: Pod>(
     account: &'a AccountInfo,
 ) -> Result<RefMut<'a, T>, ProgramError> {
@@ -70,6 +73,11 @@ pub fn load_checked<'a, T: PythAccount>(
     account: &'a AccountInfo,
     version: u32,
 ) -> Result<RefMut<'a, T>, ProgramError> {
+    pyth_assert(
+        account.data_len() >= T::minimum_size(),
+        OracleError::AccountTooSmall.into(),
+    )?;
+
     {
         let account_header = load_account_as::<AccountHeader>(account)?;
         pyth_assert(
@@ -87,6 +95,11 @@ pub fn initialize_pyth_account_checked<'a, T: PythAccount>(
     account: &'a AccountInfo,
     version: u32,
 ) -> Result<RefMut<'a, T>, ProgramError> {
+    pyth_assert(
+        account.data_len() >= T::minimum_size(),
+        OracleError::AccountTooSmall.into(),
+    )?;
+    check_valid_fresh_account(account)?;
     clear_account(account)?;
 
     {

program/rust/src/error.rs

@@ -32,6 +32,8 @@ pub enum OracleError {
     InstructionDataTooShort        = 610,
     #[error("InstructionDataSliceMisaligned")]
     InstructionDataSliceMisaligned = 611,
+    #[error("AccountTooSmall")]
+    AccountTooSmall                = 612,
 }
 
 impl From<OracleError> for ProgramError {

program/rust/src/rust_oracle.rs

@@ -22,7 +22,6 @@ use crate::c_oracle_header::{
 use crate::deserialize::{
     initialize_pyth_account_checked,
     load,
-    load_account_as_mut,
     load_checked,
 };
 use crate::instruction::{
@@ -37,7 +36,6 @@ use crate::instruction::{
 use crate::time_machine_types::PriceAccountWrapper;
 use crate::utils::{
     check_exponent_range,
-    check_valid_fresh_account,
     check_valid_funding_account,
     check_valid_signable_account,
     check_valid_writable_account,
@@ -115,7 +113,7 @@ pub fn resize_price_account(
     }?;
 
     check_valid_funding_account(funding_account_info)?;
-    check_valid_signable_account(program_id, price_account_info, size_of::<PriceAccount>())?;
+    check_valid_signable_account(program_id, price_account_info)?;
     pyth_assert(
         check_id(system_program.key),
         OracleError::InvalidSystemAccount.into(),
@@ -145,11 +143,7 @@ pub fn resize_price_account(
             price_account_info.realloc(size_of::<PriceAccountWrapper>(), false)?;
 
             // Check that everything is ok
-            check_valid_signable_account(
-                program_id,
-                price_account_info,
-                size_of::<PriceAccountWrapper>(),
-            )?;
+            check_valid_signable_account(program_id, price_account_info)?;
             let mut price_account =
                 load_checked::<PriceAccountWrapper>(price_account_info, PC_VERSION)?;
             // Initialize Time Machine
@@ -176,12 +170,7 @@ pub fn init_mapping(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(
-        program_id,
-        fresh_mapping_account,
-        size_of::<MappingAccount>(),
-    )?;
-    check_valid_fresh_account(fresh_mapping_account)?;
+    check_valid_signable_account(program_id, fresh_mapping_account)?;
 
     // Initialize by setting to zero again (just in case) and populating the account header
     let hdr = load::<CommandHeader>(instruction_data)?;
@@ -201,9 +190,8 @@ pub fn add_mapping(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, cur_mapping, size_of::<MappingAccount>())?;
-    check_valid_signable_account(program_id, next_mapping, size_of::<MappingAccount>())?;
-    check_valid_fresh_account(next_mapping)?;
+    check_valid_signable_account(program_id, cur_mapping)?;
+    check_valid_signable_account(program_id, next_mapping)?;
 
     let hdr = load::<CommandHeader>(instruction_data)?;
     let mut cur_mapping = load_checked::<MappingAccount>(cur_mapping, hdr.version)?;
@@ -240,7 +228,7 @@ pub fn upd_price(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_writable_account(program_id, price_account, size_of::<PriceAccount>())?;
+    check_valid_writable_account(program_id, price_account)?;
     // Check clock
     let clock = Clock::from_account_info(clock_account)?;
 
@@ -291,7 +279,8 @@ pub fn upd_price(
 
     let account_len = price_account.try_data_len()?;
     if aggregate_updated && account_len == PRICE_ACCOUNT_SIZE {
-        let mut price_account = load_account_as_mut::<PriceAccountWrapper>(price_account)?;
+        let mut price_account =
+            load_checked::<PriceAccountWrapper>(price_account, cmd_args.header.version)?;
         price_account.add_price_to_time_machine()?;
     }
 
@@ -358,9 +347,8 @@ pub fn add_price(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, product_account, PC_PROD_ACC_SIZE as usize)?;
-    check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
-    check_valid_fresh_account(price_account)?;
+    check_valid_signable_account(program_id, product_account)?;
+    check_valid_signable_account(program_id, price_account)?;
 
     let mut product_data =
         load_checked::<ProductAccount>(product_account, cmd_args.header.version)?;
@@ -403,8 +391,8 @@ pub fn del_price(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, product_account, PC_PROD_ACC_SIZE as usize)?;
-    check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
+    check_valid_signable_account(program_id, product_account)?;
+    check_valid_signable_account(program_id, price_account)?;
 
     {
         let cmd_args = load::<CommandHeader>(instruction_data)?;
@@ -454,7 +442,7 @@ pub fn init_price(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
+    check_valid_signable_account(program_id, price_account)?;
 
     let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
     pyth_assert(
@@ -524,7 +512,7 @@ pub fn add_publisher(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
+    check_valid_signable_account(program_id, price_account)?;
 
     let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
 
@@ -577,7 +565,7 @@ pub fn del_publisher(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
+    check_valid_signable_account(program_id, price_account)?;
 
     let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
 
@@ -613,13 +601,8 @@ pub fn add_product(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(
-        program_id,
-        tail_mapping_account,
-        size_of::<MappingAccount>(),
-    )?;
-    check_valid_signable_account(program_id, new_product_account, PC_PROD_ACC_SIZE as usize)?;
-    check_valid_fresh_account(new_product_account)?;
+    check_valid_signable_account(program_id, tail_mapping_account)?;
+    check_valid_signable_account(program_id, new_product_account)?;
 
     let hdr = load::<CommandHeader>(instruction_data)?;
     let mut mapping_data = load_checked::<MappingAccount>(tail_mapping_account, hdr.version)?;
@@ -658,7 +641,7 @@ pub fn upd_product(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, product_account, try_convert(PC_PROD_ACC_SIZE)?)?;
+    check_valid_signable_account(program_id, product_account)?;
 
     let hdr = load::<CommandHeader>(instruction_data)?;
     {
@@ -723,7 +706,7 @@ pub fn set_min_pub(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, price_account, size_of::<PriceAccount>())?;
+    check_valid_signable_account(program_id, price_account)?;
 
     let mut price_account_data = load_checked::<PriceAccount>(price_account, cmd.header.version)?;
     price_account_data.min_pub_ = cmd.minimum_publishers;
@@ -749,8 +732,8 @@ pub fn del_product(
     }?;
 
     check_valid_funding_account(funding_account)?;
-    check_valid_signable_account(program_id, mapping_account, size_of::<MappingAccount>())?;
-    check_valid_signable_account(program_id, product_account, PC_PROD_ACC_SIZE as usize)?;
+    check_valid_signable_account(program_id, mapping_account)?;
+    check_valid_signable_account(program_id, product_account)?;
 
     {
         let cmd_args = load::<CommandHeader>(instruction_data)?;

program/rust/src/tests/test_add_product.rs

@@ -131,7 +131,7 @@ fn test_add_product() {
             ],
             instruction_data
         ),
-        Err(OracleError::InvalidSignableAccount.into())
+        Err(OracleError::AccountTooSmall.into())
     );
 
     // test fill up of mapping table

program/rust/src/tests/test_init_mapping.rs

@@ -134,7 +134,7 @@ fn test_init_mapping() {
             &[funding_account.clone(), mapping_account.clone()],
             instruction_data
         ),
-        Err(OracleError::InvalidSignableAccount.into())
+        Err(OracleError::AccountTooSmall.into())
     );
 
     mapping_account.data = prev_data;

program/rust/src/utils.rs

@@ -47,25 +47,19 @@ pub fn check_valid_funding_account(account: &AccountInfo) -> Result<(), ProgramE
     )
 }
 
-pub fn valid_signable_account(
-    program_id: &Pubkey,
-    account: &AccountInfo,
-    minimum_size: usize,
-) -> bool {
+pub fn valid_signable_account(program_id: &Pubkey, account: &AccountInfo) -> bool {
     account.is_signer
         && account.is_writable
         && account.owner == program_id
-        && account.data_len() >= minimum_size
         && Rent::default().is_exempt(account.lamports(), account.data_len())
 }
 
 pub fn check_valid_signable_account(
     program_id: &Pubkey,
     account: &AccountInfo,
-    minimum_size: usize,
 ) -> Result<(), ProgramError> {
     pyth_assert(
-        valid_signable_account(program_id, account, minimum_size),
+        valid_signable_account(program_id, account),
         OracleError::InvalidSignableAccount.into(),
     )
 }
@@ -138,20 +132,18 @@ pub fn read_pc_str_t(source: &[u8]) -> Result<&[u8], ProgramError> {
     }
 }
 
-fn valid_writable_account(program_id: &Pubkey, account: &AccountInfo, minimum_size: usize) -> bool {
+fn valid_writable_account(program_id: &Pubkey, account: &AccountInfo) -> bool {
     account.is_writable
         && account.owner == program_id
-        && account.data_len() >= minimum_size
         && Rent::default().is_exempt(account.lamports(), account.data_len())
 }
 
 pub fn check_valid_writable_account(
     program_id: &Pubkey,
     account: &AccountInfo,
-    minimum_size: usize,
 ) -> Result<(), ProgramError> {
     pyth_assert(
-        valid_writable_account(program_id, account, minimum_size),
+        valid_writable_account(program_id, account),
         OracleError::InvalidWritableAccount.into(),
     )
 }