Securing Solana record program using types.

entrypoint.rs

//! Program entrypoint

use solana_program::{
    account_info::AccountInfo, entrypoint, entrypoint::ProgramResult, pubkey::Pubkey,
};

entrypoint!(process_instruction);
fn process_instruction<'a>(
    program_id: &Pubkey,
    accounts: &'a [AccountInfo<'a>],
    instruction_data: &[u8],
) -> ProgramResult {
    crate::processor::process_instruction(*program_id, accounts, instruction_data)
}

processor.rs

//! Program state processor

use {
    crate::{error::RecordError, instruction::RecordInstruction, state::RecordData},
    borsh::{BorshDeserialize, BorshSerialize},
    solana_program::{
        account_info::{next_account_info, AccountInfo},
        entrypoint::ProgramResult,
        msg,
        program_error::ProgramError,
        program_pack::IsInitialized,
        pubkey::Pubkey,
    },
};

mod api {
    use super::*;

    pub mod authority_info {
        use super::*;
        #[must_use]
        pub struct Type<'a> {
            pub info: &'a AccountInfo<'a>,
            _secret: (),
        }
        
        pub fn from<'a>(
            iter: &mut std::slice::Iter<'a, AccountInfo<'a>>,
        ) -> Result<Type<'a>, ProgramError> {
            let info = next_account_info(iter)?;
            let obj = Type { info, _secret: () };
            Ok(obj)
        }
    }

    pub mod record_info {
        use super::*;
        #[must_use]
        pub struct Type<'a> {
            pub info: &'a AccountInfo<'a>,
            _secret: (),
        }
        pub fn from<'a>(
            iter: &mut std::slice::Iter<'a, AccountInfo<'a>>,
            program_id: Pubkey,
        ) -> Result<Type<'a>, ProgramError> {
            let info = next_account_info(iter)?;
            if info.owner != &program_id {
                msg!("Record account not owned by program");
                return Err(ProgramError::IllegalOwner);
            }
            let obj = Type { info, _secret: () };
            Ok(obj)
        }
    }

    pub mod record_deserialized {
        use super::*;
        #[must_use]
        pub struct Type<'a> {
            pub bytes: std::cell::RefMut<'a, &'a mut [u8]>,
            pub state: RecordData,
            _secret: (),
        }
        pub fn from<'a>(record: record_info::Type<'a>) -> Result<Type<'a>, ProgramError> {
            let bytes = record
                .info
                .data
                .try_borrow_mut()
                .ok()
                .ok_or(ProgramError::AccountBorrowFailed)?;
            let state = RecordData::try_from_slice(&*bytes)?;
            let obj = Type {
                bytes,
                state,
                _secret: (),
            };
            Ok(obj)
        }
    }
    fn serialize<'a>(record: &mut record_deserialized::Type<'a>) -> ProgramResult {
        let writer = &mut *record.bytes;
        record.state.serialize(writer).map_err(|e| e.into())
    }

    pub mod record_uninitialized {
        use super::*;
        #[must_use]
        pub struct Type<'a> {
            data: record_deserialized::Type<'a>,
        }
        pub fn from<'a>(record: record_deserialized::Type<'a>) -> Result<Type<'a>, ProgramError> {
            if record.state.is_initialized() {
                msg!("Record account already initialized");
                return Err(ProgramError::AccountAlreadyInitialized);
            }
            let obj = Type { data: record };
            Ok(obj)
        }
        impl<'a> Type<'a> {
            pub fn initialize(mut self, authority: authority_info::Type<'a>) -> ProgramResult {
                self.data.state.version = RecordData::CURRENT_VERSION;
                self.data.state.authority = *authority.info.key;
                serialize(&mut self.data)
            }
        }
    }

    pub mod record_authorized {
        use super::*;
        #[must_use]
        pub struct Type<'a> {
            pub data: record_deserialized::Type<'a>,
            _secret: (),
        }
        pub fn from<'a>(
            record: record_deserialized::Type<'a>,
            authority: authority_info::Type<'a>,
        ) -> Result<Type<'a>, ProgramError> {
            if !record.state.is_initialized() {
                msg!("Record account not initialized");
                return Err(ProgramError::UninitializedAccount);
            }
            if record.state.authority != *authority.info.key {
                msg!("Incorrect record authority provided");
                return Err(RecordError::IncorrectAuthority.into());
            }
            if !authority.info.is_signer {
                msg!("Record authority signature missing");
                return Err(ProgramError::MissingRequiredSignature);
            }

            let obj = Type {
                data: record,
                _secret: (),
            };
            Ok(obj)
        }
        impl<'a> Type<'a> {
            pub fn set_authority(
                mut self,
                iter: &mut std::slice::Iter<'a, AccountInfo<'a>>,
            ) -> ProgramResult {
                let new_auth = next_account_info(iter)?;
                self.data.state.authority = *new_auth.key;
                serialize(&mut self.data)
            }
            pub fn write(mut self, offset: u64, new_data: Vec<u8>) -> ProgramResult {
                let start = RecordData::WRITABLE_START_INDEX
                    .checked_add(offset as usize)
                    .ok_or(RecordError::Overflow)?;

                let end = start
                    .checked_add(new_data.len())
                    .ok_or(RecordError::Overflow)?;

                let slice = self
                    .data
                    .bytes
                    .get_mut(start..end)
                    .ok_or(ProgramError::AccountDataTooSmall)?;

                slice.copy_from_slice(&new_data);
                Ok(())
            }
        }
    }

    pub mod record_ready_to_close {
        use super::*;
        #[must_use]
        pub struct Type<'a> {
            bytes: std::cell::RefMut<'a, &'a mut [u8]>,
            lamports: std::cell::RefMut<'a, &'a mut u64>,
        }
        fn from_inner<'a>(
            record: record_authorized::Type<'a>,
            lamports: std::cell::RefMut<'a, &'a mut u64>,
        ) -> Result<Type<'a>, ProgramError> {
            let obj = Type {
                bytes: record.data.bytes,
                lamports,
            };
            Ok(obj)
        }
        pub fn from<'a>(
            record: record_info::Type<'a>,
            authority: authority_info::Type<'a>,
        ) -> Result<Type<'a>, ProgramError> {
            let lamports = record
                .info
                .lamports
                .try_borrow_mut()
                .ok()
                .ok_or(ProgramError::AccountBorrowFailed)?;
            let record = record_deserialized::from(record)?;
            let record = record_authorized::from(record, authority)?;
            from_inner(record, lamports)
        }
        impl<'a> Type<'a> {
            pub fn close(
                mut self,
                iter: &mut std::slice::Iter<'a, AccountInfo<'a>>,
            ) -> ProgramResult {
                let destination = next_account_info(iter)?;
                let mut lamports = destination
                    .lamports
                    .try_borrow_mut()
                    .ok()
                    .ok_or(ProgramError::AccountBorrowFailed)?;
                let new_lamports = lamports
                    .checked_add(**self.lamports)
                    .ok_or(RecordError::Overflow)?;

                **lamports = new_lamports;
                *self.bytes = &mut [];
                **self.lamports = 0;
                Ok(())
            }
        }
    }
}

/// Instruction processor
pub fn process_instruction<'a>(
    program_id: Pubkey,
    accounts: &'a [AccountInfo<'a>],
    input: &[u8],
) -> ProgramResult {
    let instruction = RecordInstruction::try_from_slice(input)?;
    let account_info_iter = &mut accounts.iter();

    match instruction {
        RecordInstruction::Initialize => {
            msg!("RecordInstruction::Initialize");

            let record = api::record_info::from(account_info_iter, program_id)?;
            let authority = api::authority_info::from(account_info_iter)?;

            let record = api::record_deserialized::from(record)?;
            let record = api::record_uninitialized::from(record)?;
            record.initialize(authority)
        }

        RecordInstruction::Write { offset, data } => {
            msg!("RecordInstruction::Write");

            let record = api::record_info::from(account_info_iter, program_id)?;
            let authority = api::authority_info::from(account_info_iter)?;

            let record = api::record_deserialized::from(record)?;
            let record = api::record_authorized::from(record, authority)?;
            record.write(offset, data)
        }

        RecordInstruction::SetAuthority => {
            msg!("RecordInstruction::SetAuthority");

            let record = api::record_info::from(account_info_iter, program_id)?;
            let authority = api::authority_info::from(account_info_iter)?;

            let record = api::record_deserialized::from(record)?;
            let record = api::record_authorized::from(record, authority)?;

            record.set_authority(account_info_iter)
        }

        RecordInstruction::CloseAccount => {
            msg!("RecordInstruction::CloseAccount");

            let record = api::record_info::from(account_info_iter, program_id)?;
            let authority = api::authority_info::from(account_info_iter)?;

            let record = api::record_ready_to_close::from(record, authority)?;
            record.close(account_info_iter)
        }
    }
}