solana: Add 2-step set token authority ixs

nvsriram · nvsriram · commit ae21768e5cec · 2024-12-03T14:23:20.000-05:00
diff --git a/solana/programs/example-native-token-transfers/src/error.rs b/solana/programs/example-native-token-transfers/src/error.rs
@@ -55,6 +55,10 @@ pub enum NTTError {
     NoRegisteredTransceivers,
     #[msg("NotPaused")]
     NotPaused,
+    #[msg("InvalidPendingTokenAuthority")]
+    InvalidPendingTokenAuthority,
+    #[msg("IncorrectRentPayer")]
+    IncorrectRentPayer,
 }
 
 impl From<ScalingError> for NTTError {
diff --git a/solana/programs/example-native-token-transfers/src/instructions/admin.rs b/solana/programs/example-native-token-transfers/src/instructions/admin.rs
@@ -10,6 +10,7 @@ use crate::{
     config::Config,
     error::NTTError,
     peer::NttManagerPeer,
+    pending_token_authority::PendingTokenAuthority,
     queue::{inbox::InboxRateLimit, outbox::OutboxRateLimit, rate_limit::RateLimitState},
     registered_transceiver::RegisteredTransceiver,
 };
@@ -174,20 +175,124 @@ pub struct SetTokenAuthority<'info> {
     /// CHECK: the mint address matches the config
     pub mint: InterfaceAccount<'info, token_interface::Mint>,
 
+    #[account(
+        seeds = [crate::TOKEN_AUTHORITY_SEED],
+        bump,
+        constraint = mint.mint_authority.unwrap() == token_authority.key() @ NTTError::InvalidMintAuthority
+    )]
+    /// CHECK: The constraints enforce this is valid mint authority
+    pub token_authority: UncheckedAccount<'info>,
+
+    /// CHECK: This account will be the signer in the [claim_token_authority] instruction.
+    pub new_authority: UncheckedAccount<'info>,
+}
+
+#[derive(Accounts)]
+pub struct SetTokenAuthorityChecked<'info> {
+    pub common: SetTokenAuthority<'info>,
+
+    #[account(mut)]
+    pub payer: Signer<'info>,
+
+    #[account(
+        init_if_needed,
+        space = 8 + PendingTokenAuthority::INIT_SPACE,
+        payer = payer,
+        seeds = [PendingTokenAuthority::SEED_PREFIX],
+        bump
+     )]
+    pub pending_token_authority: Account<'info, PendingTokenAuthority>,
+
+    pub system_program: Program<'info, System>,
+}
+
+pub fn set_token_authority(ctx: Context<SetTokenAuthorityChecked>) -> Result<()> {
+    ctx.accounts
+        .pending_token_authority
+        .set_inner(PendingTokenAuthority {
+            bump: ctx.bumps.pending_token_authority,
+            pending_authority: ctx.accounts.common.new_authority.key(),
+            rent_payer: ctx.accounts.payer.key(),
+        });
+    Ok(())
+}
+
+#[derive(Accounts)]
+pub struct SetTokenAuthorityUnchecked<'info> {
+    pub common: SetTokenAuthority<'info>,
+
     pub token_program: Interface<'info, token_interface::TokenInterface>,
+}
+
+pub fn set_token_authority_one_step_unchecked(
+    ctx: Context<SetTokenAuthorityUnchecked>,
+) -> Result<()> {
+    token_interface::set_authority(
+        CpiContext::new_with_signer(
+            ctx.accounts.token_program.to_account_info(),
+            token_interface::SetAuthority {
+                account_or_mint: ctx.accounts.common.mint.to_account_info(),
+                current_authority: ctx.accounts.common.token_authority.to_account_info(),
+            },
+            &[&[
+                crate::TOKEN_AUTHORITY_SEED,
+                &[ctx.bumps.common.token_authority],
+            ]],
+        ),
+        AuthorityType::MintTokens,
+        Some(ctx.accounts.common.new_authority.key()),
+    )
+}
+
+#[derive(Accounts)]
+pub struct ClaimTokenAuthority<'info> {
+    #[account(
+        constraint = config.paused @ NTTError::NotPaused,
+    )]
+    pub config: Account<'info, Config>,
+
+    #[account(
+        mut,
+        address = pending_token_authority.rent_payer @ NTTError::IncorrectRentPayer,
+    )]
+    pub payer: Signer<'info>,
+
+    #[account(
+        mut,
+        address = config.mint,
+    )]
+    /// CHECK: the mint address matches the config
+    pub mint: InterfaceAccount<'info, token_interface::Mint>,
 
     #[account(
         seeds = [crate::TOKEN_AUTHORITY_SEED],
         bump,
     )]
-    /// CHECK: token_program will ensure this is the valid mint_authority.
+    /// CHECK: The seeds constraint enforces that this is the correct address
     pub token_authority: UncheckedAccount<'info>,
 
-    /// CHECK: This is unsafe as is so should be used with caution
-    pub new_authority: UncheckedAccount<'info>,
+    #[account(
+        constraint = (
+            new_authority.key() == pending_token_authority.pending_authority
+            || new_authority.key() == token_authority.key()
+        ) @ NTTError::InvalidPendingTokenAuthority
+    )]
+    pub new_authority: Signer<'info>,
+
+    #[account(
+        mut,
+        seeds = [PendingTokenAuthority::SEED_PREFIX],
+        bump = pending_token_authority.bump,
+        close = payer
+     )]
+    pub pending_token_authority: Account<'info, PendingTokenAuthority>,
+
+    pub token_program: Interface<'info, token_interface::TokenInterface>,
+
+    pub system_program: Program<'info, System>,
 }
 
-pub fn set_token_authority_one_step_unchecked(ctx: Context<SetTokenAuthority>) -> Result<()> {
+pub fn claim_token_authority(ctx: Context<ClaimTokenAuthority>) -> Result<()> {
     token_interface::set_authority(
         CpiContext::new_with_signer(
             ctx.accounts.token_program.to_account_info(),
diff --git a/solana/programs/example-native-token-transfers/src/lib.rs b/solana/programs/example-native-token-transfers/src/lib.rs
@@ -18,6 +18,7 @@ pub mod error;
 pub mod instructions;
 pub mod messages;
 pub mod peer;
+pub mod pending_token_authority;
 pub mod queue;
 pub mod registered_transceiver;
 pub mod transceivers;
@@ -126,10 +127,20 @@ pub mod example_native_token_transfers {
         instructions::claim_ownership(ctx)
     }
 
-    pub fn set_token_authority_one_step_unchecked(ctx: Context<SetTokenAuthority>) -> Result<()> {
+    pub fn set_token_authority(ctx: Context<SetTokenAuthorityChecked>) -> Result<()> {
+        instructions::set_token_authority(ctx)
+    }
+
+    pub fn set_token_authority_one_step_unchecked(
+        ctx: Context<SetTokenAuthorityUnchecked>,
+    ) -> Result<()> {
         instructions::set_token_authority_one_step_unchecked(ctx)
     }
 
+    pub fn claim_token_authority(ctx: Context<ClaimTokenAuthority>) -> Result<()> {
+        instructions::claim_token_authority(ctx)
+    }
+
     pub fn set_paused(ctx: Context<SetPaused>, pause: bool) -> Result<()> {
         instructions::set_paused(ctx, pause)
     }
diff --git a/solana/programs/example-native-token-transfers/src/pending_token_authority.rs b/solana/programs/example-native-token-transfers/src/pending_token_authority.rs
@@ -0,0 +1,13 @@
+use anchor_lang::prelude::*;
+
+#[account]
+#[derive(InitSpace)]
+pub struct PendingTokenAuthority {
+    pub bump: u8,
+    pub pending_authority: Pubkey,
+    pub rent_payer: Pubkey,
+}
+
+impl PendingTokenAuthority {
+    pub const SEED_PREFIX: &'static [u8] = b"pending_token_authority";
+}
diff --git a/solana/ts/idl/3_0_0/json/example_native_token_transfers.json b/solana/ts/idl/3_0_0/json/example_native_token_transfers.json
@@ -719,17 +719,109 @@
       ],
       "args": []
     },
+    {
+      "name": "setTokenAuthority",
+      "accounts": [
+        {
+          "name": "common",
+          "accounts": [
+            {
+              "name": "config",
+              "isMut": false,
+              "isSigner": false
+            },
+            {
+              "name": "owner",
+              "isMut": false,
+              "isSigner": true
+            },
+            {
+              "name": "mint",
+              "isMut": true,
+              "isSigner": false
+            },
+            {
+              "name": "tokenAuthority",
+              "isMut": false,
+              "isSigner": false
+            },
+            {
+              "name": "newAuthority",
+              "isMut": false,
+              "isSigner": false
+            }
+          ]
+        },
+        {
+          "name": "payer",
+          "isMut": true,
+          "isSigner": true
+        },
+        {
+          "name": "pendingTokenAuthority",
+          "isMut": true,
+          "isSigner": false
+        },
+        {
+          "name": "systemProgram",
+          "isMut": false,
+          "isSigner": false
+        }
+      ],
+      "args": []
+    },
     {
       "name": "setTokenAuthorityOneStepUnchecked",
+      "accounts": [
+        {
+          "name": "common",
+          "accounts": [
+            {
+              "name": "config",
+              "isMut": false,
+              "isSigner": false
+            },
+            {
+              "name": "owner",
+              "isMut": false,
+              "isSigner": true
+            },
+            {
+              "name": "mint",
+              "isMut": true,
+              "isSigner": false
+            },
+            {
+              "name": "tokenAuthority",
+              "isMut": false,
+              "isSigner": false
+            },
+            {
+              "name": "newAuthority",
+              "isMut": false,
+              "isSigner": false
+            }
+          ]
+        },
+        {
+          "name": "tokenProgram",
+          "isMut": false,
+          "isSigner": false
+        }
+      ],
+      "args": []
+    },
+    {
+      "name": "claimTokenAuthority",
       "accounts": [
         {
           "name": "config",
           "isMut": false,
           "isSigner": false
         },
         {
-          "name": "owner",
-          "isMut": false,
+          "name": "payer",
+          "isMut": true,
           "isSigner": true
         },
         {
@@ -738,17 +830,27 @@
           "isSigner": false
         },
         {
-          "name": "tokenProgram",
+          "name": "tokenAuthority",
           "isMut": false,
           "isSigner": false
         },
         {
-          "name": "tokenAuthority",
+          "name": "newAuthority",
+          "isMut": false,
+          "isSigner": true
+        },
+        {
+          "name": "pendingTokenAuthority",
+          "isMut": true,
+          "isSigner": false
+        },
+        {
+          "name": "tokenProgram",
           "isMut": false,
           "isSigner": false
         },
         {
-          "name": "newAuthority",
+          "name": "systemProgram",
           "isMut": false,
           "isSigner": false
         }
@@ -1412,6 +1514,26 @@
         ]
       }
     },
+    {
+      "name": "PendingTokenAuthority",
+      "type": {
+        "kind": "struct",
+        "fields": [
+          {
+            "name": "bump",
+            "type": "u8"
+          },
+          {
+            "name": "pendingAuthority",
+            "type": "publicKey"
+          },
+          {
+            "name": "rentPayer",
+            "type": "publicKey"
+          }
+        ]
+      }
+    },
     {
       "name": "InboxItem",
       "type": {
@@ -2060,6 +2182,16 @@
       "code": 6024,
       "name": "NotPaused",
       "msg": "NotPaused"
+    },
+    {
+      "code": 6025,
+      "name": "InvalidPendingTokenAuthority",
+      "msg": "InvalidPendingTokenAuthority"
+    },
+    {
+      "code": 6026,
+      "name": "IncorrectRentPayer",
+      "msg": "IncorrectRentPayer"
     }
   ]
 }
diff --git a/solana/ts/idl/3_0_0/ts/example_native_token_transfers.ts b/solana/ts/idl/3_0_0/ts/example_native_token_transfers.ts
