arcium-macros 0.5.3

//! 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];

            const LIGHTHOUSE_PROGRAM_ID: Pubkey = ::anchor_lang::solana_program::pubkey::Pubkey::from_str_const(
                "L2TExMFKdjpN9kozasaurPirfHy9P8sbXoAN1qA3S95",
            );

            // Verify the instruction sequence:
            // - Previous instruction must be from Arcium program (callback ix)
            // - Any subsequent instructions must be lighthouse assertions (max 2)

            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
            );

            // After callback instructions, only allow lighthouse assertions or no more instructions
            // Check all remaining instructions - they must all be lighthouse assertions
            let mut check_index = (curr_ix_index as usize) + 1;
            let mut lighthouse_count = 0;

            while let Ok(ix) = ::anchor_lang::solana_program::sysvar::instructions::load_instruction_at_checked(
                check_index,
                instructions_sysvar,
            ) {
                // All instructions after callbacks must be lighthouse assertions
                require!(
                    ix.program_id == LIGHTHOUSE_PROGRAM_ID,
                    CallbackError::InvalidCallbackTx
                );
                lighthouse_count += 1;
                // Allow up to 2 lighthouse assertions
                require!(
                    lighthouse_count <= 2,
                    CallbackError::InvalidCallbackTx
                );
                check_index += 1;
            }

            Ok(())
        }
    };

    quote! {
        #errors
        #signer_account_struct
        #validate_callback_ixs

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