Line data    Source code

       1             : #include "fd_zksdk_private.h"
       2             : #include "../../fd_account.h"
       3             : #include "../../fd_system_ids.h"
       4             : 
       5             : /* fd_zksdk_process_close_context_state is equivalent to process_close_proof_context()
       6             :    https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L127 */
       7             : int
       8          36 : fd_zksdk_process_close_context_state( fd_exec_instr_ctx_t * ctx ) {
       9          57 : #define ACC_IDX_PROOF (0UL)
      10          36 : #define ACC_IDX_DEST  (1UL)
      11          36 : #define ACC_IDX_OWNER (2UL)
      12             : 
      13          36 :   fd_pubkey_t owner_pubkey[1];
      14          36 :   fd_pubkey_t proof_pubkey[1];
      15          36 :   fd_pubkey_t dest_pubkey[1];
      16             : 
      17             :   /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L131-L139 */
      18          36 :   FD_BORROWED_ACCOUNT_DECL( owner_acc );
      19         144 :   FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, ACC_IDX_OWNER, owner_acc ) {
      20          36 :     if( FD_UNLIKELY( !fd_instr_acc_is_signer_idx( ctx->instr, ACC_IDX_OWNER ) ) ) {
      21           3 :       return FD_EXECUTOR_INSTR_ERR_MISSING_REQUIRED_SIGNATURE;
      22           3 :     }
      23          33 :     fd_memcpy( owner_pubkey, owner_acc->pubkey, sizeof(fd_pubkey_t) );
      24          36 :   } FD_BORROWED_ACCOUNT_DROP( owner_acc );
      25             : 
      26             :   /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L141-L149
      27             :      Note: following exactly Agave's behavior, we can certainly simplify the borrowings. */
      28          33 :   FD_BORROWED_ACCOUNT_DECL( proof_acc );
      29         132 :   FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, ACC_IDX_PROOF, proof_acc ) {
      30          33 :     fd_memcpy( proof_pubkey, proof_acc->pubkey, sizeof(fd_pubkey_t) );
      31          33 :   } FD_BORROWED_ACCOUNT_DROP( proof_acc );
      32             : 
      33          33 :   FD_BORROWED_ACCOUNT_DECL( dest_acc );
      34         132 :   FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, ACC_IDX_DEST, dest_acc ) {
      35          33 :     fd_memcpy( dest_pubkey, dest_acc->pubkey, sizeof(fd_pubkey_t) );
      36          33 :   } FD_BORROWED_ACCOUNT_DROP( dest_acc );
      37             : 
      38          33 :   if( FD_UNLIKELY( fd_memeq( proof_pubkey, dest_pubkey, sizeof(fd_pubkey_t) ) ) ) {
      39           3 :     return FD_EXECUTOR_INSTR_ERR_INVALID_INSTR_DATA;
      40           3 :   }
      41             : 
      42             :   /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L151-L152 */
      43         120 :   FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, ACC_IDX_PROOF, proof_acc ) {
      44             : 
      45             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L161-L162
      46             :        Note: data also contains context data, but we only need the initial 33 bytes. */
      47          30 :     if( FD_UNLIKELY( proof_acc->const_meta->dlen < sizeof(fd_zksdk_proof_ctx_state_meta_t) ) ) {
      48           3 :       return FD_EXECUTOR_INSTR_ERR_INVALID_ACC_DATA;
      49           3 :     }
      50          27 :     fd_zksdk_proof_ctx_state_meta_t const * proof_ctx_state_meta = fd_type_pun_const( proof_acc->const_data );
      51             : 
      52             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L155 */
      53          27 :     fd_pubkey_t const * expected_owner_addr = &proof_ctx_state_meta->ctx_state_authority;
      54             : 
      55             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L157-L159 */
      56          27 :     if( FD_UNLIKELY( !fd_memeq( owner_pubkey, expected_owner_addr, sizeof(fd_pubkey_t) ) ) ) {
      57           3 :       return FD_EXECUTOR_INSTR_ERR_INVALID_ACC_OWNER;
      58           3 :     }
      59             : 
      60             :   /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L161-L162 */
      61          96 :   FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, ACC_IDX_DEST, dest_acc ) {
      62             : 
      63             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L163-L166 */
      64          24 :     int err = 0;
      65          24 :     err = fd_account_checked_add_lamports( ctx, ACC_IDX_DEST, proof_acc->const_meta->info.lamports );
      66          24 :     if( FD_UNLIKELY( err ) ) {
      67           3 :       return err;
      68           3 :     }
      69          21 :     err = fd_account_set_lamports( ctx, ACC_IDX_PROOF, 0UL );
      70          21 :     if( FD_UNLIKELY( err ) ) {
      71           3 :       return err;
      72           3 :     }
      73          18 :     err = fd_account_set_data_length( ctx, ACC_IDX_PROOF, 0UL );
      74          18 :     if( FD_UNLIKELY( err ) ) {
      75           0 :       return err;
      76           0 :     }
      77          18 :     err = fd_account_set_owner( ctx, ACC_IDX_PROOF, &fd_solana_system_program_id );
      78          18 :     if( FD_UNLIKELY( err ) ) {
      79           0 :       return err;
      80           0 :     }
      81             : 
      82          24 :   } FD_BORROWED_ACCOUNT_DROP( dest_acc );
      83          30 :   } FD_BORROWED_ACCOUNT_DROP( proof_acc );
      84             : 
      85          18 :   return FD_EXECUTOR_INSTR_SUCCESS;
      86          30 : }
      87             : 
      88             : /* fd_zksdk_process_verify_proof is equivalent to process_verify_proof()
      89             :    and calls specific functions inside instructions/ to verify each
      90             :    individual ZKP.
      91             :    https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L32 */
      92             : int
      93         585 : fd_zksdk_process_verify_proof( fd_exec_instr_ctx_t * ctx ) {
      94         585 :   uchar const * instr_data = ctx->instr->data;
      95         585 :   ulong instr_acc_cnt      = ctx->instr->acct_cnt;
      96         585 :   uchar instr_id = instr_data[0]; /* instr_data_sz already checked by the caller */
      97             : 
      98             :   /* ProofContextState "header" size, ie. 1 authority pubkey + 1 proof_type byte */
      99         624 : #define CTX_HEAD_SZ 33UL
     100             : 
     101             :   /* Aux memory buffer.
     102             :      When proof data is taken from ix data we can access it directly,
     103             :      but when it's taken from account data we need to copy it to release
     104             :      the borrow. The largest ZKP is for range_proof_u256.
     105             :      Moreover, when storing context to an account, we need to serialize
     106             :      the ProofContextState struct that has 33 bytes of header -- we include
     107             :      them here so we can do a single memcpy. */
     108         585 : #define MAX_SZ (sizeof(fd_zksdk_range_proof_u256_proof_t)+sizeof(fd_zksdk_batched_range_proof_context_t))
     109         585 :   uchar buffer[ CTX_HEAD_SZ+MAX_SZ ];
     110             : 
     111             :   /* Specific instruction function */
     112         585 :   int (*fd_zksdk_instr_verify_proof)( void const *, void const * ) = NULL;
     113         585 :   switch( instr_id ) {
     114         126 :   case FD_ZKSDK_INSTR_VERIFY_ZERO_CIPHERTEXT:
     115         126 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_zero_ciphertext;
     116         126 :     break;
     117          63 :   case FD_ZKSDK_INSTR_VERIFY_CIPHERTEXT_CIPHERTEXT_EQUALITY:
     118          63 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_ciphertext_ciphertext_equality;
     119          63 :     break;
     120          39 :   case FD_ZKSDK_INSTR_VERIFY_CIPHERTEXT_COMMITMENT_EQUALITY:
     121          39 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_ciphertext_commitment_equality;
     122          39 :     break;
     123          33 :   case FD_ZKSDK_INSTR_VERIFY_PUBKEY_VALIDITY:
     124          33 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_pubkey_validity;
     125          33 :     break;
     126          45 :   case FD_ZKSDK_INSTR_VERIFY_PERCENTAGE_WITH_CAP:
     127          45 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_percentage_with_cap;
     128          45 :     break;
     129          48 :   case FD_ZKSDK_INSTR_VERIFY_BATCHED_RANGE_PROOF_U64:
     130          48 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_batched_range_proof_u64;
     131          48 :     break;
     132          51 :   case FD_ZKSDK_INSTR_VERIFY_BATCHED_RANGE_PROOF_U128:
     133          51 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_batched_range_proof_u128;
     134          51 :     break;
     135          42 :   case FD_ZKSDK_INSTR_VERIFY_BATCHED_RANGE_PROOF_U256:
     136          42 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_batched_range_proof_u256;
     137          42 :     break;
     138          36 :   case FD_ZKSDK_INSTR_VERIFY_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY:
     139          36 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_grouped_ciphertext_2_handles_validity;
     140          36 :     break;
     141          39 :   case FD_ZKSDK_INSTR_VERIFY_BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY:
     142          39 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_batched_grouped_ciphertext_2_handles_validity;
     143          39 :     break;
     144          30 :   case FD_ZKSDK_INSTR_VERIFY_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY:
     145          30 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_grouped_ciphertext_3_handles_validity;
     146          30 :     break;
     147          33 :   case FD_ZKSDK_INSTR_VERIFY_BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY:
     148          33 :     fd_zksdk_instr_verify_proof = &fd_zksdk_instr_verify_proof_batched_grouped_ciphertext_3_handles_validity;
     149          33 :     break;
     150           0 :   default:
     151           0 :     return FD_EXECUTOR_INSTR_ERR_INVALID_INSTR_DATA;
     152         585 :   }
     153             : 
     154             :   /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L42 */
     155         585 :   uint accessed_accounts = 0U;
     156         585 :   uchar const * context = NULL;
     157             :   /* Note: instr_id is guaranteed to be valid, to access values in the arrays. */
     158         585 :   ulong context_sz = fd_zksdk_context_sz[instr_id];
     159         585 :   ulong proof_data_sz = context_sz + fd_zksdk_proof_sz[instr_id];
     160             : 
     161         585 :   if( ctx->instr->data_sz == 5UL ) {
     162             :     /* Case 1. Proof data from account data. */
     163             : 
     164             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L46-L47 */
     165         237 :     FD_BORROWED_ACCOUNT_DECL( proof_data_acc );
     166         939 :     FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, 0UL, proof_data_acc ) {
     167             : 
     168             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L48 */
     169         234 :       accessed_accounts = 1U;
     170             : 
     171             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L50-L61
     172             :          Note: it doesn't look like the ref code can throw any error. */
     173         234 :       uint proof_data_offset = fd_uint_load_4_fast(&instr_data[1]);
     174             : 
     175             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L62-L65 */
     176         234 :       if( proof_data_offset+proof_data_sz > proof_data_acc->const_meta->dlen ) {
     177           9 :         return FD_EXECUTOR_INSTR_ERR_INVALID_ACC_DATA;
     178           9 :       }
     179         225 :       context = fd_memcpy( buffer+CTX_HEAD_SZ, &proof_data_acc->const_data[proof_data_offset], proof_data_sz );
     180             : 
     181         234 :     } FD_BORROWED_ACCOUNT_DROP( proof_data_acc );
     182         348 :   } else {
     183             :     /* Case 2. Proof data from ix data. */
     184             : 
     185             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L78-L82
     186             :        Note: instr_id is guaranteed to be valid, to access values in the arrays. */
     187         348 :     if (ctx->instr->data_sz != 1 + proof_data_sz) {
     188           6 :       fd_log_collector_msg_literal( ctx, "invalid proof data" );
     189           6 :       return FD_EXECUTOR_INSTR_ERR_INVALID_INSTR_DATA;
     190           6 :     }
     191         342 :     context = instr_data + 1;
     192         342 :   }
     193             : 
     194             :   /* Verify individual ZKP
     195             :      https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L83-L86 */
     196         567 :   void const * proof = context + fd_zksdk_context_sz[instr_id];
     197         567 :   int err = (*fd_zksdk_instr_verify_proof)( context, proof );
     198         567 :   if( FD_UNLIKELY( err ) ) {
     199             :     //TODO: full log, including err
     200         162 :     fd_log_collector_msg_literal( ctx, "proof_verification failed" );
     201         162 :     return FD_EXECUTOR_INSTR_ERR_INVALID_INSTR_DATA;
     202         162 :   }
     203             : 
     204             :   /* Create context state if accounts are provided with the instruction
     205             :      https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L92 */
     206         405 :   if( instr_acc_cnt > accessed_accounts ) {
     207             : 
     208             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L93-L98 */
     209         303 :     fd_pubkey_t context_state_authority[1];
     210         303 :     FD_BORROWED_ACCOUNT_DECL( _acc );
     211        1194 :     FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, accessed_accounts+1, _acc ) {
     212         297 :       fd_memcpy( context_state_authority, _acc->pubkey, sizeof(fd_pubkey_t) );
     213         297 :     } FD_BORROWED_ACCOUNT_DROP( _acc );
     214             : 
     215             :     /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L100-L101 */
     216         297 :     FD_BORROWED_ACCOUNT_DECL( proof_context_acc );
     217        1188 :     FD_BORROWED_ACCOUNT_TRY_BORROW_IDX( ctx, accessed_accounts, proof_context_acc ) {
     218             : 
     219             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L103-L105 */
     220         297 :       if( FD_UNLIKELY( !fd_memeq( proof_context_acc->const_meta->info.owner, &fd_solana_zk_elgamal_proof_program_id, sizeof(fd_pubkey_t) ) ) ) {
     221          15 :         return FD_EXECUTOR_INSTR_ERR_INVALID_ACC_OWNER;
     222          15 :       }
     223             : 
     224             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L107-L112 */
     225         282 :       if( FD_UNLIKELY( proof_context_acc->const_meta->dlen >= CTX_HEAD_SZ && proof_context_acc->const_data[32] != 0 ) ) {
     226           6 :         return FD_EXECUTOR_INSTR_ERR_ACC_ALREADY_INITIALIZED;
     227           6 :       }
     228             : 
     229             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L114-L115
     230             :          Note: nothing to do. */
     231             : 
     232             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L117-L119 */
     233         276 :       ulong context_data_sx = CTX_HEAD_SZ + context_sz;
     234         276 :       if( FD_UNLIKELY( proof_context_acc->const_meta->dlen != context_data_sx ) ) {
     235          12 :         return FD_EXECUTOR_INSTR_ERR_INVALID_ACC_DATA;
     236          12 :       }
     237             : 
     238             :       /* https://github.com/anza-xyz/agave/blob/v2.0.1/programs/zk-elgamal-proof/src/lib.rs#L121 */
     239         264 :       fd_memcpy( buffer, context_state_authority, sizeof(fd_pubkey_t) ); // buffer[0..31]
     240         264 :       buffer[ 32 ] = instr_id;                                           // buffer[32]
     241         264 :       if( ctx->instr->data_sz != 5UL ) {                                  // buffer[33..]
     242         123 :         fd_memcpy( buffer+CTX_HEAD_SZ, context, context_sz );
     243         123 :       }
     244         264 :       err = fd_account_set_data_from_slice( ctx, accessed_accounts, buffer, context_data_sx );
     245         264 :       if( FD_UNLIKELY( err ) ) {
     246           3 :         return err;
     247           3 :       }
     248             : 
     249         297 :     } FD_BORROWED_ACCOUNT_DROP( proof_context_acc );
     250         297 :   }
     251             : 
     252         363 :   return FD_EXECUTOR_INSTR_SUCCESS;
     253         405 : }