Add try_find_program_address syscall (bp #14118) (#14421)

* Add try_find_program_address syscall (#14118)

(cherry picked from commit ab98c1f2d4)

# Conflicts:
#	programs/bpf_loader/src/syscalls.rs
#	sdk/program/src/pubkey.rs
#	sdk/src/feature_set.rs

* fix conflicts

Co-authored-by: Jack May <jack@solana.com>

programs/bpf/c/src/invoke/invoke.c

@@ -158,6 +158,20 @@ extern uint64_t entrypoint(const uint8_t *input) {
       sol_assert(SolPubkey_same(&address, accounts[DERIVED_KEY1_INDEX].key));
     }
 
+    sol_log("Test try_find_program_address");
+    {
+      uint8_t seed[] = {'Y', 'o', 'u', ' ', 'p', 'a', 's', 's',
+                        ' ', 'b', 'u', 't', 't', 'e', 'r'};
+      const SolSignerSeed seeds[] = {{seed, SOL_ARRAY_SIZE(seed)}};
+      SolPubkey address;
+      uint8_t bump_seed;
+      sol_assert(SUCCESS == sol_try_find_program_address(
+                                seeds, SOL_ARRAY_SIZE(seeds), params.program_id,
+                                &address, &bump_seed));
+      sol_assert(SolPubkey_same(&address, accounts[DERIVED_KEY1_INDEX].key));
+      sol_assert(bump_seed == bump_seed1);
+    }
+
     sol_log("Test derived signers");
     {
       sol_assert(!accounts[DERIVED_KEY1_INDEX].is_signer);

programs/bpf/rust/invoke/src/lib.rs

@@ -252,6 +252,19 @@ fn process_instruction(
                 );
             }
 
+            msg!("Test try_find_program_address");
+            {
+                let (address, bump_seed) =
+                    Pubkey::try_find_program_address(&[b"You pass butter"], program_id).unwrap();
+                assert_eq!(&address, accounts[DERIVED_KEY1_INDEX].key);
+                assert_eq!(bump_seed, bump_seed1);
+                assert_eq!(
+                    Pubkey::create_program_address(&[b"You pass butter"], &Pubkey::default())
+                        .unwrap_err(),
+                    PubkeyError::InvalidSeeds
+                );
+            }
+
             msg!("Test derived signers");
             {
                 assert!(!accounts[DERIVED_KEY1_INDEX].is_signer);

programs/bpf_loader/src/syscalls.rs

@@ -18,6 +18,7 @@ use solana_sdk::{
     feature_set::{
         limit_cpi_loader_invoke, pubkey_log_syscall_enabled, ristretto_mul_syscall_enabled,
         sha256_syscall_enabled, sol_log_compute_units_syscall,
+        try_find_program_address_syscall_enabled,
     },
     hash::{Hasher, HASH_BYTES},
     instruction::{AccountMeta, Instruction, InstructionError},
@@ -181,6 +182,17 @@ pub fn register_syscalls<'a>(
         }),
     )?;
 
+    if invoke_context.is_feature_active(&try_find_program_address_syscall_enabled::id()) {
+        vm.register_syscall_with_context_ex(
+            "sol_try_find_program_address",
+            Box::new(SyscallTryFindProgramAddress {
+                cost: bpf_compute_budget.create_program_address_units,
+                compute_meter: invoke_context.get_compute_meter(),
+                loader_id,
+            }),
+        )?;
+    }
+
     // Cross-program invocation syscalls
 
     let invoke_context = Rc::new(RefCell::new(invoke_context));
@@ -519,6 +531,34 @@ impl SyscallObject<BPFError> for SyscallAllocFree {
     }
 }
 
+fn translate_program_address_inputs<'a>(
+    seeds_addr: u64,
+    seeds_len: u64,
+    program_id_addr: u64,
+    ro_regions: &[MemoryRegion],
+    loader_id: &Pubkey,
+) -> Result<(Vec<&'a [u8]>, &'a Pubkey), EbpfError<BPFError>> {
+    let untranslated_seeds =
+        translate_slice!(&[&u8], seeds_addr, seeds_len, ro_regions, loader_id)?;
+    if untranslated_seeds.len() > MAX_SEEDS {
+        return Err(SyscallError::BadSeeds(PubkeyError::MaxSeedLengthExceeded).into());
+    }
+    let seeds = untranslated_seeds
+        .iter()
+        .map(|untranslated_seed| {
+            translate_slice!(
+                u8,
+                untranslated_seed.as_ptr() as *const _ as u64,
+                untranslated_seed.len() as u64,
+                ro_regions,
+                loader_id
+            )
+        })
+        .collect::<Result<Vec<_>, EbpfError<BPFError>>>()?;
+    let program_id = translate_type!(Pubkey, program_id_addr, ro_regions, loader_id)?;
+    Ok((seeds, program_id))
+}
+
 /// Create a program address
 struct SyscallCreateProgramAddress<'a> {
     cost: u64,
@@ -537,26 +577,16 @@ impl<'a> SyscallObject<BPFError> for SyscallCreateProgramAddress<'a> {
         rw_regions: &[MemoryRegion],
     ) -> Result<u64, EbpfError<BPFError>> {
         self.compute_meter.consume(self.cost)?;
-        // TODO need ref?
-        let untranslated_seeds =
-            translate_slice!(&[&u8], seeds_addr, seeds_len, ro_regions, self.loader_id)?;
-        if untranslated_seeds.len() > MAX_SEEDS {
-            return Ok(1);
-        }
-        let seeds = untranslated_seeds
-            .iter()
-            .map(|untranslated_seed| {
-                translate_slice!(
-                    u8,
-                    untranslated_seed.as_ptr(),
-                    untranslated_seed.len(),
-                    ro_regions,
-                    self.loader_id
-                )
-            })
-            .collect::<Result<Vec<_>, EbpfError<BPFError>>>()?;
-        let program_id = translate_type!(Pubkey, program_id_addr, ro_regions, self.loader_id)?;
 
+        let (seeds, program_id) = translate_program_address_inputs(
+            seeds_addr,
+            seeds_len,
+            program_id_addr,
+            ro_regions,
+            self.loader_id,
+        )?;
+
+        self.compute_meter.consume(self.cost)?;
         let new_address = match Pubkey::create_program_address(&seeds, program_id) {
             Ok(address) => address,
             Err(_) => return Ok(1),
@@ -567,6 +597,56 @@ impl<'a> SyscallObject<BPFError> for SyscallCreateProgramAddress<'a> {
     }
 }
 
+/// Create a program address
+struct SyscallTryFindProgramAddress<'a> {
+    cost: u64,
+    compute_meter: Rc<RefCell<dyn ComputeMeter>>,
+    loader_id: &'a Pubkey,
+}
+impl<'a> SyscallObject<BPFError> for SyscallTryFindProgramAddress<'a> {
+    fn call(
+        &mut self,
+        seeds_addr: u64,
+        seeds_len: u64,
+        program_id_addr: u64,
+        address_addr: u64,
+        bump_seed_addr: u64,
+        ro_regions: &[MemoryRegion],
+        rw_regions: &[MemoryRegion],
+    ) -> Result<u64, EbpfError<BPFError>> {
+        let (seeds, program_id) = translate_program_address_inputs(
+            seeds_addr,
+            seeds_len,
+            program_id_addr,
+            ro_regions,
+            self.loader_id,
+        )?;
+
+        let mut bump_seed = [std::u8::MAX];
+        for _ in 0..std::u8::MAX {
+            {
+                let mut seeds_with_bump = seeds.to_vec();
+                seeds_with_bump.push(&bump_seed);
+
+                self.compute_meter.consume(self.cost)?;
+                if let Ok(new_address) =
+                    Pubkey::create_program_address(&seeds_with_bump, program_id)
+                {
+                    let bump_seed_ref =
+                        translate_type_mut!(u8, bump_seed_addr, rw_regions, self.loader_id)?;
+                    let address =
+                        translate_slice_mut!(u8, address_addr, 32, rw_regions, self.loader_id)?;
+                    *bump_seed_ref = bump_seed[0];
+                    address.copy_from_slice(new_address.as_ref());
+                    return Ok(0);
+                }
+            }
+            bump_seed[0] -= 1;
+        }
+        Ok(1)
+    }
+}
+
 /// SHA256
 pub struct SyscallSha256<'a> {
     sha256_base_cost: u64,

sdk/bpf/c/inc/solana_sdk.h

@@ -479,14 +479,31 @@ typedef struct {
  *
  * @param seeds Seed bytes used to sign program accounts
  * @param seeds_len Length of the seeds array
- * @param Progam id of the signer
+ * @param program_id Program id of the signer
- * @param Program address created, filled on return
+ * @param program_address Program address created, filled on return
  */
 static uint64_t sol_create_program_address(
     const SolSignerSeed *seeds,
     int seeds_len,
     const SolPubkey *program_id,
-    const SolPubkey *address
+    const SolPubkey *program_address
+);
+
+/**
+ * Try to find a program address and return corresponding bump seed
+ *
+ * @param seeds Seed bytes used to sign program accounts
+ * @param seeds_len Length of the seeds array
+ * @param program_id Program id of the signer
+ * @param program_address Program address created, filled on return
+ * @param bump_seed Bump seed required to create a valid program address
+ */
+static uint64_t sol_try_find_program_address(
+    const SolSignerSeed *seeds,
+    int seeds_len,
+    const SolPubkey *program_id,
+    const SolPubkey *program_address,
+    const uint8_t *bump_seed
 );
 
 /**

sdk/program/src/pubkey.rs

@@ -196,22 +196,66 @@ impl Pubkey {
         }
     }
 
-    /// Find a valid program address and its corresponding bump seed which must be passed
+    /// Find a valid program address and its corresponding bump seed which must
-    /// as an additional seed when calling `invoke_signed`
+    /// be passed as an additional seed when calling `invoke_signed`.
-    #[allow(clippy::same_item_push)]
+    ///
+    /// Panics in the very unlikely event that the additional seed could not be
+    /// found.
     pub fn find_program_address(seeds: &[&[u8]], program_id: &Pubkey) -> (Pubkey, u8) {
-        let mut bump_seed = [std::u8::MAX];
+        Self::try_find_program_address(seeds, program_id)
-        for _ in 0..std::u8::MAX {
+            .unwrap_or_else(|| panic!("Unable to find a viable program address bump seed"))
-            {
+    }
-                let mut seeds_with_bump = seeds.to_vec();
+
-                seeds_with_bump.push(&bump_seed);
+    /// Find a valid program address and its corresponding bump seed which must
-                if let Ok(address) = Self::create_program_address(&seeds_with_bump, program_id) {
+    /// be passed as an additional seed when calling `invoke_signed`
-                    return (address, bump_seed[0]);
+    #[allow(clippy::same_item_push)]
+    pub fn try_find_program_address(seeds: &[&[u8]], program_id: &Pubkey) -> Option<(Pubkey, u8)> {
+        // Perform the calculation inline, calling this from within a program is
+        // not supported
+        #[cfg(not(target_arch = "bpf"))]
+        {
+            let mut bump_seed = [std::u8::MAX];
+            for _ in 0..std::u8::MAX {
+                {
+                    let mut seeds_with_bump = seeds.to_vec();
+                    seeds_with_bump.push(&bump_seed);
+                    if let Ok(address) = Self::create_program_address(&seeds_with_bump, program_id)
+                    {
+                        return Some((address, bump_seed[0]));
+                    }
                 }
+                bump_seed[0] -= 1;
+            }
+            None
+        }
+        // Call via a system call to perform the calculation
+        #[cfg(target_arch = "bpf")]
+        {
+            extern "C" {
+                fn sol_try_find_program_address(
+                    seeds_addr: *const u8,
+                    seeds_len: u64,
+                    program_id_addr: *const u8,
+                    address_bytes_addr: *const u8,
+                    bump_seed_addr: *const u8,
+                ) -> u64;
+            };
+            let mut bytes = [0; 32];
+            let mut bump_seed = std::u8::MAX;
+            let result = unsafe {
+                sol_try_find_program_address(
+                    seeds as *const _ as *const u8,
+                    seeds.len() as u64,
+                    program_id as *const _ as *const u8,
+                    &mut bytes as *mut _ as *mut u8,
+                    &mut bump_seed as *mut _ as *mut u8,
+                )
+            };
+            match result {
+                crate::entrypoint::SUCCESS => Some((Pubkey::new(&bytes), bump_seed)),
+                _ => None,
             }
-            bump_seed[0] -= 1;
         }
-        panic!("Unable to find a viable program address bump seed");
     }
 
     pub fn to_bytes(self) -> [u8; 32] {

sdk/src/feature_set.rs

@@ -118,6 +118,10 @@ pub mod limit_cpi_loader_invoke {
     solana_sdk::declare_id!("xGbcW7EEC7zMRJ6LaJCob65EJxKryWjwM4rv8f57SRM");
 }
 
+pub mod try_find_program_address_syscall_enabled {
+    solana_sdk::declare_id!("EMsMNadQNhCYDyGpYH5Tx6dGHxiUqKHk782PU5XaWfmi");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -149,6 +153,7 @@ lazy_static! {
         (bpf_loader_upgradeable_program::id(), "upgradeable bpf loader"),
         (max_cpi_instruction_size_ipv6_mtu::id(), "Max cross-program invocation size 1280"),
         (limit_cpi_loader_invoke::id(), "Loader not authorized via CPI"),
+        (try_find_program_address_syscall_enabled::id(), "add try_find_program_address syscall"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()