logos-execution-zone/integration_tests/tests/keys.rs

#![expect(
    clippy::shadow_unrelated,
    clippy::tests_outside_test_module,
    reason = "We don't care about these in tests"
)]

use std::{str::FromStr as _, time::Duration};

use anyhow::{Context as _, Result};
use integration_tests::{
    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx, private_mention,
    public_mention, verify_commitment_is_in_state,
};
use key_protocol::key_management::key_tree::chain_index::ChainIndex;
use lee::AccountId;
use log::info;
use sequencer_service_rpc::RpcClient as _;
use tokio::test;
use wallet::{
    cli::{
        CliAccountMention, Command, SubcommandReturnValue,
        account::{AccountSubcommand, NewSubcommand},
        programs::native_token_transfer::AuthTransferSubcommand,
    },
    storage::key_chain::FoundPrivateAccount,
};

/// Create a private or public account at the given chain index and return its ID.
async fn new_account(ctx: &mut TestContext, private: bool, cci: ChainIndex) -> Result<AccountId> {
    let subcommand = if private {
        NewSubcommand::Private {
            cci: Some(cci),
            label: None,
        }
    } else {
        NewSubcommand::Public {
            cci: Some(cci),
            label: None,
        }
    };
    let result = wallet::cli::execute_subcommand(
        ctx.wallet_mut(),
        Command::Account(AccountSubcommand::New(subcommand)),
    )
    .await?;
    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
        anyhow::bail!("Expected RegisterAccount return value");
    };
    Ok(account_id)
}

/// Send `amount` from `from` to `to` via an authenticated transfer (identifier 0).
async fn send(
    ctx: &mut TestContext,
    from: CliAccountMention,
    to: CliAccountMention,
    amount: u128,
) -> Result<()> {
    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from,
        to: Some(to),
        to_npk: None,
        to_vpk: None,
        to_keys: None,
        to_identifier: Some(0),
        amount,
    });
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    Ok(())
}

/// Look up the restored private account for `account_id`, asserting it exists.
fn restored_private_account<'a>(
    ctx: &'a TestContext,
    account_id: AccountId,
    label: &str,
) -> FoundPrivateAccount<'a> {
    ctx.wallet()
        .storage()
        .key_chain()
        .private_account(account_id)
        .unwrap_or_else(|| panic!("{label} should be restored"))
}

/// Assert that a restored public account's signing key exists.
fn assert_public_account_restored(ctx: &TestContext, account_id: AccountId, label: &str) {
    ctx.wallet()
        .storage()
        .key_chain()
        .pub_account_signing_key(account_id)
        .unwrap_or_else(|| panic!("{label} should be restored"));
}

#[test]
async fn sync_private_account_with_non_zero_chain_index() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let from: AccountId = ctx.existing_private_accounts()[0];

    // Create a new private account
    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
        cci: None,
        label: None,
    }));

    for _ in 0..3 {
        // Key Tree shift
        // This way we have account with child index > 0.
        let result = wallet::cli::execute_subcommand(
            ctx.wallet_mut(),
            Command::Account(AccountSubcommand::New(NewSubcommand::Private {
                cci: None,
                label: None,
            })),
        )
        .await?;
        let SubcommandReturnValue::RegisterAccount { account_id: _ } = result else {
            anyhow::bail!("Expected RegisterAccount return value");
        };
    }

    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::RegisterAccount {
        account_id: to_account_id,
    } = sub_ret
    else {
        anyhow::bail!("Expected RegisterAccount return value");
    };

    // Get the keys for the newly created account
    let to_account = ctx
        .wallet()
        .storage()
        .key_chain()
        .private_account(to_account_id)
        .context("Failed to get private account")?;

    // Send to this account using claiming path (using npk and vpk instead of account ID)
    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from: private_mention(from),
        to: None,
        to_npk: Some(hex::encode(to_account.key_chain.nullifier_public_key.0)),
        to_vpk: Some(hex::encode(
            to_account.key_chain.viewing_public_key.to_bytes(),
        )),
        to_keys: None,
        to_identifier: Some(to_account.kind.identifier()),
        amount: 100,
    });

    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash } = sub_ret else {
        anyhow::bail!("Expected PrivacyPreservingTransfer return value");
    };

    let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash).await;

    // Sync the wallet to claim the new account
    let command = Command::Account(AccountSubcommand::SyncPrivate {});
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;

    let new_commitment1 = ctx
        .wallet()
        .get_private_account_commitment(from)
        .context("Failed to get private account commitment for sender")?;
    assert_eq!(tx.message.new_commitments[0], new_commitment1);

    assert_eq!(tx.message.new_commitments.len(), 2);
    for commitment in tx.message.new_commitments {
        assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await);
    }

    let to_res_acc = ctx
        .wallet()
        .get_account_private(to_account_id)
        .context("Failed to get recipient's private account")?;
    assert_eq!(to_res_acc.balance, 100);

    info!("Successfully transferred using claiming path");

    Ok(())
}

#[test]
async fn restore_keys_from_seed() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let from: AccountId = ctx.existing_private_accounts()[0];

    // Create private accounts at root and /0
    let to_account_id1 = new_account(&mut ctx, true, ChainIndex::root()).await?;
    let to_account_id2 = new_account(&mut ctx, true, ChainIndex::from_str("/0")?).await?;

    // Send to both private accounts
    send(
        &mut ctx,
        private_mention(from),
        private_mention(to_account_id1),
        100,
    )
    .await?;
    send(
        &mut ctx,
        private_mention(from),
        private_mention(to_account_id2),
        101,
    )
    .await?;

    let from: AccountId = ctx.existing_public_accounts()[0];

    // Create public accounts at root and /0
    let to_account_id3 = new_account(&mut ctx, false, ChainIndex::root()).await?;
    let to_account_id4 = new_account(&mut ctx, false, ChainIndex::from_str("/0")?).await?;

    // Send to both public accounts
    send(
        &mut ctx,
        public_mention(from),
        public_mention(to_account_id3),
        102,
    )
    .await?;
    send(
        &mut ctx,
        public_mention(from),
        public_mention(to_account_id4),
        103,
    )
    .await?;

    info!("Preparation complete, performing keys restoration");

    // Restore keys from seed
    wallet::cli::execute_keys_restoration(ctx.wallet_mut(), 10).await?;

    // Verify restored private accounts
    let acc1 = restored_private_account(&ctx, to_account_id1, "Acc 1");
    let acc2 = restored_private_account(&ctx, to_account_id2, "Acc 2");

    // Verify restored public accounts
    assert_public_account_restored(&ctx, to_account_id3, "Acc 3");
    assert_public_account_restored(&ctx, to_account_id4, "Acc 4");

    assert_eq!(
        acc1.account.program_owner,
        programs::authenticated_transfer().id()
    );
    assert_eq!(
        acc2.account.program_owner,
        programs::authenticated_transfer().id()
    );

    assert_eq!(acc1.account.balance, 100);
    assert_eq!(acc2.account.balance, 101);

    info!("Tree checks passed, testing restored accounts can transact");

    // Test that restored accounts can send transactions
    send(
        &mut ctx,
        private_mention(to_account_id1),
        private_mention(to_account_id2),
        10,
    )
    .await?;
    send(
        &mut ctx,
        public_mention(to_account_id3),
        public_mention(to_account_id4),
        11,
    )
    .await?;

    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;

    // Verify commitments exist for private accounts
    let comm1 = ctx
        .wallet()
        .get_private_account_commitment(to_account_id1)
        .expect("Acc 1 commitment should exist");
    let comm2 = ctx
        .wallet()
        .get_private_account_commitment(to_account_id2)
        .expect("Acc 2 commitment should exist");

    assert!(verify_commitment_is_in_state(comm1, ctx.sequencer_client()).await);
    assert!(verify_commitment_is_in_state(comm2, ctx.sequencer_client()).await);

    // Verify public account balances
    let acc3 = ctx
        .sequencer_client()
        .get_account_balance(to_account_id3)
        .await?;
    let acc4 = ctx
        .sequencer_client()
        .get_account_balance(to_account_id4)
        .await?;

    assert_eq!(acc3, 91); // 102 - 11
    assert_eq!(acc4, 114); // 103 + 11

    info!("Successfully restored keys and verified transactions");

    Ok(())
}