lssa/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::{KeyChain, key_tree::chain_index::ChainIndex};
use log::info;
use nssa::{AccountId, Data, program::Program};
use nssa_core::account::Nonce;
use sequencer_service_rpc::RpcClient as _;
use tokio::test;
use wallet::{
    account::HumanReadableAccount,
    cli::{
        Command, SubcommandReturnValue,
        account::{AccountSubcommand, NewSubcommand},
        import::ImportSubcommand,
        programs::native_token_transfer::AuthTransferSubcommand,
    },
};

#[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_keys, _) = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_private_account(to_account_id)
        .cloned()
        .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_keys.nullifier_public_key.0)),
        to_vpk: Some(hex::encode(to_keys.viewing_public_key.0)),
        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 first private account at root
    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
        cci: Some(ChainIndex::root()),
        label: None,
    }));
    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::RegisterAccount {
        account_id: to_account_id1,
    } = result
    else {
        anyhow::bail!("Expected RegisterAccount return value");
    };

    // Create second private account at /0
    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
        cci: Some(ChainIndex::from_str("/0")?),
        label: None,
    }));
    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::RegisterAccount {
        account_id: to_account_id2,
    } = result
    else {
        anyhow::bail!("Expected RegisterAccount return value");
    };

    // Send to first private account
    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from: private_mention(from),
        to: Some(private_mention(to_account_id1)),
        to_npk: None,
        to_vpk: None,
        amount: 100,
    });
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;

    // Send to second private account
    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from: private_mention(from),
        to: Some(private_mention(to_account_id2)),
        to_npk: None,
        to_vpk: None,
        amount: 101,
    });
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;

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

    // Create first public account at root
    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
        cci: Some(ChainIndex::root()),
        label: None,
    }));
    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::RegisterAccount {
        account_id: to_account_id3,
    } = result
    else {
        anyhow::bail!("Expected RegisterAccount return value");
    };

    // Create second public account at /0
    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
        cci: Some(ChainIndex::from_str("/0")?),
        label: None,
    }));
    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::RegisterAccount {
        account_id: to_account_id4,
    } = result
    else {
        anyhow::bail!("Expected RegisterAccount return value");
    };

    // Send to first public account
    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from: public_mention(from),
        to: Some(public_mention(to_account_id3)),
        to_npk: None,
        to_vpk: None,
        amount: 102,
    });
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;

    // Send to second public account
    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from: public_mention(from),
        to: Some(public_mention(to_account_id4)),
        to_npk: None,
        to_vpk: None,
        amount: 103,
    });
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).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 = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_private_account(to_account_id1)
        .expect("Acc 1 should be restored");

    let acc2 = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_private_account(to_account_id2)
        .expect("Acc 2 should be restored");

    // Verify restored public accounts
    let _acc3 = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_pub_account_signing_key(to_account_id3)
        .expect("Acc 3 should be restored");

    let _acc4 = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_pub_account_signing_key(to_account_id4)
        .expect("Acc 4 should be restored");

    assert_eq!(
        acc1.1.program_owner,
        Program::authenticated_transfer_program().id()
    );
    assert_eq!(
        acc2.1.program_owner,
        Program::authenticated_transfer_program().id()
    );

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

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

    // Test that restored accounts can send transactions
    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from: private_mention(to_account_id1),
        to: Some(private_mention(to_account_id2)),
        to_npk: None,
        to_vpk: None,
        amount: 10,
    });
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;

    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
        from: public_mention(to_account_id3),
        to: Some(public_mention(to_account_id4)),
        to_npk: None,
        to_vpk: None,
        amount: 11,
    });
    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).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(())
}

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

    let private_key = nssa::PrivateKey::new_os_random();
    let account_id = nssa::AccountId::from(&nssa::PublicKey::new_from_private_key(&private_key));

    let command = Command::Import(ImportSubcommand::Public { private_key });
    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::Empty = sub_ret else {
        anyhow::bail!("Expected Empty return value");
    };

    let imported_key = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_pub_account_signing_key(account_id);
    assert!(
        imported_key.is_some(),
        "Imported public account should be present"
    );

    Ok(())
}

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

    let key_chain = KeyChain::new_os_random();
    let account_id = nssa::AccountId::from(&key_chain.nullifier_public_key);
    let account = nssa::Account {
        program_owner: Program::authenticated_transfer_program().id(),
        balance: 777,
        data: Data::default(),
        nonce: Nonce::default(),
    };

    let key_chain_json = serde_json::to_string(&key_chain)
        .context("Failed to serialize key chain for private import")?;
    let account_state = HumanReadableAccount::from(account.clone());

    let command = Command::Import(ImportSubcommand::Private {
        key_chain_json,
        account_state,
    });
    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::Empty = sub_ret else {
        anyhow::bail!("Expected Empty return value");
    };

    let imported = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_private_account(account_id)
        .context("Imported private account should be present")?;

    assert_eq!(
        imported.0.nullifier_public_key,
        key_chain.nullifier_public_key
    );
    assert_eq!(imported.1, account);

    Ok(())
}

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

    let key_chain = KeyChain::new_os_random();
    let account_id = nssa::AccountId::from(&key_chain.nullifier_public_key);
    let key_chain_json =
        serde_json::to_string(&key_chain).context("Failed to serialize key chain")?;

    let initial_account = nssa::Account {
        program_owner: Program::authenticated_transfer_program().id(),
        balance: 100,
        data: Data::default(),
        nonce: Nonce::default(),
    };

    // First import
    wallet::cli::execute_subcommand(
        ctx.wallet_mut(),
        Command::Import(ImportSubcommand::Private {
            key_chain_json: key_chain_json.clone(),
            account_state: HumanReadableAccount::from(initial_account),
        }),
    )
    .await?;

    let updated_account = nssa::Account {
        program_owner: Program::authenticated_transfer_program().id(),
        balance: 999,
        data: Data::default(),
        nonce: Nonce::default(),
    };

    // Second import with different account data (same key chain)
    wallet::cli::execute_subcommand(
        ctx.wallet_mut(),
        Command::Import(ImportSubcommand::Private {
            key_chain_json,
            account_state: HumanReadableAccount::from(updated_account.clone()),
        }),
    )
    .await?;

    let imported = ctx
        .wallet()
        .storage()
        .key_chain()
        .get_private_account(account_id)
        .context("Imported private account should be present")?;

    assert_eq!(
        imported.1, updated_account,
        "Second import should override account data"
    );

    Ok(())
}