arcium-macros 0.4.0

//! Implementation of the `#[arcium_program]` attribute macro.
//!
//! This module provides the core functionality for transforming a standard Anchor program module
//! into an Arcium-compatible program. It wraps Anchor's `#[program]` attribute and injects
//! additional code required for Arcium integration.
//!
//! ## Generated Code
//!
//! The `#[arcium_program]` macro generates three key components:
//!
//! 1. `CallbackError` enum: Custom error types for callback validation
//! 2. `SignerAccount` struct: PDA signer account (owned by user program, not Arcium)
//! 3. `validate_callback_ixs` function: Security validation for callback transactions
//!
//! ## Integration with Anchor
//!
//! This macro expands to include Anchor's `#[program]` attribute, so programs using
//! `#[arcium_program]` get all Anchor functionality plus Arcium-specific additions.

use proc_macro::TokenStream;
use quote::quote;
use syn::Item;

pub fn arcium_program_macro(input: &mut Item) -> TokenStream {
    let errors = quote! {
        #[error_code]
        pub enum CallbackError {
            #[msg("Invalid callback transaction")]
            InvalidCallbackTx,
        }
    };

    let signer_account_struct = quote! {
        // We declare this here instead of just importing it from arcium, as otherwise
        // anchor automatically enforces its owner to be the arcium program, which is not what we want.
        #[::anchor_lang::prelude::account]
        pub struct SignerAccount {
            bump: u8,
        }
    };

    let validate_callback_ixs = quote! {
        /// Validates the transaction structure for `arcium_callback` instructions.
        ///
        /// ## Security Rationale
        ///
        /// This prevents malicious actors from calling callback instructions directly. Callbacks must
        /// only execute after Arcium's callback_computation instruction verifies the computation result.
        ///
        /// ## Required Transaction Structure
        /// 1. Previous instruction must be Arcium's `callback_computation` instruction
        /// 2. No instructions can follow the callback (prevents re-entrancy or additional malicious logic)
        /// 3. Callback cannot be the first instruction (must follow Arcium callback computation)
        fn validate_callback_ixs(instructions_sysvar: &AccountInfo, arcium_program: &Pubkey) -> Result<()> {
            const ARCIUM_CALLBACK_COMPUTATION_DISCRIMINATOR: [u8; 8] =
            [11, 224, 42, 236, 0, 154, 74, 163];

            let curr_ix_index = ::anchor_lang::solana_program::sysvar::instructions::load_current_index_checked(instructions_sysvar)?;
            require!(curr_ix_index != 0, CallbackError::InvalidCallbackTx);

            // The previous ix must be a callback computation ix in the Arcium program
            let prev_ix = ::anchor_lang::solana_program::sysvar::instructions::load_instruction_at_checked((curr_ix_index as usize) - 1, instructions_sysvar)?;

            require!(
                prev_ix.program_id == *arcium_program,
                CallbackError::InvalidCallbackTx,
            );
            require!(
                prev_ix.data[0..8] == ARCIUM_CALLBACK_COMPUTATION_DISCRIMINATOR,
                CallbackError::InvalidCallbackTx
            );

            // There must be no more ixs after this (and we didn't fail for some other reason)
            require!(
                ::anchor_lang::solana_program::sysvar::instructions::load_instruction_at_checked((curr_ix_index as usize) + 1, instructions_sysvar)
                    .is_err_and(|e| { e == ProgramError::InvalidArgument }),
                CallbackError::InvalidCallbackTx
            );

            Ok(())
        }
    };

    quote! {
        #errors
        #signer_account_struct
        #validate_callback_ixs

        #[program]
        #input
    }
    .into()
}