Discriminator

src/processor/create_proof.rs

@@ -5,7 +5,13 @@ use solana_program::{
     program_error::ProgramError, pubkey::Pubkey, system_program,
 };
 
-use crate::{instruction::CreateProofArgs, loaders::*, state::Proof, utils::create_pda, PROOF};
+use crate::{
+    instruction::CreateProofArgs,
+    loaders::*,
+    state::{Discriminator, Proof},
+    utils::create_pda,
+    PROOF,
+};
 
 pub fn process_create_proof<'a, 'info>(
     _program_id: &Pubkey,
@@ -27,12 +33,13 @@ pub fn process_create_proof<'a, 'info>(
     create_pda(
         proof_info,
         &crate::id(),
-        size_of::<Proof>(),
+        8 + size_of::<Proof>(),
         &[PROOF, signer.key.as_ref(), &[args.bump]],
         system_program,
         signer,
     )?;
     let mut proof_data = proof_info.data.borrow_mut();
+    proof_data[0] = Proof::discriminator() as u8;
     let mut proof = Proof::try_from_bytes_mut(&mut proof_data)?;
     proof.bump = args.bump as u64;
     proof.authority = *signer.key;

src/processor/initialize.rs

@@ -8,6 +8,7 @@ use solana_program::{
 };
 use spl_token::state::Mint;
 
+use crate::state::Discriminator;
 use crate::{instruction::*, BUS, INITIAL_DIFFICULTY, MINT_ADDRESS, TREASURY_ADDRESS};
 use crate::{
     loaders::*,
@@ -70,31 +71,30 @@ pub fn process_initialize<'a, 'info>(
         create_pda(
             bus_infos[i],
             &crate::id(),
-            size_of::<Bus>(),
+            8 + size_of::<Bus>(),
             &[BUS, &[i as u8], &[bus_bumps[i]]],
             system_program,
             signer,
         )?;
-        bus_infos[i].try_borrow_mut_data()?.copy_from_slice(
-            Bus {
-                bump: bus_bumps[i] as u32,
-                id: i as u32,
-                available_rewards: 0,
-            }
-            .to_bytes(),
-        );
+        let mut bus_data = bus_infos[i].try_borrow_mut_data()?;
+        bus_data[0] = Bus::discriminator() as u8;
+        let mut bus = Bus::try_from_bytes_mut(&mut bus_data)?;
+        bus.bump = bus_bumps[i] as u32;
+        bus.id = i as u32;
+        bus.available_rewards = 0;
     }
 
     // Initialize treasury
     create_pda(
         treasury_info,
         &crate::id(),
-        size_of::<Treasury>(),
+        8 + size_of::<Treasury>(),
         &[TREASURY, &[args.treasury_bump]],
         system_program,
         signer,
     )?;
     let mut treasury_data = treasury_info.data.borrow_mut();
+    treasury_data[0] = Treasury::discriminator() as u8;
     let mut treasury = Treasury::try_from_bytes_mut(&mut treasury_data)?;
     treasury.bump = args.treasury_bump as u64;
     treasury.admin = *signer.key;

src/state/bus.rs

@@ -2,6 +2,8 @@ use bytemuck::{Pod, Zeroable};
 
 use crate::{impl_account_from_bytes, impl_to_bytes};
 
+use super::{AccountDiscriminator, Discriminator};
+
 #[repr(C)]
 #[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
 pub struct Bus {
@@ -15,5 +17,11 @@ pub struct Bus {
     pub available_rewards: u64,
 }
 
+impl Discriminator for Bus {
+    fn discriminator() -> super::AccountDiscriminator {
+        AccountDiscriminator::Bus
+    }
+}
+
 impl_to_bytes!(Bus);
 impl_account_from_bytes!(Bus);

src/state/hash.rs

@@ -3,7 +3,7 @@ use std::mem::transmute;
 use bytemuck::{Pod, Zeroable};
 use solana_program::keccak::{Hash as KeccakHash, HASH_BYTES};
 
-use crate::{impl_account_from_bytes, impl_to_bytes};
+use crate::impl_to_bytes;
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
@@ -24,4 +24,4 @@ impl From<Hash> for KeccakHash {
 }
 
 impl_to_bytes!(Hash);
-impl_account_from_bytes!(Hash);
+// impl_account_from_bytes!(Hash);

src/state/mod.rs

@@ -5,5 +5,18 @@ mod treasury;
 
 pub use bus::*;
 pub use hash::*;
+use num_enum::{IntoPrimitive, TryFromPrimitive};
 pub use proof::*;
 pub use treasury::*;
+
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq, IntoPrimitive, TryFromPrimitive)]
+pub enum AccountDiscriminator {
+    Bus = 100,
+    Proof = 101,
+    Treasury = 102,
+}
+
+pub trait Discriminator {
+    fn discriminator() -> AccountDiscriminator;
+}

src/state/proof.rs

@@ -3,7 +3,7 @@ use solana_program::pubkey::Pubkey;
 
 use crate::{impl_account_from_bytes, impl_to_bytes};
 
-use super::Hash;
+use super::{AccountDiscriminator, Discriminator, Hash};
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
@@ -27,5 +27,11 @@ pub struct Proof {
     pub total_rewards: u64,
 }
 
+impl Discriminator for Proof {
+    fn discriminator() -> super::AccountDiscriminator {
+        AccountDiscriminator::Proof
+    }
+}
+
 impl_to_bytes!(Proof);
 impl_account_from_bytes!(Proof);

src/state/treasury.rs

@@ -3,7 +3,7 @@ use solana_program::pubkey::Pubkey;
 
 use crate::{impl_account_from_bytes, impl_to_bytes};
 
-use super::Hash;
+use super::{AccountDiscriminator, Discriminator, Hash};
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
@@ -27,5 +27,11 @@ pub struct Treasury {
     pub total_claimed_rewards: u64,
 }
 
+impl Discriminator for Treasury {
+    fn discriminator() -> super::AccountDiscriminator {
+        AccountDiscriminator::Treasury
+    }
+}
+
 impl_to_bytes!(Treasury);
 impl_account_from_bytes!(Treasury);

src/utils.rs

@@ -50,14 +50,20 @@ macro_rules! impl_account_from_bytes {
             pub fn try_from_bytes(
                 data: &[u8],
             ) -> Result<&Self, solana_program::program_error::ProgramError> {
-                bytemuck::try_from_bytes::<Self>(data).or(Err(
+                if (Self::discriminator() as u8).ne(&data[0]) {
+                    return Err(solana_program::program_error::ProgramError::InvalidAccountData);
+                }
+                bytemuck::try_from_bytes::<Self>(&data[8..]).or(Err(
                     solana_program::program_error::ProgramError::InvalidAccountData,
                 ))
             }
             pub fn try_from_bytes_mut(
                 data: &mut [u8],
             ) -> Result<&mut Self, solana_program::program_error::ProgramError> {
-                bytemuck::try_from_bytes_mut::<Self>(data).or(Err(
+                if (Self::discriminator() as u8).ne(&data[0]) {
+                    return Err(solana_program::program_error::ProgramError::InvalidAccountData);
+                }
+                bytemuck::try_from_bytes_mut::<Self>(&mut data[8..]).or(Err(
                     solana_program::program_error::ProgramError::InvalidAccountData,
                 ))
             }