arcium-macros 0.9.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. `ArciumSignerAccount` 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! {
        #[derive(Clone, Copy, Debug)]
        pub enum CallbackError {
            InvalidCallbackTx,
        }

        impl CallbackError {
            pub fn name(&self) -> String {
                match self {
                    CallbackError::InvalidCallbackTx => "InvalidCallbackTx".to_string(),
                }
            }
        }

        impl ::core::fmt::Display for CallbackError {
            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
                f.write_str("Invalid callback transaction")
            }
        }

        impl From<CallbackError> for u32 {
            fn from(_: CallbackError) -> u32 {
                9999
            }
        }

        impl From<CallbackError> for ::anchor_lang::error::Error {
            fn from(e: CallbackError) -> Self {
                ::anchor_lang::error::Error::from(
                    ::anchor_lang::error::AnchorError {
                        error_name: e.name(),
                        error_code_number: e.into(),
                        error_msg: e.to_string(),
                        error_origin: None,
                        compared_values: None,
                    }
                )
            }
        }
    };

    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 ArciumSignerAccount {
            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()
}