lssa/integration_tests/tests/wallet_ffi.rs

use std::{
    collections::HashSet,
    ffi::{CString, c_char},
    io::Write,
};

use anyhow::Result;
use integration_tests::{ACC_SENDER, TestContext};
use log::info;
use nssa::{Account, AccountId, PublicKey, program::Program};
use tempfile::tempdir;
use wallet::WalletCore;
use wallet_ffi::{
    FfiAccount, FfiAccountList, FfiBytes32, FfiPublicAccountKey, WalletHandle, error,
};

unsafe extern "C" {
    fn wallet_ffi_create_new(
        config_path: *const c_char,
        storage_path: *const c_char,
        password: *const c_char,
    ) -> *mut WalletHandle;

    fn wallet_ffi_create_account_public(
        handle: *mut WalletHandle,
        out_account_id: *mut FfiBytes32,
    ) -> error::WalletFfiError;

    fn wallet_ffi_create_account_private(
        handle: *mut WalletHandle,
        out_account_id: *mut FfiBytes32,
    ) -> error::WalletFfiError;

    fn wallet_ffi_list_accounts(
        handle: *mut WalletHandle,
        out_list: *mut FfiAccountList,
    ) -> error::WalletFfiError;

    fn wallet_ffi_free_account_list(list: *mut FfiAccountList);

    fn wallet_ffi_get_balance(
        handle: *mut WalletHandle,
        account_id: *const FfiBytes32,
        is_public: bool,
        out_balance: *mut [u8; 16],
    ) -> error::WalletFfiError;

    fn wallet_ffi_get_account_public(
        handle: *mut WalletHandle,
        account_id: *const FfiBytes32,
        out_account: *mut FfiAccount,
    ) -> error::WalletFfiError;

    fn wallet_ffi_free_account_data(account: *mut FfiAccount);

    fn wallet_ffi_get_public_account_key(
        handle: *mut WalletHandle,
        account_id: *const FfiBytes32,
        out_public_key: *mut FfiPublicAccountKey,
    ) -> error::WalletFfiError;
}

fn new_wallet_ffi_with_test_context_config(ctx: &TestContext) -> *mut WalletHandle {
    let tempdir = tempfile::tempdir().unwrap();
    let config_path = tempdir.path().join("wallet_config.json");
    let storage_path = tempdir.path().join("storage.json");
    let mut config = ctx.wallet().config().to_owned();
    if let Some(config_overrides) = ctx.wallet().config_overrides().clone() {
        config.apply_overrides(config_overrides);
    }
    let mut file = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .open(&config_path)
        .unwrap();

    let config_with_overrides_serialized = serde_json::to_vec_pretty(&config).unwrap();

    file.write_all(&config_with_overrides_serialized).unwrap();

    let config_path = CString::new(config_path.to_str().unwrap()).unwrap();
    let storage_path = CString::new(storage_path.to_str().unwrap()).unwrap();
    let password = CString::new(ctx.wallet_password()).unwrap();

    unsafe {
        wallet_ffi_create_new(
            config_path.as_ptr(),
            storage_path.as_ptr(),
            password.as_ptr(),
        )
    }
}

fn new_wallet_ffi_with_default_config(password: &str) -> *mut WalletHandle {
    let tempdir = tempdir().unwrap();
    let config_path = tempdir.path().join("wallet_config.json");
    let storage_path = tempdir.path().join("storage.json");
    let config_path_c = CString::new(config_path.to_str().unwrap()).unwrap();
    let storage_path_c = CString::new(storage_path.to_str().unwrap()).unwrap();
    let password = CString::new(password).unwrap();

    unsafe {
        wallet_ffi_create_new(
            config_path_c.as_ptr(),
            storage_path_c.as_ptr(),
            password.as_ptr(),
        )
    }
}

fn new_wallet_rust_with_default_config(password: &str) -> WalletCore {
    let tempdir = tempdir().unwrap();
    let config_path = tempdir.path().join("wallet_config.json");
    let storage_path = tempdir.path().join("storage.json");

    WalletCore::new_init_storage(
        config_path.to_path_buf(),
        storage_path.to_path_buf(),
        None,
        password.to_string(),
    )
    .unwrap()
}

#[test]
fn test_wallet_ffi_create_public_accounts() {
    let password = "password_for_tests";
    let n_accounts = 10;
    // First `n_accounts` public accounts created with Rust wallet
    let new_public_account_ids_rust = {
        let mut account_ids = Vec::new();

        let mut wallet_rust = new_wallet_rust_with_default_config(password);
        for _ in 0..n_accounts {
            let account_id = wallet_rust.create_new_account_public(None).0;
            account_ids.push(*account_id.value());
        }
        account_ids
    };

    // First `n_accounts` public accounts created with wallet FFI
    let new_public_account_ids_ffi = unsafe {
        let mut account_ids = Vec::new();

        let wallet_ffi_handle = new_wallet_ffi_with_default_config(password);
        for _ in 0..n_accounts {
            let mut out_account_id = FfiBytes32::from_bytes([0; 32]);
            wallet_ffi_create_account_public(
                wallet_ffi_handle,
                (&mut out_account_id) as *mut FfiBytes32,
            );
            account_ids.push(out_account_id.data);
        }
        account_ids
    };

    assert_eq!(new_public_account_ids_ffi, new_public_account_ids_rust)
}

#[test]
fn test_wallet_ffi_create_private_accounts() {
    let password = "password_for_tests";
    let n_accounts = 10;
    // First `n_accounts` private accounts created with Rust wallet
    let new_private_account_ids_rust = {
        let mut account_ids = Vec::new();

        let mut wallet_rust = new_wallet_rust_with_default_config(password);
        for _ in 0..n_accounts {
            let account_id = wallet_rust.create_new_account_private(None).0;
            account_ids.push(*account_id.value());
        }
        account_ids
    };

    // First `n_accounts` private accounts created with wallet FFI
    let new_private_account_ids_ffi = unsafe {
        let mut account_ids = Vec::new();

        let wallet_ffi_handle = new_wallet_ffi_with_default_config(password);
        for _ in 0..n_accounts {
            let mut out_account_id = FfiBytes32::from_bytes([0; 32]);
            wallet_ffi_create_account_private(
                wallet_ffi_handle,
                (&mut out_account_id) as *mut FfiBytes32,
            );
            account_ids.push(out_account_id.data);
        }
        account_ids
    };

    assert_eq!(new_private_account_ids_ffi, new_private_account_ids_rust)
}

#[test]
fn test_wallet_ffi_list_accounts() {
    let password = "password_for_tests";

    // Create the wallet FFI
    let wallet_ffi_handle = unsafe {
        let handle = new_wallet_ffi_with_default_config(password);
        // Create 5 public accounts and 5 private accounts
        for _ in 0..5 {
            let mut out_account_id = FfiBytes32::from_bytes([0; 32]);
            wallet_ffi_create_account_public(handle, (&mut out_account_id) as *mut FfiBytes32);
            wallet_ffi_create_account_private(handle, (&mut out_account_id) as *mut FfiBytes32);
        }

        handle
    };

    // Create the wallet Rust
    let wallet_rust = {
        let mut wallet = new_wallet_rust_with_default_config(password);
        // Create 5 public accounts and 5 private accounts
        for _ in 0..5 {
            wallet.create_new_account_public(None);
            wallet.create_new_account_private(None);
        }
        wallet
    };

    // Get the account list with FFI method
    let mut wallet_ffi_account_list = unsafe {
        let mut out_list = FfiAccountList::default();
        wallet_ffi_list_accounts(wallet_ffi_handle, (&mut out_list) as *mut FfiAccountList);
        out_list
    };

    let wallet_rust_account_ids = wallet_rust
        .storage()
        .user_data
        .account_ids()
        .collect::<Vec<_>>();

    // Assert same number of elements between Rust and FFI result
    assert_eq!(wallet_rust_account_ids.len(), wallet_ffi_account_list.count);

    let wallet_ffi_account_list_slice = unsafe {
        core::slice::from_raw_parts(
            wallet_ffi_account_list.entries,
            wallet_ffi_account_list.count,
        )
    };

    // Assert same account ids between Rust and FFI result
    assert_eq!(
        wallet_rust_account_ids
            .iter()
            .map(|id| id.value())
            .collect::<HashSet<_>>(),
        wallet_ffi_account_list_slice
            .iter()
            .map(|entry| &entry.account_id.data)
            .collect::<HashSet<_>>()
    );

    // Assert `is_pub` flag is correct in the FFI result
    for entry in wallet_ffi_account_list_slice.iter() {
        let account_id = AccountId::new(entry.account_id.data);
        let is_pub_default_in_rust_wallet = wallet_rust
            .storage()
            .user_data
            .default_pub_account_signing_keys
            .contains_key(&account_id);
        let is_pub_key_tree_wallet_rust = wallet_rust
            .storage()
            .user_data
            .public_key_tree
            .account_id_map
            .contains_key(&account_id);

        let is_public_in_rust_wallet = is_pub_default_in_rust_wallet || is_pub_key_tree_wallet_rust;

        assert_eq!(entry.is_public, is_public_in_rust_wallet);
    }

    unsafe { wallet_ffi_free_account_list((&mut wallet_ffi_account_list) as *mut FfiAccountList) };
}

#[test]
fn test_wallet_ffi_get_balance_public() -> Result<()> {
    let ctx = TestContext::new_blocking()?;
    let account_id: AccountId = ACC_SENDER.parse().unwrap();
    let wallet_ffi_handle = new_wallet_ffi_with_test_context_config(&ctx);

    let balance = unsafe {
        let mut out_balance: [u8; 16] = [0; 16];
        let ffi_account_id = FfiBytes32::from(&account_id);
        let _result = wallet_ffi_get_balance(
            wallet_ffi_handle,
            (&ffi_account_id) as *const FfiBytes32,
            true,
            (&mut out_balance) as *mut [u8; 16],
        );
        out_balance
    };
    let balance = u128::from_le_bytes(balance);
    assert_eq!(balance, 10000);

    info!("Successfully retrieved account balance");

    Ok(())
}

#[test]
fn test_wallet_ffi_get_account_public() -> Result<()> {
    let ctx = TestContext::new_blocking()?;
    let account_id: AccountId = ACC_SENDER.parse().unwrap();
    let wallet_ffi_handle = new_wallet_ffi_with_test_context_config(&ctx);
    let mut out_account = FfiAccount::default();

    let account: Account = unsafe {
        let ffi_account_id = FfiBytes32::from(&account_id);
        let _result = wallet_ffi_get_account_public(
            wallet_ffi_handle,
            (&ffi_account_id) as *const FfiBytes32,
            (&mut out_account) as *mut FfiAccount,
        );
        (&out_account).try_into().unwrap()
    };

    assert_eq!(
        account.program_owner,
        Program::authenticated_transfer_program().id()
    );
    assert_eq!(account.balance, 10000);
    assert!(account.data.is_empty());
    assert_eq!(account.nonce, 0);

    unsafe {
        wallet_ffi_free_account_data((&mut out_account) as *mut FfiAccount);
    }

    info!("Successfully retrieved account with correct details");

    Ok(())
}

#[test]
fn test_wallet_ffi_get_public_account_keys() -> Result<()> {
    let ctx = TestContext::new_blocking()?;
    let account_id: AccountId = ACC_SENDER.parse().unwrap();
    let wallet_ffi_handle = new_wallet_ffi_with_test_context_config(&ctx);
    let mut out_key = FfiPublicAccountKey::default();

    let key: PublicKey = unsafe {
        let ffi_account_id = FfiBytes32::from(&account_id);
        let _result = wallet_ffi_get_public_account_key(
            wallet_ffi_handle,
            (&ffi_account_id) as *const FfiBytes32,
            (&mut out_key) as *mut FfiPublicAccountKey,
        );
        (&out_key).try_into().unwrap()
    };

    let expected_key = {
        let private_key = ctx
            .wallet()
            .get_account_public_signing_key(&account_id)
            .unwrap();
        PublicKey::new_from_private_key(private_key)
    };

    assert_eq!(key, expected_key);

    info!("Successfully retrieved account key");

    Ok(())
}