add tests for configure account with registry

token/program-2022-test/Cargo.toml

@@ -24,6 +24,7 @@ solana-program = "2.0.3"
 solana-program-test = "2.0.3"
 solana-sdk = "2.0.3"
 spl-associated-token-account = { version = "5.0.1", path = "../../associated-token-account/program" }
+spl-elgamal-registry = { version = "0.1.0", path = "../confidential-transfer/elgamal-registry" }
 spl-memo = { version = "5.0.0", path = "../../memo/program", features = [
   "no-entrypoint",
 ] }
@@ -34,6 +35,7 @@ spl-record = { version = "0.2.1", path = "../../record/program", features = [
 spl-token-2022 = { version = "5.0.2", path = "../program-2022", features = [
   "no-entrypoint",
 ] }
+spl-token-confidential-transfer-proof-extraction = { version = "0.1.0", path = "../confidential-transfer/proof-extraction" }
 spl-token-confidential-transfer-proof-generation = { version = "0.1.0", path = "../confidential-transfer/proof-generation" }
 spl-instruction-padding = { version = "0.2.0", path = "../../instruction-padding/program", features = [
   "no-entrypoint",

token/program-2022-test/tests/confidential_transfer.rs

@@ -13,14 +13,15 @@ use {
         pubkey::Pubkey,
         signature::Signer,
         signer::{keypair::Keypair, signers::Signers},
-        transaction::TransactionError,
+        transaction::{Transaction, TransactionError},
         transport::TransportError,
     },
     spl_record::state::RecordData,
     spl_token_2022::{
         error::TokenError,
         extension::{
             confidential_transfer::{
+                self,
                 account_info::{EmptyAccountAccountInfo, TransferAccountInfo, WithdrawAccountInfo},
                 ConfidentialTransferAccount, MAXIMUM_DEPOSIT_TRANSFER_AMOUNT,
             },
@@ -38,6 +39,7 @@ use {
             TokenResult,
         },
     },
+    spl_token_confidential_transfer_proof_extraction::{ProofData, ProofLocation},
     spl_token_confidential_transfer_proof_generation::{
         transfer::TransferProofData, transfer_with_fee::TransferWithFeeProofData,
         withdraw::WithdrawProofData,
@@ -2810,3 +2812,122 @@ async fn confidential_transfer_transfer_with_fee_and_memo_option(
         )
         .await;
 }
+
+#[tokio::test]
+async fn confidential_transfer_configure_token_account_with_registry() {
+    let authority = Keypair::new();
+    let auto_approve_new_accounts = false;
+    let auditor_elgamal_keypair = ElGamalKeypair::new_rand();
+    let auditor_elgamal_pubkey = (*auditor_elgamal_keypair.pubkey()).into();
+
+    let mut context = TestContext::new().await;
+    context
+        .init_token_with_mint(vec![
+            ExtensionInitializationParams::ConfidentialTransferMint {
+                authority: Some(authority.pubkey()),
+                auto_approve_new_accounts,
+                auditor_elgamal_pubkey: Some(auditor_elgamal_pubkey),
+            },
+        ])
+        .await
+        .unwrap();
+
+    let TokenContext { token, alice, .. } = context.token_context.unwrap();
+    let alice_account_keypair = Keypair::new();
+    token
+        .create_auxiliary_token_account_with_extension_space(
+            &alice_account_keypair,
+            &alice.pubkey(),
+            vec![ExtensionType::ConfidentialTransferAccount],
+        )
+        .await
+        .unwrap();
+    let elgamal_keypair = ElGamalKeypair::new_rand();
+
+    // create ElGamal registry
+    let mut ctx = context.context.lock().await;
+    let proof_data =
+        confidential_transfer::instruction::PubkeyValidityProofData::new(&elgamal_keypair).unwrap();
+    let proof_location = ProofLocation::InstructionOffset(
+        1.try_into().unwrap(),
+        ProofData::InstructionData(&proof_data),
+    );
+
+    let instructions = spl_elgamal_registry::instruction::create_registry(
+        &ctx.payer.pubkey(),
+        &alice.pubkey(),
+        proof_location,
+    )
+    .unwrap();
+    let tx = Transaction::new_signed_with_payer(
+        &instructions,
+        Some(&ctx.payer.pubkey()),
+        &[&ctx.payer],
+        ctx.last_blockhash,
+    );
+    ctx.banks_client.process_transaction(tx).await.unwrap();
+
+    // update ElGamal registry
+    let new_elgamal_keypair =
+        ElGamalKeypair::new_from_signer(&alice, &alice_account_keypair.pubkey().to_bytes())
+            .unwrap();
+    let proof_data =
+        confidential_transfer::instruction::PubkeyValidityProofData::new(&new_elgamal_keypair)
+            .unwrap();
+    let proof_location = ProofLocation::InstructionOffset(
+        1.try_into().unwrap(),
+        ProofData::InstructionData(&proof_data),
+    );
+
+    let elgamal_registry_address = spl_elgamal_registry::get_elgamal_registry_address(
+        &alice.pubkey(),
+        &spl_elgamal_registry::id(),
+    );
+
+    let instructions = spl_elgamal_registry::instruction::update_registry(
+        &elgamal_registry_address,
+        &alice.pubkey(),
+        proof_location,
+    )
+    .unwrap();
+    let tx = Transaction::new_signed_with_payer(
+        &instructions,
+        Some(&ctx.payer.pubkey()),
+        &[&ctx.payer, &alice],
+        ctx.last_blockhash,
+    );
+    ctx.banks_client.process_transaction(tx).await.unwrap();
+    drop(ctx);
+
+    // configure account using ElGamal registry
+    let alice_account_keypair = Keypair::new();
+    let alice_token_account = alice_account_keypair.pubkey();
+    token
+        .create_auxiliary_token_account_with_extension_space(
+            &alice_account_keypair,
+            &alice_token_account,
+            vec![ExtensionType::ConfidentialTransferAccount],
+        )
+        .await
+        .unwrap();
+
+    token
+        .confidential_transfer_configure_token_account_with_registry(
+            &alice_account_keypair.pubkey(),
+            &elgamal_registry_address,
+            &alice.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    let state = token.get_account_info(&alice_token_account).await.unwrap();
+    let extension = state
+        .get_extension::<ConfidentialTransferAccount>()
+        .unwrap();
+    assert!(!bool::from(&extension.approved));
+    assert!(bool::from(&extension.allow_confidential_credits));
+    assert_eq!(
+        extension.elgamal_pubkey,
+        (*new_elgamal_keypair.pubkey()).into()
+    );
+}

token/program-2022-test/tests/program_test.rs

@@ -50,6 +50,11 @@ impl TestContext {
             spl_record::id(),
             processor!(spl_record::processor::process_instruction),
         );
+        program_test.add_program(
+            "spl_elgamal_registry",
+            spl_elgamal_registry::id(),
+            processor!(spl_elgamal_registry::processor::process_instruction),
+        );
         let context = program_test.start_with_context().await;
         let context = Arc::new(Mutex::new(context));