armaniferrante před 4 roky
rodič
revize
fa91b55639
3 změnil soubory, kde provedl 64 přidání a 7 odebrání
  1. 13 1
      cli/src/main.rs
  2. 18 6
      lang/src/fuzzing.rs
  3. 33 0
      lang/syn/src/codegen/program.rs

+ 13 - 1
cli/src/main.rs

@@ -2,7 +2,7 @@
 
 use crate::config::{read_all_programs, Config, Program};
 use anchor_lang::idl::{IdlAccount, IdlInstruction};
-use anchor_lang::{AccountDeserialize, AnchorDeserialize, AnchorSerialize};
+use anchor_lang::{AccountDeserialize, AnchorDeserialize, AnchorSerialize, Discriminator};
 use anchor_syn::idl::Idl;
 use anyhow::{anyhow, Context, Result};
 use clap::Clap;
@@ -572,6 +572,18 @@ fn fetch_idl(idl_addr: Pubkey) -> Result<Idl> {
             .map_or(Err(anyhow!("Account not found")), Ok)?;
     }
 
+    // Validate the account discriminator.
+    // TODO: use the anchor_client instead of manually checking this.
+    let discriminator = {
+        let mut disc = [0u8; 8];
+        disc.copy_from_slice(&account.data[..8]);
+        disc
+    };
+    if discriminator != IdlAccount::discriminator() {
+        println!("Error: the IDL's account discriminator is invalid.");
+        std::process::exit(1);
+    }
+
     // Cut off account discriminator.
     let mut d: &[u8] = &account.data[8..];
     let idl_account: IdlAccount = AnchorDeserialize::deserialize(&mut d)?;

+ 18 - 6
lang/src/fuzzing.rs

@@ -49,6 +49,10 @@ impl Environment {
         env
     }
 
+    pub fn accounts_mut(&mut self) -> &mut AccountStore {
+        &mut self.accounts
+    }
+
     // Registers the program on the environment so that it can be invoked via
     // CPI.
     pub fn register(&mut self, program: Box<dyn Program>) {
@@ -62,15 +66,23 @@ impl Environment {
         accounts: &[AccountInfo<'info>],
         seeds: &[&[&[u8]]],
     ) -> ProgramResult {
-        let current_program = self.current_program.unwrap();
-
         // If seeds were given, then calculate the expected PDA.
-        let pda = match seeds.len() > 0 {
-            false => None,
-            true => Some(Pubkey::create_program_address(seeds[0], &current_program).unwrap()),
+        let pda = {
+            match self.current_program {
+                None => None,
+                Some(current_program) => match seeds.len() > 0 {
+                    false => None,
+                    true => {
+                        Some(Pubkey::create_program_address(seeds[0], &current_program).unwrap())
+                    }
+                },
+            }
         };
 
+        // Set the current program.
         self.current_program = Some(ix.program_id);
+
+        // Invoke the current program.
         let program = self.programs.get(&ix.program_id).unwrap();
         let account_infos: Vec<AccountInfo> = ix
             .accounts
@@ -94,7 +106,7 @@ impl Environment {
     }
 }
 
-struct AccountStore {
+pub struct AccountStore {
     // Storage bytes.
     storage: Bump,
 }

+ 33 - 0
lang/syn/src/codegen/program.rs

@@ -39,6 +39,39 @@ pub fn generate(program: Program) -> proc_macro2::TokenStream {
         #methods
 
         #cpi
+
+        #[cfg(fuzzing)]
+        pub mod fuzzing {
+            use super::*;
+
+            #[derive(Debug)]
+            pub struct Program {
+                id: Pubkey,
+            }
+
+            impl Program {
+                pub fn new(id: Pubkey) -> Self {
+                    Self {
+                        id
+                    }
+                }
+            }
+
+            impl anchor_lang::fuzzing::Program for Program {
+                fn entry(
+                    &self,
+                    program_id: &Pubkey,
+                    accounts: &[AccountInfo],
+                    ix_data: &[u8],
+                ) -> ProgramResult {
+                    entry(program_id, accounts, ix_data)
+                }
+
+                fn id(&self) -> Pubkey {
+                    self.id
+                }
+            }
+        }
     }
 }