miden-node-store 0.17.2

use std::time::{SystemTime, UNIX_EPOCH};

use assert_matches::assert_matches;
use miden_node_db::migration::{SchemaHash, SchemaHashes};
use miden_protocol::account::AccountId;
use miden_protocol::testing::account_id::{
    ACCOUNT_ID_PRIVATE_SENDER,
    ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
    ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE,
};
use tempfile::TempDir;

use super::{
    AccountAllowlist,
    AllowlistError,
    InvitationCode,
    InvitationEntry,
    InvitationImportOutcome,
    InvitationStatus,
    RegistrationOutcome,
};
use crate::DataDirectory;

fn setup() -> (TempDir, AccountAllowlist) {
    let dir = tempfile::tempdir().unwrap();
    AccountAllowlist::bootstrap(data_directory(&dir).allowlist_database_path()).unwrap();
    let registry = reopen(&dir);
    (dir, registry)
}

fn reopen(dir: &TempDir) -> AccountAllowlist {
    AccountAllowlist::load(data_directory(dir).allowlist_database_path()).unwrap()
}

fn data_directory(dir: &TempDir) -> DataDirectory {
    DataDirectory::load(dir.path().to_path_buf()).unwrap()
}

fn account(index: usize) -> AccountId {
    [
        ACCOUNT_ID_PRIVATE_SENDER,
        ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE,
        ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
    ][index]
        .try_into()
        .unwrap()
}

fn invitation(value: u8) -> InvitationCode {
    InvitationCode::new(&format!("invitation-{value}")).unwrap()
}

fn entry(value: u8, account_id: Option<AccountId>) -> InvitationEntry {
    InvitationEntry {
        invitation_code: invitation(value),
        account_id,
    }
}

#[test]
fn migration_schema_hashes_are_stable() {
    const EXPECTED: [SchemaHash; 1] = [SchemaHash::from_hex(
        "e46c4bb6724dfd7b0da02c11b471ef4535bf4acc577ebe8306cbd14fd2ced399",
    )];
    let migrator = super::migrations::migrator().unwrap();
    pretty_assertions::assert_eq!(migrator.schema_hashes(), SchemaHashes(&EXPECTED));
}

#[tokio::test]
async fn registrations_and_allowlist_timestamps_persist() {
    const ALLOWLISTED_AT: i64 = 946_684_800;

    let (dir, registry) = setup();
    let before =
        i64::try_from(SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs()).unwrap();
    registry.import_invitation(entry(1, None)).await.unwrap();
    registry.import_invitation(entry(2, Some(account(1)))).await.unwrap();
    registry.add_account(account(2)).await.unwrap();
    let after =
        i64::try_from(SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs()).unwrap();

    let path = data_directory(&dir).allowlist_database_path();
    let timestamps = registry
        .reader
        .db
        .read("allowlist_timestamps", |tx| {
            tx.query("SELECT allowlisted_at FROM account_allowlist", &[], |row| row.get::<i64>(0))
        })
        .await
        .unwrap();
    assert_eq!(timestamps.len(), 3);
    assert!(timestamps.iter().all(|timestamp| (before..=after).contains(timestamp)));
    // An earlier timestamp detects replacement without a clock delay.
    registry
        .writer
        .write("set_allowlist_timestamps", |tx| {
            tx.execute("UPDATE account_allowlist SET allowlisted_at = ?1", &[&ALLOWLISTED_AT])
        })
        .await
        .unwrap();

    registry.register_account(invitation(1), account(0)).await.unwrap();
    registry.import_invitation(entry(1, None)).await.unwrap();
    registry.import_invitation(entry(2, Some(account(1)))).await.unwrap();
    for index in 0..3 {
        assert!(!registry.add_account(account(index)).await.unwrap());
    }
    drop(registry);

    AccountAllowlist::migrate(&path).unwrap();
    let registry = reopen(&dir);
    for index in 0..3 {
        assert!(registry.contains_account(account(index)).await.unwrap());
    }
    for (code, id) in [(1, account(0)), (2, account(1))] {
        assert_eq!(
            registry.invitation_status(invitation(code)).await.unwrap(),
            InvitationStatus::Registered(id)
        );
    }
    let timestamps = registry
        .reader
        .db
        .read("allowlist_timestamps", |tx| {
            tx.query("SELECT allowlisted_at FROM account_allowlist", &[], |row| row.get::<i64>(0))
        })
        .await
        .unwrap();
    assert_eq!(timestamps, vec![ALLOWLISTED_AT; 3]);
}

#[tokio::test]
async fn registration_rules() {
    let (_dir, registry) = setup();
    let reader = registry.reader();

    assert!(!reader.contains_account(account(0)).await.unwrap());
    assert_eq!(
        reader.invitation_status(invitation(1)).await.unwrap(),
        InvitationStatus::Unknown
    );
    assert_matches!(
        registry.register_account(invitation(1), account(0)).await,
        Err(AllowlistError::InvitationNotFound)
    );

    registry.import_invitation(entry(1, None)).await.unwrap();
    registry.import_invitation(entry(2, None)).await.unwrap();
    assert_eq!(reader.invitation_status(invitation(1)).await.unwrap(), InvitationStatus::Unused);
    assert!(!reader.contains_account(account(0)).await.unwrap());

    assert_eq!(
        registry.register_account(invitation(1), account(0)).await.unwrap(),
        RegistrationOutcome::Registered
    );
    assert_eq!(
        registry.register_account(invitation(1), account(0)).await.unwrap(),
        RegistrationOutcome::AlreadyRegistered
    );
    assert_matches!(
        registry.register_account(invitation(1), account(1)).await,
        Err(AllowlistError::InvitationAlreadyUsed)
    );
    assert_matches!(
        registry.register_account(invitation(2), account(0)).await,
        Err(AllowlistError::AccountAlreadyRegistered(id)) if id == account(0)
    );
    assert!(!registry.contains_account(account(1)).await.unwrap());
    assert!(!registry.add_account(account(0)).await.unwrap());
    assert!(registry.add_account(account(1)).await.unwrap());
    assert!(!registry.add_account(account(1)).await.unwrap());
    assert_matches!(
        registry.register_account(invitation(2), account(1)).await,
        Err(AllowlistError::AccountAlreadyRegistered(id)) if id == account(1)
    );
    assert_eq!(
        reader.invitation_status(invitation(1)).await.unwrap(),
        InvitationStatus::Registered(account(0))
    );
    assert_eq!(reader.invitation_status(invitation(2)).await.unwrap(), InvitationStatus::Unused);
    assert!(reader.contains_account(account(0)).await.unwrap());
    assert!(reader.contains_account(account(1)).await.unwrap());
}

#[tokio::test]
async fn invitation_import_preserves_registrations_on_conflicts_and_retries() {
    let (_dir, registry) = setup();
    assert_eq!(
        registry.import_invitation(entry(1, None)).await.unwrap(),
        InvitationImportOutcome {
            invitation_added: true,
            registered_account: None
        }
    );
    assert_eq!(
        registry.import_invitation(entry(2, Some(account(0)))).await.unwrap(),
        InvitationImportOutcome {
            invitation_added: true,
            registered_account: Some(account(0))
        }
    );

    assert_matches!(
        registry.import_invitation(entry(2, Some(account(1)))).await,
        Err(AllowlistError::InvitationAlreadyUsed)
    );
    for code in [1, 3] {
        assert_matches!(
            registry.import_invitation(entry(code, Some(account(0)))).await,
            Err(AllowlistError::AccountAlreadyRegistered(id)) if id == account(0)
        );
    }

    assert_eq!(
        registry.invitation_status(invitation(1)).await.unwrap(),
        InvitationStatus::Unused
    );
    assert_eq!(
        registry.invitation_status(invitation(2)).await.unwrap(),
        InvitationStatus::Registered(account(0))
    );
    assert_eq!(
        registry.invitation_status(invitation(3)).await.unwrap(),
        InvitationStatus::Unknown
    );
    assert!(!registry.contains_account(account(1)).await.unwrap());
    assert!(!registry.contains_account(account(2)).await.unwrap());

    assert_eq!(
        registry.import_invitation(entry(1, Some(account(2)))).await.unwrap(),
        InvitationImportOutcome {
            invitation_added: false,
            registered_account: Some(account(2))
        }
    );
    assert_eq!(
        registry.import_invitation(entry(3, Some(account(1)))).await.unwrap(),
        InvitationImportOutcome {
            invitation_added: true,
            registered_account: Some(account(1))
        }
    );
    for entry in [entry(1, Some(account(2))), entry(3, Some(account(1)))] {
        assert_eq!(
            registry.import_invitation(entry).await.unwrap(),
            InvitationImportOutcome {
                invitation_added: false,
                registered_account: None
            }
        );
    }
    assert_eq!(
        registry.invitation_status(invitation(1)).await.unwrap(),
        InvitationStatus::Registered(account(2))
    );
    assert_eq!(
        registry.invitation_status(invitation(3)).await.unwrap(),
        InvitationStatus::Registered(account(1))
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn concurrent_accounts_cannot_claim_the_same_invitation() {
    let (dir, registry) = setup();
    registry.import_invitation(entry(1, None)).await.unwrap();
    let other = reopen(&dir);

    let (first, second) = tokio::join!(
        registry.register_account(invitation(1), account(0)),
        other.register_account(invitation(1), account(1)),
    );
    let winner = match (first, second) {
        (Ok(RegistrationOutcome::Registered), Err(AllowlistError::InvitationAlreadyUsed)) => 0,
        (Err(AllowlistError::InvitationAlreadyUsed), Ok(RegistrationOutcome::Registered)) => 1,
        results => panic!("expected one successful registration, got {results:?}"),
    };
    assert_eq!(
        registry.invitation_status(invitation(1)).await.unwrap(),
        InvitationStatus::Registered(account(winner))
    );
    assert!(registry.contains_account(account(winner)).await.unwrap());
    assert!(!registry.contains_account(account(1 - winner)).await.unwrap());
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn concurrent_invitations_cannot_register_the_same_account() {
    let (dir, registry) = setup();
    registry.import_invitation(entry(1, None)).await.unwrap();
    registry.import_invitation(entry(2, None)).await.unwrap();
    let other = reopen(&dir);

    let (first, second) = tokio::join!(
        registry.register_account(invitation(1), account(0)),
        other.register_account(invitation(2), account(0)),
    );
    let winner = match (first, second) {
        (Ok(RegistrationOutcome::Registered), Err(AllowlistError::AccountAlreadyRegistered(_))) => {
            1
        },
        (Err(AllowlistError::AccountAlreadyRegistered(_)), Ok(RegistrationOutcome::Registered)) => {
            2
        },
        results => panic!("expected one successful registration, got {results:?}"),
    };
    assert_eq!(
        registry.invitation_status(invitation(winner)).await.unwrap(),
        InvitationStatus::Registered(account(0))
    );
    assert_eq!(
        registry.invitation_status(invitation(3 - winner)).await.unwrap(),
        InvitationStatus::Unused
    );
}

#[test]
fn invitation_codes_preserve_text_and_hide_debug_values() {
    assert!(InvitationCode::new("").is_err());
    let code = InvitationCode::new("code").unwrap();
    for different in ["Code", " code ", "code\n"] {
        assert_ne!(code, InvitationCode::new(different).unwrap());
    }
    assert_ne!(InvitationCode::new("\u{e9}").unwrap(), InvitationCode::new("e\u{301}").unwrap());
    let invitation = InvitationCode::new("private invitation code").unwrap();
    assert_eq!(
        InvitationCode::from_hex_digest(&invitation.to_hex_digest()).unwrap(),
        invitation
    );
    let debug = format!("{invitation:?}");
    assert!(!debug.contains("private invitation code"));
    assert!(!debug.contains(&hex::encode(invitation.digest())));
}