feat: init feed index when creating prices and init permission account

Riateche · Riateche · commit 38ada8a5ff10 · 2024-08-05T12:27:32.000+01:00
diff --git a/program/rust/src/accounts.rs b/program/rust/src/accounts.rs
@@ -105,6 +105,11 @@ pub trait PythAccount: Pod {
     /// have. `INITIAL_SIZE` <= `minimum_size()`
     const MINIMUM_SIZE: usize = size_of::<Self>();
 
+    /// Size of the account data on creation. Usually this is the same as `MINIMUM_SIZE` but it's
+    /// different for `PermissionAccount` because we've added new fields to it. In this case
+    /// we cannot increase `MINIMUM_SIZE` because that would break reading the account.
+    const NEW_ACCOUNT_SPACE: usize = Self::MINIMUM_SIZE;
+
     /// Given an `AccountInfo`, verify it is sufficiently large and has the correct discriminator.
     fn initialize<'a>(
         account: &'a AccountInfo,
@@ -139,15 +144,15 @@ pub trait PythAccount: Pod {
         seeds: &[&[u8]],
         version: u32,
     ) -> Result<(), ProgramError> {
-        let target_rent = get_rent()?.minimum_balance(Self::MINIMUM_SIZE);
+        let target_rent = get_rent()?.minimum_balance(Self::NEW_ACCOUNT_SPACE);
 
         if account.data_len() == 0 {
             create(
                 funding_account,
                 account,
                 system_program,
                 program_id,
-                Self::MINIMUM_SIZE,
+                Self::NEW_ACCOUNT_SPACE,
                 target_rent,
                 seeds,
             )?;
diff --git a/program/rust/src/accounts/permission.rs b/program/rust/src/accounts/permission.rs
@@ -78,7 +78,6 @@ impl PermissionAccount {
 
 impl PythAccount for PermissionAccount {
     const ACCOUNT_TYPE: u32 = PC_ACCTYPE_PERMISSIONS;
-    // TODO: change?
-    // TODO: add feed_index when creating account
-    const INITIAL_SIZE: u32 = size_of::<PermissionAccount>() as u32;
+    const INITIAL_SIZE: u32 = Self::MIN_SIZE_WITH_LAST_FEED_INDEX as u32;
+    const NEW_ACCOUNT_SPACE: usize = Self::MIN_SIZE_WITH_LAST_FEED_INDEX;
 }
diff --git a/program/rust/src/error.rs b/program/rust/src/error.rs
@@ -54,6 +54,8 @@ pub enum OracleError {
     NeedsSuccesfulAggregation      = 620,
     #[error("MaxLastFeedIndexReached")]
     MaxLastFeedIndexReached        = 621,
+    #[error("FeedIndexAlreadyInitialized")]
+    FeedIndexAlreadyInitialized    = 622,
 }
 
 impl From<OracleError> for ProgramError {
diff --git a/program/rust/src/processor.rs b/program/rust/src/processor.rs
@@ -1,14 +1,25 @@
 use {
     crate::{
+        accounts::{
+            AccountHeader,
+            PermissionAccount,
+        },
+        deserialize::load_account_as_mut,
         error::OracleError,
         instruction::{
             load_command_header_checked,
             OracleCommand,
         },
+        utils::{
+            pyth_assert,
+            try_convert,
+        },
     },
     solana_program::{
         entrypoint::ProgramResult,
+        program::invoke,
         pubkey::Pubkey,
+        system_instruction,
         sysvar::slot_history::AccountInfo,
     },
 };
@@ -33,7 +44,6 @@ pub use add_publisher::{
     DISABLE_ACCUMULATOR_V2,
     ENABLE_ACCUMULATOR_V2,
 };
-use init_price_feed_index::init_price_feed_index;
 pub use {
     add_price::add_price,
     add_product::add_product,
@@ -55,6 +65,14 @@ pub use {
     },
     upd_product::upd_product,
 };
+use {
+    init_price_feed_index::init_price_feed_index,
+    solana_program::{
+        program_error::ProgramError,
+        rent::Rent,
+        sysvar::Sysvar,
+    },
+};
 
 /// Dispatch to the right instruction in the oracle.
 pub fn process_instruction(
@@ -90,3 +108,41 @@ pub fn process_instruction(
         InitPriceFeedIndex => init_price_feed_index(program_id, accounts, instruction_data),
     }
 }
+
+fn reserve_new_price_feed_index<'a>(
+    funding_account: &AccountInfo<'a>,
+    permissions_account: &AccountInfo<'a>,
+    system_program: &AccountInfo<'a>,
+) -> Result<u32, ProgramError> {
+    if permissions_account.data_len() < PermissionAccount::MIN_SIZE_WITH_LAST_FEED_INDEX {
+        let new_size = PermissionAccount::MIN_SIZE_WITH_LAST_FEED_INDEX;
+        let rent = Rent::get()?;
+        let new_minimum_balance = rent.minimum_balance(new_size);
+        let lamports_diff = new_minimum_balance.saturating_sub(permissions_account.lamports());
+        if lamports_diff > 0 {
+            invoke(
+                &system_instruction::transfer(
+                    funding_account.key,
+                    permissions_account.key,
+                    lamports_diff,
+                ),
+                &[
+                    funding_account.clone(),
+                    permissions_account.clone(),
+                    system_program.clone(),
+                ],
+            )?;
+        }
+
+        permissions_account.realloc(new_size, true)?;
+        let mut header = load_account_as_mut::<AccountHeader>(permissions_account)?;
+        header.size = try_convert(new_size)?;
+    }
+    let mut last_feed_index = PermissionAccount::load_last_feed_index_mut(permissions_account)?;
+    *last_feed_index += 1;
+    pyth_assert(
+        *last_feed_index < (1 << 28),
+        OracleError::MaxLastFeedIndexReached.into(),
+    )?;
+    Ok(*last_feed_index)
+}
diff --git a/program/rust/src/processor/add_price.rs b/program/rust/src/processor/add_price.rs
@@ -1,4 +1,5 @@
 use {
+    super::reserve_new_price_feed_index,
     crate::{
         accounts::{
             PriceAccount,
@@ -18,6 +19,7 @@ use {
             check_exponent_range,
             check_permissioned_funding_account,
             check_valid_funding_account,
+            check_valid_writable_account,
             pyth_assert,
         },
         OracleError,
@@ -31,9 +33,11 @@ use {
 };
 
 /// Add new price account to a product account
-// account[0] funding account       [signer writable]
-// account[1] product account       [signer writable]
-// account[2] new price account     [signer writable]
+// account[0] funding account        [signer writable]
+// account[1] product account        [writable]
+// account[2] new price account      [writable]
+// account[3] permissions account    [writable]
+// account[4] system program account []
 pub fn add_price(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
@@ -48,10 +52,11 @@ pub fn add_price(
     )?;
 
 
-    let (funding_account, product_account, price_account, permissions_account) = match accounts {
-        [x, y, z, p] => Ok((x, y, z, p)),
-        _ => Err(OracleError::InvalidNumberOfAccounts),
-    }?;
+    let (funding_account, product_account, price_account, permissions_account, system_program) =
+        match accounts {
+            [x, y, z, p, q] => Ok((x, y, z, p, q)),
+            _ => Err(OracleError::InvalidNumberOfAccounts),
+        }?;
 
     check_valid_funding_account(funding_account)?;
     check_permissioned_funding_account(
@@ -68,16 +73,24 @@ pub fn add_price(
         permissions_account,
         &cmd_args.header,
     )?;
+    check_valid_writable_account(program_id, permissions_account)?;
+    pyth_assert(
+        solana_program::system_program::check_id(system_program.key),
+        OracleError::InvalidSystemAccount.into(),
+    )?;
 
     let mut product_data =
         load_checked::<ProductAccount>(product_account, cmd_args.header.version)?;
 
+    // TODO: where is it created???
     let mut price_data = PriceAccount::initialize(price_account, cmd_args.header.version)?;
     price_data.exponent = cmd_args.exponent;
     price_data.price_type = cmd_args.price_type;
     price_data.product_account = *product_account.key;
     price_data.next_price_account = product_data.first_price_account;
     price_data.min_pub_ = PRICE_ACCOUNT_DEFAULT_MIN_PUB;
+    price_data.feed_index =
+        reserve_new_price_feed_index(funding_account, permissions_account, system_program)?;
     product_data.first_price_account = *price_account.key;
 
     Ok(())
diff --git a/program/rust/src/processor/init_price_feed_index.rs b/program/rust/src/processor/init_price_feed_index.rs
@@ -1,9 +1,7 @@
 use {
+    super::reserve_new_price_feed_index,
     crate::{
-        accounts::{
-            PermissionAccount,
-            PriceAccount,
-        },
+        accounts::PriceAccount,
         deserialize::{
             load,
             load_checked,
@@ -12,28 +10,25 @@ use {
         utils::{
             check_permissioned_funding_account,
             check_valid_funding_account,
+            check_valid_writable_account,
             pyth_assert,
         },
         OracleError,
     },
     solana_program::{
         account_info::AccountInfo,
         entrypoint::ProgramResult,
-        program::invoke,
         program_error::ProgramError,
         pubkey::Pubkey,
-        rent::Rent,
-        system_instruction,
-        sysvar::Sysvar,
     },
     std::mem::size_of,
 };
 
 /// Init price feed index
-// account[0] funding account       [signer writable]
-// account[1] price account         [writable]
-// account[2] permissions account   [writable]
-// account[3] system program account
+// account[0] funding account        [signer writable]
+// account[1] price account          [writable]
+// account[2] permissions account    [writable]
+// account[3] system program account []
 pub fn init_price_feed_index(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
@@ -59,42 +54,19 @@ pub fn init_price_feed_index(
         permissions_account,
         cmd,
     )?;
+    check_valid_writable_account(program_id, permissions_account)?;
     pyth_assert(
         solana_program::system_program::check_id(system_program.key),
         OracleError::InvalidSystemAccount.into(),
     )?;
 
-    if permissions_account.data_len() < PermissionAccount::MIN_SIZE_WITH_LAST_FEED_INDEX {
-        let new_size = PermissionAccount::MIN_SIZE_WITH_LAST_FEED_INDEX;
-        let rent = Rent::get()?;
-        let new_minimum_balance = rent.minimum_balance(new_size);
-        let lamports_diff = new_minimum_balance.saturating_sub(permissions_account.lamports());
-        if lamports_diff > 0 {
-            invoke(
-                &system_instruction::transfer(
-                    funding_account.key,
-                    permissions_account.key,
-                    lamports_diff,
-                ),
-                &[
-                    funding_account.clone(),
-                    permissions_account.clone(),
-                    system_program.clone(),
-                ],
-            )?;
-        }
-
-        permissions_account.realloc(new_size, true)?;
-    }
-    let mut last_feed_index = PermissionAccount::load_last_feed_index_mut(permissions_account)?;
-    *last_feed_index += 1;
+    let mut price_account_data = load_checked::<PriceAccount>(price_account, cmd.version)?;
     pyth_assert(
-        *last_feed_index < (1 << 28),
-        OracleError::MaxLastFeedIndexReached.into(),
+        price_account_data.feed_index == 0,
+        OracleError::FeedIndexAlreadyInitialized.into(),
     )?;
-
-    let mut price_account_data = load_checked::<PriceAccount>(price_account, cmd.version)?;
-    price_account_data.feed_index = *last_feed_index;
+    price_account_data.feed_index =
+        reserve_new_price_feed_index(funding_account, permissions_account, system_program)?;
 
     Ok(())
 }
diff --git a/program/rust/src/tests/pyth_simulator.rs b/program/rust/src/tests/pyth_simulator.rs
@@ -338,6 +338,7 @@ impl PythSimulator {
                 AccountMeta::new(product_keypair.pubkey(), true),
                 AccountMeta::new(price_keypair.pubkey(), true),
                 AccountMeta::new(self.get_permissions_pubkey(), false),
+                AccountMeta::new(system_program::id(), false),
             ],
         );
 
diff --git a/program/rust/src/tests/test_add_price.rs b/program/rust/src/tests/test_add_price.rs
@@ -62,6 +62,9 @@ fn test_add_price() {
     let mut permissions_setup = AccountSetup::new_permission(&program_id);
     let permissions_account = permissions_setup.as_account_info();
 
+    let mut system_program = AccountSetup::new_system_program();
+    let system_program_account = system_program.as_account_info();
+
     {
         let mut permissions_account_data =
             PermissionAccount::initialize(&permissions_account, PC_VERSION).unwrap();
@@ -82,17 +85,18 @@ fn test_add_price() {
     )
     .is_ok());
 
-    assert!(process_instruction(
+    process_instruction(
         &program_id,
         &[
             funding_account.clone(),
             product_account.clone(),
             price_account.clone(),
             permissions_account.clone(),
+            system_program_account.clone(),
         ],
-        instruction_data_add_price
+        instruction_data_add_price,
     )
-    .is_ok());
+    .unwrap();
 
     {
         let price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
@@ -105,17 +109,18 @@ fn test_add_price() {
         assert!(product_data.first_price_account == *price_account.key);
     }
 
-    assert!(process_instruction(
+    process_instruction(
         &program_id,
         &[
             funding_account.clone(),
             product_account.clone(),
             price_account_2.clone(),
             permissions_account.clone(),
+            system_program_account.clone(),
         ],
-        instruction_data_add_price
+        instruction_data_add_price,
     )
-    .is_ok());
+    .unwrap();
 
     {
         let price_data_2 = load_checked::<PriceAccount>(&price_account_2, PC_VERSION).unwrap();
@@ -137,6 +142,7 @@ fn test_add_price() {
                 product_account.clone(),
                 price_account.clone(),
                 permissions_account.clone(),
+                system_program_account.clone(),
                 permissions_account.clone(),
             ],
             instruction_data_add_price
@@ -153,6 +159,7 @@ fn test_add_price() {
                 product_account.clone(),
                 price_account.clone(),
                 permissions_account.clone(),
+                system_program_account.clone(),
             ],
             instruction_data_add_price
         ),
@@ -177,6 +184,7 @@ fn test_add_price() {
                 product_account.clone(),
                 price_account.clone(),
                 permissions_account.clone(),
+                system_program_account.clone(),
             ],
             instruction_data_add_price
         ),
@@ -202,6 +210,7 @@ fn test_add_price() {
                 product_account.clone(),
                 price_account.clone(),
                 permissions_account.clone(),
+                system_program_account.clone(),
             ],
             instruction_data_add_price
         ),
diff --git a/program/rust/src/tests/test_permission_migration.rs b/program/rust/src/tests/test_permission_migration.rs
diff --git a/program/rust/src/tests/test_upd_permissions.rs b/program/rust/src/tests/test_upd_permissions.rs
diff --git a/program/rust/src/tests/test_utils.rs b/program/rust/src/tests/test_utils.rs

Original file line number	Diff line number	Diff line change
`@@ -78,7 +78,6 @@ impl PermissionAccount {`
`78`	`78`
`79`	`79`	`impl PythAccount for PermissionAccount {`
`80`	`80`	`const ACCOUNT_TYPE: u32 = PC_ACCTYPE_PERMISSIONS;`
`81`		`- // TODO: change?`
`82`		`- // TODO: add feed_index when creating account`
`83`		`- const INITIAL_SIZE: u32 = size_of::<PermissionAccount>() as u32;`
	`81`	`+ const INITIAL_SIZE: u32 = Self::MIN_SIZE_WITH_LAST_FEED_INDEX as u32;`
	`82`	`+ const NEW_ACCOUNT_SPACE: usize = Self::MIN_SIZE_WITH_LAST_FEED_INDEX;`
`84`	`83`	`}`
Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,8 @@ pub enum OracleError {`
`54`	`54`	`NeedsSuccesfulAggregation = 620,`
`55`	`55`	`#[error("MaxLastFeedIndexReached")]`
`56`	`56`	`MaxLastFeedIndexReached = 621,`
	`57`	`+ #[error("FeedIndexAlreadyInitialized")]`
	`58`	`+ FeedIndexAlreadyInitialized = 622,`
`57`	`59`	`}`
`58`	`60`
`59`	`61`	`impl From<OracleError> for ProgramError {`