solana: allow pools to specify writable accounts

chains/solana/contracts/programs/ccip-router/src/instructions/v1/pools.rs

@@ -188,29 +188,48 @@ pub fn validate_and_parse_token_accounts<'info>(
                 .map_err(|_| CcipRouterError::InvalidInputsLookupTableAccounts)?;
 
         // reconstruct + validate expected values in token pool lookup table
-        // base set of constant accounts (8)
+        // base set of constant accounts (9)
         // + additional constant accounts (remaining_accounts) that are not required but may be used for additional token pool functionality (like CPI)
-        let mut expected_entries = vec![
-            lookup_table.key(),
-            token_admin_registry.key(),
-            pool_program.key(),
-            pool_config.key(),
-            pool_token_account.key(),
-            pool_signer.key(),
-            token_program.key(),
-            mint.key(),
-            fee_token_config.key(),
+        let required_entries = [
+            lookup_table,
+            token_admin_registry,
+            pool_program,
+            pool_config,
+            pool_token_account,
+            pool_signer,
+            token_program,
+            mint,
+            fee_token_config,
         ];
-        let mut remaining_keys: Vec<Pubkey> = remaining_accounts.iter().map(|x| x.key()).collect();
-        expected_entries.append(&mut remaining_keys);
-        require!(
-            lookup_table_account.addresses.len() == expected_entries.len(),
-            CcipRouterError::InvalidInputsLookupTableAccounts
-        );
-        require!(
-            lookup_table_account.addresses.as_ref() == expected_entries,
-            CcipRouterError::InvalidInputsLookupTableAccounts
-        );
+        {
+            // validate pool addresses
+            let mut expected_keys: Vec<Pubkey> = required_entries.iter().map(|x| x.key()).collect();
+            let mut remaining_keys: Vec<Pubkey> =
+                remaining_accounts.iter().map(|x| x.key()).collect();
+            expected_keys.append(&mut remaining_keys);
+            require!(
+                lookup_table_account.addresses.len() == expected_keys.len(),
+                CcipRouterError::InvalidInputsLookupTableAccounts
+            );
+            require!(
+                lookup_table_account.addresses.as_ref() == expected_keys,
+                CcipRouterError::InvalidInputsLookupTableAccounts
+            );
+        }
+        {
+            // validate pool address writable
+            let mut expected_is_writable: Vec<bool> =
+                required_entries.iter().map(|x| x.is_writable).collect();
+            let mut remaining_is_writable: Vec<bool> =
+                remaining_accounts.iter().map(|x| x.is_writable).collect();
+            expected_is_writable.append(&mut remaining_is_writable);
+            for (i, is_writable) in expected_is_writable.iter().enumerate() {
+                require!(
+                    token_admin_registry_account.is_writable(i as u8) == *is_writable,
+                    CcipRouterError::InvalidInputsLookupTableAccountWritable
+                );
+            }
+        }
     }
 
     Ok(TokenAccounts {

chains/solana/contracts/programs/ccip-router/src/instructions/v1/token_admin_registry.rs

@@ -1,9 +1,16 @@
 use anchor_lang::prelude::*;
+use anchor_spl::{
+    associated_token::get_associated_token_address_with_program_id, token::spl_token::native_mint,
+};
+use bytemuck::Zeroable;
+use solana_program::{address_lookup_table::state::AddressLookupTable, log::sol_log};
 
 use crate::{
     AcceptAdminRoleTokenAdminRegistry, AdministratorRegistered, AdministratorTransferRequested,
     AdministratorTransferred, CcipRouterError, ModifyTokenAdminRegistry, PoolSet,
     RegisterTokenAdminRegistryViaGetCCIPAdmin, RegisterTokenAdminRegistryViaOwner,
+    SetPoolTokenAdminRegistry, CCIP_TOKENPOOL_CONFIG, CCIP_TOKENPOOL_SIGNER,
+    FEE_BILLING_TOKEN_CONFIG, TOKEN_ADMIN_REGISTRY_SEED,
 };
 
 pub fn register_token_admin_registry_via_get_ccip_admin(
@@ -53,20 +60,91 @@ pub fn register_token_admin_registry_via_owner(
 }
 
 pub fn set_pool(
-    ctx: Context<ModifyTokenAdminRegistry>,
+    ctx: Context<SetPoolTokenAdminRegistry>,
     mint: Pubkey,
-    pool_lookup_table: Pubkey,
+    writable_indexes: Vec<u8>,
 ) -> Result<()> {
+    // set new lookup table
     let token_admin_registry = &mut ctx.accounts.token_admin_registry;
     let previous_pool = token_admin_registry.lookup_table;
-    token_admin_registry.lookup_table = pool_lookup_table;
+    let new_pool = ctx.accounts.pool_lookuptable.key();
+    token_admin_registry.lookup_table = new_pool;
 
-    // TODO: Validate here that the lookup table has everything
+    // set writable indexes
+    token_admin_registry.reset_writable();
+    for ind in writable_indexes {
+        token_admin_registry.set_writable(ind)
+    }
+
+    // validate lookup table contains minium required accounts if not zero address
+    if new_pool != Pubkey::zeroed() {
+        // deserialize lookup table account
+        let lookup_table_data = &mut &ctx.accounts.pool_lookuptable.data.borrow()[..];
+        let lookup_table_account: AddressLookupTable =
+            AddressLookupTable::deserialize(lookup_table_data)
+                .map_err(|_| CcipRouterError::InvalidInputsLookupTableAccounts)?;
+        require!(
+            lookup_table_account.addresses.len() >= 9,
+            CcipRouterError::InvalidInputsLookupTableAccounts
+        );
+
+        // calculate or retrieve expected addresses
+        let (token_admin_registry, _) = Pubkey::find_program_address(
+            &[TOKEN_ADMIN_REGISTRY_SEED, mint.as_ref()],
+            ctx.program_id,
+        );
+        let pool_program = lookup_table_account.addresses[2]; // cannot be calculated, can be custom per pool
+        let token_program = lookup_table_account.addresses[6]; // cannot be calculated, can be custom per token
+        let (pool_config, _) =
+            Pubkey::find_program_address(&[CCIP_TOKENPOOL_CONFIG, mint.as_ref()], &pool_program);
+        let (pool_signer, _) =
+            Pubkey::find_program_address(&[CCIP_TOKENPOOL_SIGNER, mint.as_ref()], &pool_program);
+        let (fee_billing_config, _) = Pubkey::find_program_address(
+            &[
+                FEE_BILLING_TOKEN_CONFIG,
+                if mint.key() == Pubkey::default() {
+                    native_mint::ID.as_ref() // pre-2022 WSOL
+                } else {
+                    mint.as_ref()
+                },
+            ],
+            ctx.program_id,
+        );
+
+        let min_accounts = [
+            ctx.accounts.pool_lookuptable.key(),
+            token_admin_registry,
+            pool_program,
+            pool_config,
+            get_associated_token_address_with_program_id(
+                &pool_signer.key(),
+                &mint.key(),
+                &token_program.key(),
+            ),
+            pool_signer,
+            token_program,
+            mint,
+            fee_billing_config,
+        ];
+
+        for (i, acc) in min_accounts.iter().enumerate() {
+            // for easier debugging UX
+            if lookup_table_account.addresses[i] != *acc {
+                sol_log(&i.to_string());
+                sol_log(&acc.to_string());
+                sol_log(&lookup_table_account.addresses[i].to_string());
+            }
+            require!(
+                lookup_table_account.addresses[i] == *acc,
+                CcipRouterError::InvalidInputsLookupTableAccounts
+            );
+        }
+    }
 
     emit!(PoolSet {
         token: mint,
         previous_pool_lookup_table: previous_pool,
-        new_pool_lookup_table: pool_lookup_table,
+        new_pool_lookup_table: new_pool,
     });
 
     Ok(())

chains/solana/contracts/programs/ccip-router/src/lib.rs

@@ -327,12 +327,13 @@ pub mod ccip_router {
     /// * `ctx` - The context containing the accounts required for setting the pool.
     /// * `mint` - The public key of the token mint.
     /// * `pool_lookup_table` - The public key of the pool lookup table, this address will be used for validations when interacting with the pool.
+    /// * `is_writable` - index of account in lookup table that is writable
     pub fn set_pool(
-        ctx: Context<ModifyTokenAdminRegistry>,
+        ctx: Context<SetPoolTokenAdminRegistry>,
         mint: Pubkey,
-        pool_lookup_table: Pubkey,
+        writable_indexes: Vec<u8>,
     ) -> Result<()> {
-        v1::token_admin_registry::set_pool(ctx, mint, pool_lookup_table)
+        v1::token_admin_registry::set_pool(ctx, mint, writable_indexes)
     }
 
     /// Transfers the admin role of the token admin registry to a new admin.
@@ -628,6 +629,8 @@ pub enum CcipRouterError {
     InvalidInputsTokenAdminRegistryAccounts,
     #[msg("Invalid LookupTable account")]
     InvalidInputsLookupTableAccounts,
+    #[msg("Invalid LookupTable account writable access")]
+    InvalidInputsLookupTableAccountWritable,
     #[msg("Cannot send zero tokens")]
     InvalidInputsTokenAmount,
     #[msg("Release or mint balance mismatch")]

chains/solana/contracts/programs/ccip-router/src/token_context.rs

@@ -13,6 +13,35 @@ pub struct TokenAdminRegistry {
     pub administrator: Pubkey,
     pub pending_administrator: Pubkey,
     pub lookup_table: Pubkey,
+    // binary representation of indexes that are writable in token pool lookup table
+    // lookup table can store 256 addresses
+    pub writable_indexes: [u128; 2],
+}
+
+impl TokenAdminRegistry {
+    // set writable inserts bits from left to right
+    // index 0 is left-most bit
+    pub fn set_writable(&mut self, index: u8) {
+        match index < 128 {
+            true => {
+                self.writable_indexes[0] |= 1 << (127 - index);
+            }
+            false => {
+                self.writable_indexes[1] |= 1 << (255 - index);
+            }
+        }
+    }
+
+    pub fn is_writable(&self, index: u8) -> bool {
+        match index < 128 {
+            true => self.writable_indexes[0] & 1 << (127 - index) != 0,
+            false => self.writable_indexes[1] & 1 << (255 - index) != 0,
+        }
+    }
+
+    pub fn reset_writable(&mut self) {
+        self.writable_indexes = [0, 0];
+    }
 }
 
 #[derive(Accounts)]
@@ -73,6 +102,22 @@ pub struct ModifyTokenAdminRegistry<'info> {
     pub authority: Signer<'info>,
 }
 
+#[derive(Accounts)]
+#[instruction(mint: Pubkey)]
+pub struct SetPoolTokenAdminRegistry<'info> {
+    #[account(
+        mut,
+        seeds = [TOKEN_ADMIN_REGISTRY_SEED, mint.as_ref()],
+        bump,
+        constraint = valid_version(token_admin_registry.version, MAX_TOKEN_REGISTRY_V) @ CcipRouterError::InvalidInputs,
+    )]
+    pub token_admin_registry: Account<'info, TokenAdminRegistry>,
+    /// CHECK: anchor does not support automatic lookup table deserialization
+    pub pool_lookuptable: UncheckedAccount<'info>,
+    #[account(mut, address = token_admin_registry.administrator @ CcipRouterError::Unauthorized)]
+    pub authority: Signer<'info>,
+}
+
 #[derive(Accounts)]
 #[instruction(mint: Pubkey)]
 pub struct AcceptAdminRoleTokenAdminRegistry<'info> {
@@ -112,3 +157,73 @@ pub struct SetTokenBillingConfig<'info> {
     pub authority: Signer<'info>,
     pub system_program: Program<'info, System>,
 }
+
+#[cfg(test)]
+mod tests {
+    use bytemuck::Zeroable;
+    use solana_program::pubkey::Pubkey;
+
+    use crate::TokenAdminRegistry;
+    #[test]
+    fn set_writable() {
+        let mut state = TokenAdminRegistry {
+            version: 0,
+            administrator: Pubkey::zeroed(),
+            pending_administrator: Pubkey::zeroed(),
+            lookup_table: Pubkey::zeroed(),
+            writable_indexes: [0, 0],
+        };
+
+        state.set_writable(0);
+        state.set_writable(128);
+        assert_eq!(state.writable_indexes[0], 2u128.pow(127));
+        assert_eq!(state.writable_indexes[1], 2u128.pow(127));
+
+        state.reset_writable();
+        assert_eq!(state.writable_indexes[0], 0);
+        assert_eq!(state.writable_indexes[1], 0);
+
+        state.set_writable(0);
+        state.set_writable(2);
+        state.set_writable(127);
+        state.set_writable(128);
+        state.set_writable(2 + 128);
+        state.set_writable(255);
+        assert_eq!(
+            state.writable_indexes[0],
+            2u128.pow(127) + 2u128.pow(127 - 2) + 2u128.pow(0)
+        );
+        assert_eq!(
+            state.writable_indexes[1],
+            2u128.pow(127) + 2u128.pow(127 - 2) + 2u128.pow(0)
+        );
+    }
+
+    #[test]
+    fn check_writable() {
+        let state = TokenAdminRegistry {
+            version: 0,
+            administrator: Pubkey::zeroed(),
+            pending_administrator: Pubkey::zeroed(),
+            lookup_table: Pubkey::zeroed(),
+            writable_indexes: [
+                2u128.pow(127 - 7) + 2u128.pow(127 - 2) + 2u128.pow(127 - 4),
+                2u128.pow(127 - 8) + 2u128.pow(127 - 56) + 2u128.pow(127 - 100),
+            ],
+        };
+
+        assert_eq!(state.is_writable(0), false);
+        assert_eq!(state.is_writable(128), false);
+        assert_eq!(state.is_writable(255), false);
+
+        assert_eq!(state.writable_indexes[0].count_ones(), 3);
+        assert_eq!(state.writable_indexes[1].count_ones(), 3);
+
+        assert_eq!(state.is_writable(7), true);
+        assert_eq!(state.is_writable(2), true);
+        assert_eq!(state.is_writable(4), true);
+        assert_eq!(state.is_writable(128 + 8), true);
+        assert_eq!(state.is_writable(128 + 56), true);
+        assert_eq!(state.is_writable(128 + 100), true);
+    }
+}