add condition filtering for txn ctx account getters

cali-jumptrading · cali-jumptrading · commit c8b03c64335e · 2025-03-21T22:00:47.000Z
diff --git a/src/flamenco/runtime/context/fd_exec_instr_ctx.c b/src/flamenco/runtime/context/fd_exec_instr_ctx.c
@@ -147,7 +147,8 @@ fd_exec_instr_ctx_try_borrow_last_program_account( fd_exec_instr_ctx_t const * c
   fd_txn_account_t * program_account = NULL;
   fd_exec_txn_ctx_get_account_at_index( ctx->txn_ctx, 
                                         ctx->instr->program_id,
-                                        &program_account );
+                                        &program_account, 
+                                        NULL );
 
   return fd_exec_instr_ctx_try_borrow_account( ctx, 
                                                ctx->instr->program_id,
diff --git a/src/flamenco/runtime/context/fd_exec_instr_ctx.h b/src/flamenco/runtime/context/fd_exec_instr_ctx.h
@@ -41,7 +41,7 @@ struct __attribute__((aligned(8UL))) fd_exec_instr_ctx {
 /* Be careful when using this macro. There may be places where the error
    will need to be handled differently. */
 #define FD_TRY_BORROW_INSTR_ACCOUNT_DEFAULT_ERR_CHECK( ctx, idx, acc ) do {    \
-  int err = fd_exec_instr_ctx_try_borrow_instr_account( ctx, idx, acc ); \
+  int err = fd_exec_instr_ctx_try_borrow_instr_account( ctx, idx, acc );       \
   if( FD_UNLIKELY( err ) ) return err;                                         \
 } while (0)
 
diff --git a/src/flamenco/runtime/context/fd_exec_txn_ctx.c b/src/flamenco/runtime/context/fd_exec_txn_ctx.c
@@ -86,39 +86,56 @@ fd_exec_txn_ctx_delete( void * mem ) {
 int
 fd_exec_txn_ctx_get_account_at_index( fd_exec_txn_ctx_t *  ctx,
                                       uchar                idx,
-                                      fd_txn_account_t * * account ) {
+                                      fd_txn_account_t * * account,
+                                      int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                        int                 idx,
+                                                        fd_txn_account_t *  acc) ) {
   if( FD_UNLIKELY( idx>=ctx->accounts_cnt ) ) {
     return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
   }
 
   fd_txn_account_t * txn_account = &ctx->accounts[idx];
   *account = txn_account;
 
+  if( condition != NULL ) {
+    if( FD_UNLIKELY( !condition( ctx, idx, *account ) ) ) 
+      return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
+  }
+
   return FD_ACC_MGR_SUCCESS;
 }
 
 int
 fd_exec_txn_ctx_get_account_with_key( fd_exec_txn_ctx_t *  ctx,
                                       fd_pubkey_t const *  pubkey,
-                                      fd_txn_account_t * * account ) {
+                                      fd_txn_account_t * * account,
+                                      int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                        int                 idx,
+                                                        fd_txn_account_t *  acc) ) {
   int index = fd_exec_txn_ctx_find_index_of_account( ctx, pubkey );
   if( FD_UNLIKELY( index==-1 ) ) {
     return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
   }
 
-  return fd_exec_txn_ctx_get_account_at_index( ctx, (uchar)index, account );
+  return fd_exec_txn_ctx_get_account_at_index( ctx, 
+                                               (uchar)index, 
+                                               account, 
+                                               condition );
 }
 
 int
 fd_exec_txn_ctx_get_executable_account( fd_exec_txn_ctx_t *  ctx,
                                         fd_pubkey_t const *  pubkey,
-                                        fd_txn_account_t * * account ) {
+                                        fd_txn_account_t * * account,
+                                        int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                          int                 idx,
+                                                          fd_txn_account_t *  acc) ) {
   /* First try to fetch the executable account from the existing borrowed accounts.
      If the pubkey is in the account keys, then we want to re-use that
      borrowed account since it reflects changes from prior instructions. Referencing the
      read-only executable accounts list is incorrect behavior when the program
      data account is written to in a prior instruction (e.g. program upgrade + invoke within the same txn) */
-  int err = fd_exec_txn_ctx_get_account_with_key( ctx, pubkey, account );
+  int err = fd_exec_txn_ctx_get_account_with_key( ctx, pubkey, account, condition );
   if( FD_UNLIKELY( err==FD_ACC_MGR_SUCCESS ) ) {
     return FD_ACC_MGR_SUCCESS;
   }
@@ -128,8 +145,10 @@ fd_exec_txn_ctx_get_executable_account( fd_exec_txn_ctx_t *  ctx,
       fd_txn_account_t * txn_account = &ctx->executable_accounts[i];
       *account = txn_account;
 
-      if( FD_UNLIKELY( !fd_acc_exists( txn_account->const_meta ) ) )
-        return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
+      if( condition != NULL ) {
+        if( FD_UNLIKELY( !condition( ctx, (int)i, *account ) ) ) 
+          return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
+      }
 
       return FD_ACC_MGR_SUCCESS;
     }
diff --git a/src/flamenco/runtime/context/fd_exec_txn_ctx.h b/src/flamenco/runtime/context/fd_exec_txn_ctx.h
@@ -240,7 +240,7 @@ fd_exec_txn_ctx_find_index_of_account( fd_exec_txn_ctx_t * ctx,
                                       fd_pubkey_t const *  pubkey ) {
   for( ulong i=ctx->accounts_cnt; i>0UL; i-- ) {
     if( 0==memcmp( pubkey, &ctx->account_keys[ i-1UL ], sizeof(fd_pubkey_t) ) ) {
-      return (int)((ushort)i);
+      return (int)(i-1UL);
     }
   }
   return -1;
@@ -253,17 +253,26 @@ fd_exec_txn_ctx_find_index_of_account( fd_exec_txn_ctx_t * ctx,
 int
 fd_exec_txn_ctx_get_account_at_index( fd_exec_txn_ctx_t *  ctx,
                                       uchar                idx,
-                                      fd_txn_account_t * * account );
+                                      fd_txn_account_t * * account,
+                                      int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                        int                 idx,
+                                                        fd_txn_account_t *  acc) );
 
 int
 fd_exec_txn_ctx_get_account_with_key( fd_exec_txn_ctx_t *  ctx,
                                       fd_pubkey_t const *  pubkey,
-                                      fd_txn_account_t * * account );
+                                      fd_txn_account_t * * account,
+                                      int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                        int                 idx,
+                                                        fd_txn_account_t *  acc) );
 
 int
 fd_exec_txn_ctx_get_executable_account( fd_exec_txn_ctx_t *  ctx,
                                         fd_pubkey_t const *  pubkey,
-                                        fd_txn_account_t * * account );
+                                        fd_txn_account_t * * account,
+                                        int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                          int                 idx,
+                                                          fd_txn_account_t *  acc) );
 
 void
 fd_exec_txn_ctx_reset_return_data( fd_exec_txn_ctx_t * txn_ctx );
@@ -283,6 +292,32 @@ fd_exec_txn_ctx_reset_return_data( fd_exec_txn_ctx_t * txn_ctx );
 int
 fd_exec_txn_ctx_account_is_writable_idx( fd_exec_txn_ctx_t const * txn_ctx, int idx );
 
+/* Conditional filter functions */
+static inline int
+fd_txn_account_exists( fd_exec_txn_ctx_t * ctx,
+                       int                 idx,
+                       fd_txn_account_t * acc ) {
+  (void) ctx;
+  (void) idx;
+  return fd_acc_exists( acc->const_meta );
+}
+
+static inline int
+fd_txn_account_is_writable( fd_exec_txn_ctx_t * ctx,
+                            int                 idx,
+                            fd_txn_account_t * acc ) {
+  (void) acc;
+  return fd_exec_txn_ctx_account_is_writable_idx( ctx, idx );
+}
+
+static inline int
+fd_txn_account_fee_payer_writable( fd_exec_txn_ctx_t * ctx,
+                                   int                 idx,
+                                   fd_txn_account_t * acc ) {
+  (void) acc;
+  return fd_txn_is_writable( ctx->txn_descriptor, idx );
+}
+
 FD_PROTOTYPES_END
 
 #endif /* HEADER_fd_src_flamenco_runtime_context_fd_exec_txn_ctx_h */
diff --git a/src/flamenco/runtime/fd_executor.c b/src/flamenco/runtime/fd_executor.c
@@ -444,7 +444,7 @@ fd_executor_load_transaction_accounts( fd_exec_txn_ctx_t * txn_ctx ) {
   /* https://github.com/anza-xyz/agave/blob/v2.2.0/svm/src/account_loader.rs#L429-L443 */
   for( ulong i=0UL; i<txn_ctx->accounts_cnt; i++ ) {
     fd_txn_account_t * acct = &txn_ctx->accounts[i];
-    uchar unknown_acc = !!( fd_exec_txn_ctx_get_account_at_index( txn_ctx, (uchar)i, &acct ) ||
+    uchar unknown_acc = !!( fd_exec_txn_ctx_get_account_at_index( txn_ctx, (uchar)i, &acct, &fd_txn_account_exists ) ||
                             acct->const_meta->info.lamports==0UL );
     ulong acc_size    = unknown_acc ? 0UL : acct->const_meta->dlen;
     uchar is_writable = !!( fd_exec_txn_ctx_account_is_writable_idx( txn_ctx, (int)i ) );
@@ -493,7 +493,7 @@ fd_executor_load_transaction_accounts( fd_exec_txn_ctx_t * txn_ctx ) {
     /* Mimicking `load_account()` here with 0-lamport check as well.
        https://github.com/anza-xyz/agave/blob/v2.2.0/svm/src/account_loader.rs#L455-L462 */
     fd_txn_account_t * program_account = NULL;
-    int err = fd_exec_txn_ctx_get_account_at_index( txn_ctx, instr->program_id, &program_account );
+    int err = fd_exec_txn_ctx_get_account_at_index( txn_ctx, instr->program_id, &program_account, &fd_txn_account_exists );
     if( FD_UNLIKELY( err!=FD_ACC_MGR_SUCCESS || program_account->const_meta->info.lamports==0UL ) ) {
       return FD_RUNTIME_TXN_ERR_PROGRAM_ACCOUNT_NOT_FOUND;
     }
@@ -792,8 +792,8 @@ fd_executor_validate_transaction_fee_payer( fd_exec_txn_ctx_t * txn_ctx ) {
    doesn't have a writable signature.
    https://github.com/anza-xyz/agave/blob/16de8b75ebcd57022409b422de557dd37b1de8db/svm/src/transaction_processor.rs#L431-L436 */
   fd_txn_account_t * fee_payer_rec = NULL;
-  err = fd_exec_txn_ctx_get_account_at_index( txn_ctx,  FD_FEE_PAYER_TXN_IDX, &fee_payer_rec );
-  if( FD_UNLIKELY( err!=FD_ACC_MGR_SUCCESS || !fd_txn_is_writable( txn_ctx->txn_descriptor, FD_FEE_PAYER_TXN_IDX ) ) ) {
+  err = fd_exec_txn_ctx_get_account_at_index( txn_ctx,  FD_FEE_PAYER_TXN_IDX, &fee_payer_rec, &fd_txn_account_fee_payer_writable );
+  if( FD_UNLIKELY( err!=FD_ACC_MGR_SUCCESS ) ) {
     return FD_RUNTIME_TXN_ERR_ACCOUNT_NOT_FOUND;
   }
 
diff --git a/src/flamenco/runtime/program/fd_bpf_loader_program.c b/src/flamenco/runtime/program/fd_bpf_loader_program.c
@@ -130,7 +130,7 @@ read_bpf_upgradeable_loader_state_for_program( fd_exec_txn_ctx_t *
                                                uchar                               program_id,
                                                int *                               opt_err ) {
   fd_txn_account_t * rec = NULL;
-  int err = fd_exec_txn_ctx_get_account_at_index( txn_ctx, program_id, &rec );
+  int err = fd_exec_txn_ctx_get_account_at_index( txn_ctx, program_id, &rec, &fd_txn_account_exists );
   if( FD_UNLIKELY( err ) ) {
     *opt_err = err;
     return NULL;
@@ -1923,7 +1923,7 @@ fd_bpf_loader_program_execute( fd_exec_instr_ctx_t * ctx ) {
 
       fd_txn_account_t * program_data_account = NULL;
       fd_pubkey_t *      programdata_pubkey   = (fd_pubkey_t *)&program_account_state->inner.program.programdata_address;
-      err = fd_exec_txn_ctx_get_executable_account( ctx->txn_ctx, programdata_pubkey, &program_data_account );
+      err = fd_exec_txn_ctx_get_executable_account( ctx->txn_ctx, programdata_pubkey, &program_data_account, &fd_txn_account_exists );
       if( FD_UNLIKELY( err!=FD_ACC_MGR_SUCCESS ) ) {
         return FD_EXECUTOR_INSTR_ERR_INVALID_ACC_DATA;
       }
diff --git a/src/flamenco/runtime/program/fd_loader_v4_program.c b/src/flamenco/runtime/program/fd_loader_v4_program.c
@@ -777,7 +777,7 @@ fd_loader_v4_program_execute( fd_exec_instr_ctx_t * instr_ctx ) {
       /* try borrow last program account */
       /* https://github.com/anza-xyz/agave/blob/v2.1.4/programs/loader-v4/src/lib.rs#L489 */
       fd_txn_account_t * program = NULL;
-      rc = fd_exec_txn_ctx_get_account_at_index( instr_ctx->txn_ctx, instr_ctx->instr->program_id, &program );
+      rc = fd_exec_txn_ctx_get_account_at_index( instr_ctx->txn_ctx, instr_ctx->instr->program_id, &program, &fd_txn_account_exists );
       if( FD_UNLIKELY( rc ) ) {
         return rc;
       }
diff --git a/src/flamenco/runtime/sysvar/fd_sysvar_instructions.c b/src/flamenco/runtime/sysvar/fd_sysvar_instructions.c
@@ -39,7 +39,7 @@ fd_sysvar_instructions_serialize_account( fd_exec_txn_ctx_t *      txn_ctx,
   ulong serialized_sz = instructions_serialized_size( instrs, instrs_cnt );
 
   fd_txn_account_t * rec = NULL;
-  int err = fd_exec_txn_ctx_get_account_with_key( txn_ctx, &fd_sysvar_instructions_id, &rec );
+  int err = fd_exec_txn_ctx_get_account_with_key( txn_ctx, &fd_sysvar_instructions_id, &rec, &fd_txn_account_exists );
   if( FD_UNLIKELY( err != FD_ACC_MGR_SUCCESS && rec == NULL ) ) {
     /* The way we use this, this should NEVER hit since the borrowed accounts should be set up
        before this is called, and this is only called if the sysvar instructions account is in
@@ -139,13 +139,9 @@ fd_sysvar_instructions_update_current_instr_idx( fd_exec_txn_ctx_t * txn_ctx,
   }
 
   /* https://github.com/anza-xyz/agave/blob/v2.1.14/svm/src/message_processor.rs#L53-L55 */
-  int err = fd_exec_txn_ctx_get_account_at_index( txn_ctx, (uchar)index, &rec );
-  if( FD_UNLIKELY( err != FD_ACC_MGR_SUCCESS ) )
+  int err = fd_exec_txn_ctx_get_account_at_index( txn_ctx, (uchar)index, &rec, NULL );
+  if( FD_UNLIKELY( err != FD_ACC_MGR_SUCCESS ) ) {
     return FD_ACC_MGR_ERR_READ_FAILED;
-
-    /* Check if account is writable */
-  if( FD_UNLIKELY( !fd_exec_txn_ctx_account_is_writable_idx( txn_ctx, index ) ) ) {
-    return FD_ACC_MGR_ERR_WRITE_FAILED;
   }
 
   /* Store the current instruction index
diff --git a/src/flamenco/runtime/tests/fd_exec_instr_test.c b/src/flamenco/runtime/tests/fd_exec_instr_test.c
@@ -2290,7 +2290,7 @@ __wrap_fd_execute_instr( fd_exec_txn_ctx_t * txn_ctx,
            TODO: Once direct mapping is enabled we _technically_ don't need
                  this check */
 
-        if( fd_exec_txn_ctx_get_account_at_index( txn_ctx, idx_in_txn, &acct ) ) {
+        if( fd_exec_txn_ctx_get_account_at_index( txn_ctx, idx_in_txn, &acct, &fd_txn_account_exists ) ) {
           break;
         }
         if( acct->meta == NULL ){
@@ -2301,7 +2301,8 @@ __wrap_fd_execute_instr( fd_exec_txn_ctx_t * txn_ctx,
         int err = fd_exec_txn_ctx_get_account_at_index( txn_ctx,
                                                           idx_in_txn,
                                                           /* Do not reallocate if data is not going to be modified */
-                                                          &acct );
+                                                          &acct,
+                                                        &fd_txn_account_is_writable );
         if( err ) break;
 
         /* resize manually */

Original file line number	Diff line number	Diff line change
`@@ -777,7 +777,7 @@ fd_loader_v4_program_execute( fd_exec_instr_ctx_t * instr_ctx ) {`
`777`	`777`	`/* try borrow last program account */`
`778`	`778`	`/* https://github.com/anza-xyz/agave/blob/v2.1.4/programs/loader-v4/src/lib.rs#L489 */`
`779`	`779`	`fd_txn_account_t * program = NULL;`
`780`		`- rc = fd_exec_txn_ctx_get_account_at_index( instr_ctx->txn_ctx, instr_ctx->instr->program_id, &program );`
	`780`	`+ rc = fd_exec_txn_ctx_get_account_at_index( instr_ctx->txn_ctx, instr_ctx->instr->program_id, &program, &fd_txn_account_exists );`
`781`	`781`	`if( FD_UNLIKELY( rc ) ) {`
`782`	`782`	`return rc;`
`783`	`783`	`}`