magicsvm 0.2.0

use {
    borsh::BorshDeserialize,
    dlp_api::{
        args::{
            DelegateWithActionsArgs, MaybeEncryptedAccountMeta, MaybeEncryptedIxData,
            MaybeEncryptedPubkey, PostDelegationActions,
        },
        compact,
        discriminator::DlpDiscriminator,
        encryption::{self, KEY_LEN},
    },
    solana_address::Address,
    solana_message::{AccountMeta, Instruction},
    solana_transaction::{InstructionError, TransactionError},
};

/// Extracts the post-delegation actions carried by a `DelegateWithActions`
/// instruction.
///
/// Returns `None` if the instruction is not a `DelegateWithActions` call or if
/// its arguments fail to deserialize.
pub fn post_delegation_actions(
    discriminator: DlpDiscriminator,
    instruction_data: &[u8],
) -> Option<PostDelegationActions> {
    if discriminator != DlpDiscriminator::DelegateWithActions {
        return None;
    }

    DelegateWithActionsArgs::try_from_slice(instruction_data.get(8..)?)
        .ok()
        .map(|args| args.actions)
}

/// Decrypts a set of post-delegation actions into executable instructions.
///
/// The validator's ed25519 key pair is converted to x25519 and used to decrypt
/// the encrypted pubkeys, account metas and instruction data. `validator_pubkey`
/// and `validator_secret` are the validator's ed25519 public and secret keys.
///
/// Returns [`TransactionError::InstructionError`] with
/// [`InstructionError::InvalidInstructionData`] if key conversion or decryption
/// fails.
pub fn decrypt_post_delegation_instructions(
    actions: PostDelegationActions,
    validator_pubkey: &[u8; KEY_LEN],
    validator_secret: &[u8],
) -> Result<Vec<Instruction>, TransactionError> {
    let validator_x25519_pubkey =
        encryption::ed25519_pubkey_to_x25519(validator_pubkey).map_err(|_| {
            TransactionError::InstructionError(0, InstructionError::InvalidInstructionData)
        })?;
    let validator_x25519_secret =
        encryption::ed25519_secret_to_x25519(validator_secret).map_err(|_| {
            TransactionError::InstructionError(0, InstructionError::InvalidInstructionData)
        })?;
    let mut pubkeys: Vec<Address> = actions
        .signers
        .iter()
        .map(|pubkey| Address::new_from_array(*pubkey))
        .collect();
    for pubkey in actions.non_signers {
        pubkeys.push(decrypt_post_delegation_pubkey(
            pubkey,
            &validator_x25519_pubkey,
            &validator_x25519_secret,
        )?);
    }

    actions
        .instructions
        .into_iter()
        .map(|instruction| {
            let program_id = resolve_post_delegation_pubkey(&pubkeys, instruction.program_id)?;
            let accounts = instruction
                .accounts
                .into_iter()
                .map(|account| {
                    decrypt_post_delegation_account_meta(
                        &pubkeys,
                        account,
                        &validator_x25519_pubkey,
                        &validator_x25519_secret,
                    )
                })
                .collect::<Result<Vec<_>, _>>()?;
            let data = decrypt_post_delegation_data(
                instruction.data,
                &validator_x25519_pubkey,
                &validator_x25519_secret,
            )?;

            Ok(Instruction {
                program_id,
                accounts,
                data,
            })
        })
        .collect()
}

fn decrypt_post_delegation_pubkey(
    pubkey: MaybeEncryptedPubkey,
    validator_x25519_pubkey: &[u8; KEY_LEN],
    validator_x25519_secret: &[u8; KEY_LEN],
) -> Result<Address, TransactionError> {
    match pubkey {
        MaybeEncryptedPubkey::ClearText(pubkey) => Ok(Address::new_from_array(pubkey)),
        MaybeEncryptedPubkey::Encrypted(buffer) => {
            let decrypted = encryption::decrypt(
                buffer.as_bytes(),
                validator_x25519_pubkey,
                validator_x25519_secret,
            )
            .map_err(|_| {
                TransactionError::InstructionError(0, InstructionError::InvalidInstructionData)
            })?;
            let pubkey = <[u8; KEY_LEN]>::try_from(decrypted.as_slice()).map_err(|_| {
                TransactionError::InstructionError(0, InstructionError::InvalidInstructionData)
            })?;
            Ok(Address::new_from_array(pubkey))
        }
    }
}

fn decrypt_post_delegation_account_meta(
    pubkeys: &[Address],
    account: MaybeEncryptedAccountMeta,
    validator_x25519_pubkey: &[u8; KEY_LEN],
    validator_x25519_secret: &[u8; KEY_LEN],
) -> Result<AccountMeta, TransactionError> {
    let account = match account {
        MaybeEncryptedAccountMeta::ClearText(account) => account,
        MaybeEncryptedAccountMeta::Encrypted(buffer) => {
            let decrypted = encryption::decrypt(
                buffer.as_bytes(),
                validator_x25519_pubkey,
                validator_x25519_secret,
            )
            .map_err(|_| {
                TransactionError::InstructionError(0, InstructionError::InvalidInstructionData)
            })?;
            if decrypted.len() != 1 {
                return Err(TransactionError::InstructionError(
                    0,
                    InstructionError::InvalidInstructionData,
                ));
            }
            compact::AccountMeta::from_byte(decrypted[0]).ok_or(
                TransactionError::InstructionError(0, InstructionError::InvalidInstructionData),
            )?
        }
    };
    let pubkey = resolve_post_delegation_pubkey(pubkeys, account.key())?;

    Ok(if account.is_writable() {
        AccountMeta::new(pubkey, account.is_signer())
    } else {
        AccountMeta::new_readonly(pubkey, account.is_signer())
    })
}

fn decrypt_post_delegation_data(
    data: MaybeEncryptedIxData,
    validator_x25519_pubkey: &[u8; KEY_LEN],
    validator_x25519_secret: &[u8; KEY_LEN],
) -> Result<Vec<u8>, TransactionError> {
    let mut decrypted_data = data.prefix;
    if !data.suffix.as_bytes().is_empty() {
        decrypted_data.extend_from_slice(
            &encryption::decrypt(
                data.suffix.as_bytes(),
                validator_x25519_pubkey,
                validator_x25519_secret,
            )
            .map_err(|_| {
                TransactionError::InstructionError(0, InstructionError::InvalidInstructionData)
            })?,
        );
    }
    Ok(decrypted_data)
}

fn resolve_post_delegation_pubkey(
    pubkeys: &[Address],
    index: u8,
) -> Result<Address, TransactionError> {
    pubkeys
        .get(usize::from(index))
        .copied()
        .ok_or(TransactionError::InstructionError(
            0,
            InstructionError::InvalidInstructionData,
        ))
}