rialo-venus 0.2.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::cmp::Ordering;

use rialo_s_program::{
    account_info::AccountInfo,
    entrypoint::ProgramResult,
    msg,
    program::{invoke, invoke_signed},
    program_error::ProgramError,
    pubkey::Pubkey,
    rent::Rent,
    system_instruction, system_program,
    sysvar::Sysvar,
};
use rialo_types::Nonce;

// Re-export dependencies so that programs using rialo-venus don't have to declare them again.
pub mod reexports {
    pub use base64;
    pub use bincode;
    pub use rialo_events_core;
    pub use rialo_oracle_processor_interface;
    pub use rialo_oracle_registry_interface;
    pub use rialo_s_program;
    pub use rialo_subscriber_interface;
    pub use rialo_types;
    pub use serde;
}

/// Seed used to generate the PDA containing the workflow data.
/// The seed for the PDA is `RIALO_WORKFLOW_SEED` + payer key + nonce.
pub const WORKFLOW_SEED: &str = "rialo_workflow";

/// Derive the Program Derived Address (PDA) for a workflow account.
///
/// # Arguments
///
/// * `program_id` - The program ID that will own the workflow account
/// * `payer_key` - The payer key
/// * `nonce` - A nonce to discriminate between different workflows for this program
///
/// # Returns
///
/// A tuple containing the PDA and bump seed
///
/// # Example
///
/// ```rust
/// use rialo_s_program::pubkey::Pubkey;
/// use rialo_venus::derive_workflow_address;
///
/// let payer_key = Pubkey::new_unique();
/// let program_id = Pubkey::new_unique();
/// let nonce = b"My Nonce".to_vec();
/// let (workflow_address, bump) = derive_workflow_address(&program_id, &payer_key, &nonce);
/// ```
pub fn derive_workflow_address<NONCE: Into<Nonce>>(
    program_id: &Pubkey,
    payer_key: &Pubkey,
    nonce: NONCE,
) -> (Pubkey, u8) {
    Pubkey::find_program_address(
        &[
            WORKFLOW_SEED.as_bytes(),
            payer_key.as_array(),
            nonce.into().as_bytes(),
        ],
        program_id,
    )
}

/// Write workflow data to storage.
///
/// This method writes the workflow data to the PDA account, creating the account if required.
/// If the workflow data has grown larger than the allocated account
/// space, the account will be automatically resized.
///
/// # Arguments
///
/// * `program_id` - The program ID that owns the workflow account
/// * `slug` - A slug to discriminate between different workflows for this payer
/// * `payer_account` - must be a signer, used for potential resize funding
/// * `workflow_account` - PDA from `derive_workflow_address`, doesn't have to be initialised
/// * `system_program_account` - `rialo_s_program::system_program`, used for potential resize
/// * `data` - The data item to serialise and write
///
/// # Returns
///
/// * `ProgramResult` - Result of the operation
///
/// # Errors
///
/// Returns an error if:
/// * Required accounts are missing
/// * PDA derivation doesn't match
/// * Serialization fails
/// * Account data is too small and resize fails
pub fn write_to_storage<'account_info, D: serde::Serialize, NONCE: Into<Nonce>>(
    program_id: &Pubkey,
    slug: NONCE,
    payer_account: &AccountInfo<'account_info>,
    workflow_account: &AccountInfo<'account_info>,
    system_program_account: &AccountInfo<'account_info>,
    data: &D,
) -> ProgramResult {
    // Ensure the system program account is the correct one
    if !system_program::check_id(system_program_account.key) {
        msg!("Error: Incorrect system program ID");
        return Err(ProgramError::IncorrectProgramId);
    }

    let slug: Nonce = slug.into();
    let slug_bytes = slug.as_bytes();

    // Derive the PDA for the workflow account
    let (workflow_pubkey, bump_seed) =
        derive_workflow_address(program_id, payer_account.key, slug_bytes);

    // Verify the derived address matches the provided workflow account
    if workflow_pubkey != *workflow_account.key {
        msg!(
            "Error: Provided workflow account key {} doesn't match derived address {}",
            workflow_account.key,
            workflow_pubkey
        );
        return Err(ProgramError::InvalidAccountData);
    }

    // TODO: Better error code
    let data_bytes = bincode::serialize(data).map_err(|_| ProgramError::InvalidAccountData)?;
    let data_len = data_bytes.len();

    let rent = Rent::get()?;
    let minimum_balance_required = rent.minimum_balance(data_len);

    // Check if the workflow data account is initialized
    if workflow_account.kelvins() == 0 {
        // Create the workflow account
        let seeds = &[
            WORKFLOW_SEED.as_bytes(),
            payer_account.key.as_array(),
            slug_bytes,
            &[bump_seed],
        ];

        msg!(
            "Creating workflow account {} with {} bytes and {} kelvin",
            workflow_account.key,
            data_len,
            minimum_balance_required,
        );

        // Create account with PDA
        invoke_signed(
            &system_instruction::create_account(
                payer_account.key,
                workflow_account.key,
                minimum_balance_required,
                data_len as u64,
                program_id,
            ),
            &[payer_account.clone(), workflow_account.clone()],
            &[seeds],
        )?;

        msg!(
            "Workflow account {} initialized successfully",
            workflow_account.key
        );
    } else {
        // Check if resizing is needed via rent cost
        if workflow_account.try_borrow_data()?.len() != data_len {
            let original_data_len = workflow_account.data_len();

            msg!(
                "Resizing workflow account {} need {} bytes != allocated {}",
                workflow_account.key,
                data_len,
                original_data_len
            );

            // Calculate kelvin needed for rent exemption after resize
            let current_balance = workflow_account.kelvins();

            // Transfer kelvin if needed in either direction
            match minimum_balance_required.cmp(&current_balance) {
                Ordering::Greater => {
                    msg!(
                        "Workflow account {} received {} kelvin because it has grown",
                        workflow_account.key,
                        minimum_balance_required - current_balance,
                    );
                    invoke(
                        &system_instruction::transfer(
                            payer_account.key,
                            workflow_account.key,
                            minimum_balance_required - current_balance,
                        ),
                        &[
                            payer_account.clone(),
                            workflow_account.clone(),
                            system_program_account.clone(),
                        ],
                    )?;
                }

                Ordering::Less => {
                    // I wonder if this is even worthwhile in terms of compute cost vs. rent minimisation?
                    msg!(
                        "Workflow account {} returned {} kelvin because it has shrunk",
                        workflow_account.key,
                        current_balance - minimum_balance_required,
                    );

                    let from_account = workflow_account;
                    let to_account = payer_account;
                    let kelvin = current_balance - minimum_balance_required;
                    if **from_account.try_borrow_mut_kelvins()? < kelvin {
                        return Err(ProgramError::InsufficientFunds);
                    }
                    **from_account.try_borrow_mut_kelvins()? -= kelvin;
                    **to_account.try_borrow_mut_kelvins()? += kelvin;

                    // We cannot use `invoke_signed` because: "Transfer: `from` must not carry data"
                }

                Ordering::Equal => {}
            }

            /*
               TODO: realloc is actually limited to MAX_PERMITTED_DATA_INCREASE in one transaction,
               handle cases when this value is actually exceeded
            */
            workflow_account.realloc(data_len, true)?;

            msg!(
                "Resized workflow account {} from {} to {} bytes",
                workflow_account.key,
                original_data_len,
                data_len,
            );
        }
    }

    // Update the account data with the new workflow information
    workflow_account
        .try_borrow_mut_data()?
        .copy_from_slice(&data_bytes);

    Ok(())
}

/// Read workflow data from storage.
///
/// This method reads the workflow data from the PDA account.
/// It exists to hide the serialisation strategy.
///
/// # Arguments
///
/// * `data` - bytes to deserialize
///
/// # Returns
///
/// * `Result` - Result of the operation
///
/// # Errors
///
/// Returns an error if:
/// * Deserialization fails
pub fn read_from_storage<D: for<'de> serde::Deserialize<'de>>(
    data: &[u8],
) -> Result<D, ProgramError> {
    // TODO: Better error code
    bincode::deserialize::<D>(data).map_err(|_| ProgramError::InvalidAccountData)
}

/// Close the workflow account and transfer its kelvin to the payer account.
///
/// # Arguments
///
/// * `payer_account` - The account that will receive the kelvin from the workflow account
/// * `workflow_account` - The PDA account to be closed
///
/// # Returns
///
/// * `ProgramResult` - Result of the operation
/// # Errors
///
/// Returns an error if:
/// * The workflow account data cannot be borrowed for mutation
/// * The workflow account kelvin cannot be borrowed for mutation
/// * The payer account kelvin cannot be borrowed for mutation
pub fn close_account(
    payer_account: &AccountInfo<'_>,
    workflow_account: &AccountInfo<'_>,
) -> ProgramResult {
    workflow_account.try_borrow_mut_data()?.fill(0);
    let workflow_kelvin = workflow_account.kelvins();
    **workflow_account.try_borrow_mut_kelvins()? = 0;
    **payer_account.try_borrow_mut_kelvins()? += workflow_kelvin;
    Ok(())
}