Allow using custom discriminator

src/args/call_handler.rs

@@ -1,15 +1,9 @@
 use borsh::{BorshDeserialize, BorshSerialize};
 
-#[derive(BorshSerialize, BorshDeserialize, Clone, Copy)]
-pub enum Context {
-    Commit,
-    Undelegate,
-    Standalone,
-}
-
 #[derive(BorshSerialize, BorshDeserialize)]
 pub struct CallHandlerArgs {
     pub escrow_index: u8,
+    /// This is raw instruction data, it could include discriminator + args
+    /// or can be in any other custom format
     pub data: Vec<u8>,
-    pub context: Context,
 }

src/consts.rs

@@ -10,10 +10,6 @@ pub const PROTOCOL_FEES_PERCENTAGE: u8 = 10;
 /// The discriminator for the external undelegate instruction.
 pub const EXTERNAL_UNDELEGATE_DISCRIMINATOR: [u8; 8] = [196, 28, 41, 206, 48, 37, 51, 167];
 
-/// The discriminator for the external hook after finalization is complete
-/// For anchor: corresponds to function/instruction name delegation_program_call_handler
-pub const EXTERNAL_CALL_HANDLER_DISCRIMINATOR: [u8; 8] = [157, 197, 228, 30, 0, 80, 121, 135];
-
 /// The program ID of the delegation program.
 pub const DELEGATION_PROGRAM_ID: Pubkey = crate::id();
 

src/processor/call_handler.rs

@@ -1,5 +1,4 @@
 use crate::args::CallHandlerArgs;
-use crate::consts::EXTERNAL_CALL_HANDLER_DISCRIMINATOR;
 use crate::ephemeral_balance_seeds_from_payer;
 use crate::processor::utils::loaders::{
     load_initialized_validator_fees_vault, load_owned_pda, load_pda, load_signer,
@@ -109,10 +108,9 @@ pub fn process_call_handler(
             })
             .collect();
 
-    let data = [EXTERNAL_CALL_HANDLER_DISCRIMINATOR.to_vec(), data.to_vec()].concat();
     let handler_instruction = Instruction {
         program_id: *destination_program.key,
-        data,
+        data: args.data,
         accounts: accounts_meta,
     };
     let bump_slice = &[escrow_bump];

tests/integration/programs/test-delegation/src/lib.rs

@@ -1,7 +1,6 @@
 use anchor_lang::prelude::*;
 use ephemeral_rollups_sdk::anchor::{delegate, ephemeral};
 use ephemeral_rollups_sdk::cpi::DelegateConfig;
-use ephemeral_rollups_sdk::pda::ephemeral_balance_pda_from_payer;
 
 declare_id!("3vAK9JQiDsKoQNwmcfeEng4Cnv22pYuj1ASfso7U4ukF");
 
@@ -68,59 +67,42 @@ pub mod test_delegation {
     }
 
     /// Delegation program call handler
-    pub fn delegation_program_call_handler(
-        ctx: Context<DelegationProgramCallHandler>,
-        hook_args: delegation_program_utils::CallHandlerArgs,
+    #[instruction(discriminator = [1, 0, 1, 0])]
+    pub fn commit_base_action_handler(
+        ctx: Context<CommitBaseActionHandler>,
+        amount: u64,
     ) -> Result<()> {
-        let expected = ephemeral_balance_pda_from_payer(
-            ctx.accounts.escrow_authority.key,
-            hook_args.escrow_index,
+        msg!("commit_base_action_handler!");
+        let transfer_ctx = CpiContext::new(
+            ctx.accounts.system_program.to_account_info(),
+            Transfer {
+                from: ctx.accounts.escrow_account.to_account_info(),
+                to: ctx.accounts.destination_account.to_account_info(),
+            },
         );
-        if &expected != ctx.accounts.escrow_account.key {
-            Err(ProgramError::InvalidAccountData)
-        } else {
-            Ok(())
-        }?;
-
-        if !ctx.accounts.escrow_account.is_signer {
-            Err(ProgramError::MissingRequiredSignature)
-        } else {
-            Ok(())
-        }?;
-
-        match hook_args.context {
-            delegation_program_utils::Context::Commit => {
-                msg!("commit context");
-                let amount = u64::try_from_slice(&hook_args.data)?;
-                let transfer_ctx = CpiContext::new(
-                    ctx.accounts.system_program.to_account_info(),
-                    Transfer {
-                        from: ctx.accounts.escrow_account.to_account_info(),
-                        to: ctx.accounts.destination_account.to_account_info(),
-                    },
-                );
-                transfer(transfer_ctx, amount)?;
-            }
-            delegation_program_utils::Context::Undelegate => {
-                msg!("undelegate context");
-                let amount = u64::try_from_slice(&hook_args.data)?;
-                let transfer_ctx = CpiContext::new(
-                    ctx.accounts.system_program.to_account_info(),
-                    Transfer {
-                        from: ctx.accounts.escrow_account.to_account_info(),
-                        to: ctx.accounts.destination_account.to_account_info(),
-                    },
-                );
-                transfer(transfer_ctx, amount)?;
-
-                let counter_data = &mut ctx.accounts.counter.try_borrow_mut_data()?;
-                let mut counter = Counter::try_from_slice(&counter_data)?;
-                counter.count += 1;
-
-                counter_data.copy_from_slice(&counter.try_to_vec()?);
-            }
-            delegation_program_utils::Context::Standalone => msg!("standalone context"),
-        }
+
+        transfer(transfer_ctx, amount)
+    }
+
+    #[instruction(discriminator = [1, 0, 2, 0])]
+    pub fn undelegate_base_action_handler(
+        ctx: Context<UndelegateBaseActionHandler>,
+        amount: u64,
+    ) -> Result<()> {
+        msg!("undelegate_base_action_handler");
+        let transfer_ctx = CpiContext::new(
+            ctx.accounts.system_program.to_account_info(),
+            Transfer {
+                from: ctx.accounts.escrow_account.to_account_info(),
+                to: ctx.accounts.destination_account.to_account_info(),
+            },
+        );
+        transfer(transfer_ctx, amount)?;
+
+        let counter_data = &mut ctx.accounts.counter.try_borrow_mut_data()?;
+        let mut counter = Counter::try_from_slice(&counter_data)?;
+        counter.count += 1;
+        counter_data.copy_from_slice(&counter.try_to_vec()?);
 
         Ok(())
     }
@@ -194,16 +176,20 @@ pub struct Increment<'info> {
 }
 
 #[derive(Accounts)]
-#[instruction(hook_args: delegation_program_utils::CallHandlerArgs)]
-pub struct DelegationProgramCallHandler<'info> {
+pub struct CommitBaseActionHandler<'info> {
+    /// CHECK: The authority that owns the escrow account
+    pub escrow_authority: UncheckedAccount<'info>,
+    pub escrow_account: Signer<'info>,
+    /// CHECK: The destination account to transfer lamports to
+    #[account(mut)]
+    pub destination_account: AccountInfo<'info>,
+    pub system_program: Program<'info, System>,
+}
+
+#[derive(Accounts)]
+pub struct UndelegateBaseActionHandler<'info> {
     /// CHECK: The authority that owns the escrow account
     pub escrow_authority: UncheckedAccount<'info>,
-    #[account(
-        mut,
-        seeds = [b"balance", &escrow_authority.key().as_ref(), &[hook_args.escrow_index]],
-        seeds::program = delegation_program_utils::ID,
-        bump
-    )]
     pub escrow_account: Signer<'info>,
     /// CHECK: The destination account to transfer lamports to
     #[account(mut)]
@@ -222,18 +208,4 @@ mod delegation_program_utils {
     use anchor_lang::prelude::*;
 
     declare_id!("DELeGGvXpWV2fqJUhqcF5ZSYMS4JTLjteaAMARRSaeSh");
-
-    #[derive(AnchorSerialize, AnchorDeserialize)]
-    pub enum Context {
-        Commit,
-        Undelegate,
-        Standalone,
-    }
-
-    #[derive(AnchorSerialize, AnchorDeserialize)]
-    pub struct CallHandlerArgs {
-        pub escrow_index: u8,
-        pub data: Vec<u8>,
-        pub context: Context,
-    }
 }

tests/test_call_handler.rs

@@ -24,6 +24,9 @@ use solana_sdk::{
 
 mod fixtures;
 
+const COMMIT_HANDLER_DISCRIMINATOR: [u8; 4] = [1, 0, 1, 0];
+const UNDELEGATE_HANDLER_DISCRIMINATOR: [u8; 4] = [1, 0, 2, 0];
+
 // Mimic counter from test_delegation program
 #[derive(BorshSerialize, BorshDeserialize)]
 pub struct Counter {
@@ -296,13 +299,15 @@ async fn test_finalize_call_handler() {
         payer.pubkey(),         // escrow authority
         vec![
             AccountMeta::new(transfer_destination.pubkey(), false),
-            AccountMeta::new(DELEGATED_PDA_ID, false),
             AccountMeta::new_readonly(system_program::id(), false),
         ],
         CallHandlerArgs {
             escrow_index: 2, // undelegated escrow index,
-            data: to_vec(&PRIZE).unwrap(),
-            context: dlp::args::Context::Commit,
+            data: [
+                COMMIT_HANDLER_DISCRIMINATOR.to_vec(),
+                to_vec(&PRIZE).unwrap(),
+            ]
+            .concat(),
         },
     );
 
@@ -350,8 +355,11 @@ async fn test_undelegate_call_handler() {
         ],
         CallHandlerArgs {
             escrow_index: 2, // undelegated escrow index,
-            data: to_vec(&PRIZE).unwrap(),
-            context: dlp::args::Context::Undelegate,
+            data: [
+                UNDELEGATE_HANDLER_DISCRIMINATOR.to_vec(),
+                to_vec(&PRIZE).unwrap(),
+            ]
+            .concat(),
         },
     );
 
@@ -399,8 +407,7 @@ async fn test_finalize_invalid_escrow_call_handler() {
         vec![AccountMeta::new(transfer_destination.pubkey(), false)],
         CallHandlerArgs {
             escrow_index: 0,
-            data: vec![],
-            context: dlp::args::Context::Commit,
+            data: COMMIT_HANDLER_DISCRIMINATOR.to_vec(),
         },
     );
     let tx = Transaction::new_signed_with_payer(
@@ -432,8 +439,7 @@ async fn test_undelegate_invalid_escow_call_handler() {
         vec![AccountMeta::new(destination.pubkey(), false)],
         CallHandlerArgs {
             escrow_index: 0,
-            data: vec![],
-            context: dlp::args::Context::Commit,
+            data: UNDELEGATE_HANDLER_DISCRIMINATOR.to_vec(),
         },
     );
 
@@ -450,8 +456,11 @@ async fn test_undelegate_invalid_escow_call_handler() {
         vec![AccountMeta::new(destination.pubkey(), false)],
         CallHandlerArgs {
             escrow_index: 0,
-            data: to_vec(&PRIZE).unwrap(),
-            context: dlp::args::Context::Undelegate,
+            data: [
+                UNDELEGATE_HANDLER_DISCRIMINATOR.to_vec(),
+                to_vec(&PRIZE).unwrap(),
+            ]
+            .concat(),
         },
     );
     let tx = Transaction::new_signed_with_payer(