diff --git a/Cargo.lock b/Cargo.lock index 0d552ff0..a5bf60d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -274,6 +274,13 @@ dependencies = [ "soroban-sdk", ] +[[package]] +name = "creator-keys-factory" +version = "0.1.0" +dependencies = [ + "soroban-sdk", +] + [[package]] name = "crypto-bigint" version = "0.5.5" diff --git a/Cargo.toml b/Cargo.toml index 38b46c70..0deebe5f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["creator-keys"] +members = ["creator-keys", "creator-keys-factory"] resolver = "2" [workspace.package] diff --git a/creator-keys-factory/Cargo.toml b/creator-keys-factory/Cargo.toml new file mode 100644 index 00000000..3c0b7c52 --- /dev/null +++ b/creator-keys-factory/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "creator-keys-factory" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["cdylib", "rlib"] + +[dependencies] +soroban-sdk = { workspace = true } + +[dev-dependencies] +soroban-sdk = { workspace = true, features = ["testutils"] } diff --git a/creator-keys-factory/src/lib.rs b/creator-keys-factory/src/lib.rs new file mode 100644 index 00000000..2caf69fa --- /dev/null +++ b/creator-keys-factory/src/lib.rs @@ -0,0 +1,169 @@ +#![no_std] +//! Minimal factory that deploys `creator-keys` contract instances and keeps a +//! registry of the deployed addresses. + +use soroban_sdk::{ + contract, contracterror, contractimpl, contracttype, symbol_short, Address, BytesN, Env, + Symbol, Vec, +}; + +pub const KEY_DEPLOYED_EVENT_NAME: Symbol = symbol_short!("key_depl"); + +#[contracterror] +#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord)] +#[repr(u32)] +pub enum FactoryError { + AlreadyInitialised = 1, + NotInitialised = 2, + Unauthorized = 3, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +#[contracttype] +pub struct KeyDeployedEvent { + pub key: Address, + pub creator: Address, +} + +#[derive(Clone)] +#[contracttype] +enum DataKey { + Admin, + WasmHash, + Registry, + Allowed(Address), +} + +#[contract] +pub struct CreatorKeysFactory; + +fn read_admin(env: &Env) -> Result { + env.storage() + .instance() + .get(&DataKey::Admin) + .ok_or(FactoryError::NotInitialised) +} + +#[contractimpl] +impl CreatorKeysFactory { + /// One-time setup: the admin and the uploaded `creator-keys` wasm hash to deploy. + pub fn initialise(env: Env, admin: Address, wasm_hash: BytesN<32>) -> Result<(), FactoryError> { + if env.storage().instance().has(&DataKey::Admin) { + return Err(FactoryError::AlreadyInitialised); + } + admin.require_auth(); + env.storage().instance().set(&DataKey::Admin, &admin); + env.storage().instance().set(&DataKey::WasmHash, &wasm_hash); + Ok(()) + } + + /// Admin allows or disallows a creator wallet to call `deploy_key`. + pub fn set_creator_allowed( + env: Env, + admin: Address, + creator: Address, + allowed: bool, + ) -> Result<(), FactoryError> { + admin.require_auth(); + if admin != read_admin(&env)? { + return Err(FactoryError::Unauthorized); + } + env.storage() + .instance() + .set(&DataKey::Allowed(creator), &allowed); + Ok(()) + } + + /// Deploys a new key contract for `creator` and records its address. + /// Only the admin or a whitelisted creator may call it. + pub fn deploy_key( + env: Env, + caller: Address, + creator: Address, + salt: BytesN<32>, + ) -> Result { + caller.require_auth(); + let is_allowed: bool = env + .storage() + .instance() + .get(&DataKey::Allowed(caller.clone())) + .unwrap_or(false); + if caller != read_admin(&env)? && !is_allowed { + return Err(FactoryError::Unauthorized); + } + let wasm_hash: BytesN<32> = env + .storage() + .instance() + .get(&DataKey::WasmHash) + .ok_or(FactoryError::NotInitialised)?; + + let key = env + .deployer() + .with_current_contract(salt) + .deploy_v2(wasm_hash, ()); + + let mut registry = Self::get_registry(env.clone()); + registry.push_back(key.clone()); + env.storage().instance().set(&DataKey::Registry, ®istry); + + env.events().publish( + (KEY_DEPLOYED_EVENT_NAME, creator.clone()), + KeyDeployedEvent { + key: key.clone(), + creator, + }, + ); + Ok(key) + } + + /// Read-only view: every key contract address deployed by this factory. + pub fn get_registry(env: Env) -> Vec
{ + env.storage() + .instance() + .get(&DataKey::Registry) + .unwrap_or_else(|| Vec::new(&env)) + } +} + +#[cfg(test)] +mod test { + use super::*; + use soroban_sdk::testutils::Address as _; + + fn setup(env: &Env) -> (CreatorKeysFactoryClient<'_>, Address) { + env.mock_all_auths(); + let client = CreatorKeysFactoryClient::new(env, &env.register(CreatorKeysFactory, ())); + let admin = Address::generate(env); + client.initialise(&admin, &BytesN::from_array(env, &[7; 32])); + (client, admin) + } + + #[test] + fn unauthorised_deploy_is_rejected() { + let env = Env::default(); + let (client, _admin) = setup(&env); + let stranger = Address::generate(&env); + let salt = BytesN::from_array(&env, &[1; 32]); + assert_eq!( + client.try_deploy_key(&stranger, &stranger, &salt), + Err(Ok(FactoryError::Unauthorized)) + ); + assert_eq!(client.get_registry().len(), 0); + } + + #[test] + fn only_admin_manages_allowed_creators_and_init_is_once() { + let env = Env::default(); + let (client, admin) = setup(&env); + let stranger = Address::generate(&env); + assert_eq!( + client.try_set_creator_allowed(&stranger, &stranger, &true), + Err(Ok(FactoryError::Unauthorized)) + ); + client.set_creator_allowed(&admin, &stranger, &true); + assert_eq!( + client.try_initialise(&admin, &BytesN::from_array(&env, &[7; 32])), + Err(Ok(FactoryError::AlreadyInitialised)) + ); + } +} diff --git a/creator-keys/src/events.rs b/creator-keys/src/events.rs index e39caf6f..217084fb 100644 --- a/creator-keys/src/events.rs +++ b/creator-keys/src/events.rs @@ -819,6 +819,11 @@ pub const WHITELIST_ENABLED_EVENT_NAME: Symbol = symbol_short!("wl_en"); pub const WHITELIST_DISABLED_EVENT_NAME: Symbol = symbol_short!("wl_dis"); pub const ADDRESS_WHITELISTED_EVENT_NAME: Symbol = symbol_short!("wl_add"); pub const ADDRESS_REMOVED_EVENT_NAME: Symbol = symbol_short!("wl_rem"); +pub const HOLDING_CAP_UPDATED_EVENT_NAME: Symbol = symbol_short!("hold_cap"); +pub const WHITELIST_UPDATED_EVENT_NAME: Symbol = symbol_short!("wl_upd"); +pub const REFERRAL_REGISTERED_EVENT_NAME: Symbol = symbol_short!("ref_reg"); +pub const REFERRAL_REWARD_ALLOCATED_EVENT_NAME: Symbol = symbol_short!("ref_rwd"); +pub const REFERRAL_REWARDS_CLAIMED_EVENT_NAME: Symbol = symbol_short!("ref_clm"); pub const KEYS_BURNED_EVENT_NAME: Symbol = symbol_short!("burned"); pub const SELF_FREEZE_APPLIED_EVENT_NAME: Symbol = symbol_short!("sf_add"); pub const SELF_FREEZE_LIFTED_EVENT_NAME: Symbol = symbol_short!("sf_del"); @@ -895,6 +900,66 @@ pub fn address_removed_topics(creator: &Address) -> (Symbol, Address) { (ADDRESS_REMOVED_EVENT_NAME, creator.clone()) } +/// Emitted when a creator changes their per-wallet holding cap. +#[derive(Clone, Debug, Eq, PartialEq)] +#[contracttype] +pub struct HoldingCapUpdatedEvent { + pub creator: Address, + pub old_cap: Option, + pub new_cap: u32, +} + +pub fn holding_cap_updated_topics(creator: &Address) -> (Symbol, Address) { + (HOLDING_CAP_UPDATED_EVENT_NAME, creator.clone()) +} + +/// Emitted on every early-access whitelist add (`allowed = true`) and remove. +#[derive(Clone, Debug, Eq, PartialEq)] +#[contracttype] +pub struct WhitelistUpdatedEvent { + pub creator: Address, + pub wallet: Address, + pub allowed: bool, +} + +pub fn whitelist_updated_topics(creator: &Address) -> (Symbol, Address) { + (WHITELIST_UPDATED_EVENT_NAME, creator.clone()) +} + +#[derive(Clone, Debug, Eq, PartialEq)] +#[contracttype] +pub struct ReferralRegisteredEvent { + pub referee: Address, + pub referrer: Address, +} + +pub fn referral_registered_topics() -> Symbol { + REFERRAL_REGISTERED_EVENT_NAME +} + +#[derive(Clone, Debug, Eq, PartialEq)] +#[contracttype] +pub struct ReferralRewardAllocatedEvent { + pub referee: Address, + pub referrer: Address, + pub amount: i128, +} + +pub fn referral_reward_allocated_topics() -> Symbol { + REFERRAL_REWARD_ALLOCATED_EVENT_NAME +} + +#[derive(Clone, Debug, Eq, PartialEq)] +#[contracttype] +pub struct ReferralRewardsClaimedEvent { + pub referrer: Address, + pub amount: i128, +} + +pub fn referral_rewards_claimed_topics() -> Symbol { + REFERRAL_REWARDS_CLAIMED_EVENT_NAME +} + #[derive(Clone, Debug, Eq, PartialEq)] #[contracttype] pub struct KeysBurnedEvent { diff --git a/creator-keys/src/lib.rs b/creator-keys/src/lib.rs index 8854cfd3..15fde93b 100644 --- a/creator-keys/src/lib.rs +++ b/creator-keys/src/lib.rs @@ -629,6 +629,16 @@ pub mod constants { pub fn holder_cap_bps(creator: &Address) -> DataKey { DataKey::HolderCapBps(creator.clone()) } + + pub const MAX_HOLDING_BOUND: DataKey = DataKey::MaxHoldingBound; + + pub fn referrer_of(referee: &Address) -> DataKey { + DataKey::Referrer(referee.clone()) + } + + pub fn referral_settled(referee: &Address) -> DataKey { + DataKey::ReferralSettled(referee.clone()) + } pub fn quorum_bps(creator: &Address) -> DataKey { DataKey::QuorumBps(creator.clone()) } @@ -1176,6 +1186,12 @@ pub enum DataKey { EarlyExitPenaltyBps(Address), /// Maximum buy quantity per transaction for a creator. MaxBuyQuantity(Address), + /// Admin-defined upper bound for a creator's per-wallet holding cap. + MaxHoldingBound, + /// Registered referrer for a referee wallet. + Referrer(Address), + /// Set once a referee's first referred trade has paid its referral reward. + ReferralSettled(Address), } #[derive(Clone, Debug, PartialEq)] @@ -2001,6 +2017,64 @@ fn assert_before_global_deadline(env: &Env) -> Result<(), ContractError> { Ok(()) } +/// Rejects a transfer that would leave the recipient above the creator's +/// per-wallet holding cap (the same cap `buy_key` enforces). +fn assert_within_holding_cap( + env: &Env, + creator: &Address, + new_balance: u32, +) -> Result<(), ContractError> { + if let Some(cap) = env + .storage() + .persistent() + .get::(&constants::storage::max_keys_per_wallet(creator)) + { + if new_balance > cap { + return Err(ContractError::WalletCapExceeded); + } + } + Ok(()) +} + +/// Resolves the referrer for a buy. An explicit referrer wins; otherwise the +/// buyer's registered referrer is used once, on their first referred trade. +/// Returns the referrer and whether it came from the one-time registration. +fn resolve_registered_referrer( + env: &Env, + buyer: &Address, + explicit: Option
, +) -> (Option
, bool) { + if explicit.is_some() { + return (explicit, false); + } + let settled_key = constants::storage::referral_settled(buyer); + if env.storage().persistent().has(&settled_key) { + return (None, false); + } + let registered: Option
= env + .storage() + .persistent() + .get(&constants::storage::referrer_of(buyer)); + if registered.is_some() { + env.storage().persistent().set(&settled_key, &true); + extend_key_ttl_to_full_window(env, &settled_key); + } + let is_registered = registered.is_some(); + (registered, is_registered) +} + +fn assert_creator_or_admin( + env: &Env, + caller: &Address, + creator: &Address, +) -> Result<(), ContractError> { + read_registered_creator_profile(env, creator)?; + if caller == creator { + return Ok(()); + } + assert_is_admin(env, caller) +} + fn assert_is_admin(env: &Env, caller: &Address) -> Result<(), ContractError> { let admin: Address = env .storage() @@ -3582,6 +3656,8 @@ impl CreatorKeysContract { } } + let (referrer, from_registration) = resolve_registered_referrer(&env, &buyer, referrer); + let base_price: i128 = env .storage() .persistent() @@ -3815,7 +3891,21 @@ impl CreatorKeysContract { // Split protocol fee between treasury and referrer only when a referrer is provided if let Some(referrer_addr) = referrer { - let referral_amount = protocol_fee / 2; + // Registered referrals pay the admin-configured share of the + // protocol fee (default 50%); explicit referrers keep the flat split. + let referral_amount = if from_registration { + let bps: u32 = env + .storage() + .persistent() + .get(&constants::storage::referral_fee_bps()) + .unwrap_or(5_000); + protocol_fee + .checked_mul(i128::from(bps)) + .ok_or(ContractError::Overflow)? + / i128::from(fee::BPS_MAX) + } else { + protocol_fee / 2 + }; let treasury_amount = protocol_fee - referral_amount; credit_treasury_balance(&env, treasury_amount)?; @@ -3831,6 +3921,16 @@ impl CreatorKeysContract { env.storage().persistent().set(&ref_key, &new_earnings); extend_key_ttl_to_full_window(&env, &ref_key); + if from_registration { + env.events().publish( + (events::referral_reward_allocated_topics(),), + events::ReferralRewardAllocatedEvent { + referee: buyer.clone(), + referrer: referrer_addr.clone(), + amount: referral_amount, + }, + ); + } env.events().publish( (events::referral_fee_paid_topics(),), events::ReferralFeePaidEvent { @@ -6635,6 +6735,7 @@ impl CreatorKeysContract { let new_to_balance = to_balance .checked_add(amount) .ok_or(ContractError::Overflow)?; + assert_within_holding_cap(&env, &creator, new_to_balance)?; env.storage() .persistent() .set(&to_balance_key, &new_to_balance); @@ -6759,6 +6860,7 @@ impl CreatorKeysContract { // Increment the recipient balance. let new_to_balance = to_balance.checked_add(qty).ok_or(ContractError::Overflow)?; + assert_within_holding_cap(&env, &creator, new_to_balance)?; env.storage() .persistent() .set(&to_balance_key, &new_to_balance); @@ -8054,6 +8156,184 @@ impl CreatorKeysContract { Ok(()) } + /// Admin sets the upper bound a creator may choose for their per-wallet + /// holding cap via [`Self::set_holding_cap`]. + pub fn set_holding_cap_bound( + env: Env, + admin: Address, + bound: u32, + ) -> Result<(), ContractError> { + admin.require_auth(); + assert_is_admin(&env, &admin)?; + if bound == 0 { + return Err(ContractError::InvalidHolderCap); + } + let key = constants::storage::MAX_HOLDING_BOUND; + env.storage().persistent().set(&key, &bound); + extend_key_ttl_to_full_window(&env, &key); + Ok(()) + } + + /// Creator updates the maximum number of keys one wallet may hold. + /// + /// Reuses the per-wallet cap already enforced by `buy_key`, and now also + /// by `transfer_keys` and `batch_transfer_keys`. `new_cap` must be greater + /// than zero and, when the admin has set a bound, not exceed it. + pub fn set_holding_cap(env: Env, creator: Address, new_cap: u32) -> Result<(), ContractError> { + creator.require_auth(); + read_registered_creator_profile(&env, &creator)?; + if new_cap == 0 { + return Err(ContractError::InvalidHolderCap); + } + if let Some(bound) = env + .storage() + .persistent() + .get::(&constants::storage::MAX_HOLDING_BOUND) + { + if new_cap > bound { + return Err(ContractError::InvalidHolderCap); + } + } + let key = constants::storage::max_keys_per_wallet(&creator); + let old_cap: Option = env.storage().persistent().get(&key); + env.storage().persistent().set(&key, &new_cap); + extend_key_ttl_to_full_window(&env, &key); + + env.events().publish( + events::holding_cap_updated_topics(&creator), + events::HoldingCapUpdatedEvent { + creator, + old_cap, + new_cap, + }, + ); + Ok(()) + } + + /// Read-only view: the per-wallet holding cap, or `None` when uncapped. + pub fn get_holding_cap(env: Env, creator: Address) -> Option { + env.storage() + .persistent() + .get(&constants::storage::max_keys_per_wallet(&creator)) + } + + /// Creator or admin toggles early-access mode. While on, only whitelisted + /// wallets may buy; turning it off opens trading to everyone. Shares the + /// flag used by `enable_whitelist` / `disable_whitelist`. + pub fn set_early_access_mode( + env: Env, + caller: Address, + creator: Address, + enabled: bool, + ) -> Result<(), ContractError> { + caller.require_auth(); + assert_creator_or_admin(&env, &caller, &creator)?; + let mode_key = constants::storage::whitelist_mode(&creator); + env.storage().persistent().set(&mode_key, &enabled); + extend_key_ttl_to_full_window(&env, &mode_key); + Ok(()) + } + + /// Creator or admin adds (`allowed = true`) or removes a wallet from the + /// early-access whitelist and emits `WhitelistUpdatedEvent`. + pub fn update_whitelist( + env: Env, + caller: Address, + creator: Address, + wallet: Address, + allowed: bool, + ) -> Result<(), ContractError> { + caller.require_auth(); + assert_creator_or_admin(&env, &caller, &creator)?; + let entry_key = constants::storage::whitelist_entry(&creator, &wallet); + env.storage().persistent().set(&entry_key, &allowed); + extend_key_ttl_to_full_window(&env, &entry_key); + + env.events().publish( + events::whitelist_updated_topics(&creator), + events::WhitelistUpdatedEvent { + creator, + wallet, + allowed, + }, + ); + Ok(()) + } + + /// Read-only view: whether `wallet` is on `creator`'s early-access whitelist. + pub fn get_wallet_whitelist_status(env: Env, creator: Address, wallet: Address) -> bool { + env.storage() + .persistent() + .get(&constants::storage::whitelist_entry(&creator, &wallet)) + .unwrap_or(false) + } + + /// Admin sets the share of the protocol fee (in bps) paid to a referrer on + /// a referee's first registered trade. Defaults to 5000 (50%) when unset. + pub fn set_referral_fee_bps(env: Env, admin: Address, bps: u32) -> Result<(), ContractError> { + admin.require_auth(); + assert_is_admin(&env, &admin)?; + if bps > fee::BPS_MAX { + return Err(ContractError::InvalidFeeConfig); + } + let key = constants::storage::referral_fee_bps(); + env.storage().persistent().set(&key, &bps); + extend_key_ttl_to_full_window(&env, &key); + Ok(()) + } + + /// Links `referee` to `referrer`. Allowed once per referee and only before + /// the referee's referral reward has been paid. The next buy by the referee + /// without an explicit referrer pays the reward; later buys pay nothing. + pub fn register_referral( + env: Env, + referee: Address, + referrer: Address, + ) -> Result<(), ContractError> { + referee.require_auth(); + if referee == referrer { + return Err(ContractError::InvalidReferrer); + } + let key = constants::storage::referrer_of(&referee); + if env.storage().persistent().has(&key) { + return Err(ContractError::AlreadyRegistered); + } + env.storage().persistent().set(&key, &referrer); + extend_key_ttl_to_full_window(&env, &key); + + env.events().publish( + (events::referral_registered_topics(),), + events::ReferralRegisteredEvent { referee, referrer }, + ); + Ok(()) + } + + /// Read-only view: the referrer registered for `referee`, if any. + pub fn get_referrer(env: Env, referee: Address) -> Option
{ + env.storage() + .persistent() + .get(&constants::storage::referrer_of(&referee)) + } + + /// Referrer withdraws all accumulated referral rewards. Returns the amount + /// claimed and resets the balance to zero. + pub fn claim_referral_rewards(env: Env, referrer: Address) -> Result { + referrer.require_auth(); + let key = constants::storage::referral_earnings(&referrer); + let amount: i128 = env.storage().persistent().get(&key).unwrap_or(0); + if amount <= 0 { + return Err(ContractError::NotPositiveAmount); + } + env.storage().persistent().set(&key, &0i128); + extend_key_ttl_to_full_window(&env, &key); + + env.events().publish( + (events::referral_rewards_claimed_topics(),), + events::ReferralRewardsClaimedEvent { referrer, amount }, + ); + Ok(amount) + } + pub fn burn( env: Env, caller: Address, diff --git a/creator-keys/src/test_new_features.rs b/creator-keys/src/test_new_features.rs index bd0fdeb7..ad93fc84 100644 --- a/creator-keys/src/test_new_features.rs +++ b/creator-keys/src/test_new_features.rs @@ -453,3 +453,132 @@ fn test_get_key_metadata_returns_none_for_uninitialised() { let result = client.get_key_metadata(&creator); assert_eq!(result, None); } + +#[test] +fn test_holding_cap_buy_transfer_and_update() { + let (env, client, admin, _treasury) = setup_test(); + let creator = Address::generate(&env); + register_creator(&env, &client, &creator); + client.set_circuit_breaker_threshold(&admin, &500u32); + + // Only a registered creator may set the cap, within the admin bound. + let attacker = Address::generate(&env); + assert_eq!( + client.try_set_holding_cap(&attacker, &1), + Err(Ok(ContractError::NotRegistered)) + ); + client.set_holding_cap_bound(&admin, &5); + assert_eq!( + client.try_set_holding_cap(&creator, &6), + Err(Ok(ContractError::InvalidHolderCap)) + ); + assert_eq!( + client.try_set_holding_cap(&creator, &0), + Err(Ok(ContractError::InvalidHolderCap)) + ); + assert_eq!(client.get_holding_cap(&creator), None); + client.set_holding_cap(&creator, &1); + assert_eq!(client.get_holding_cap(&creator), Some(1)); + + // At-cap buy is rejected. + let buyer = Address::generate(&env); + client.buy_key(&creator, &buyer, &1000i128, &None); + assert_eq!( + client.try_buy_key(&creator, &buyer, &1000i128, &None), + Err(Ok(ContractError::WalletCapExceeded)) + ); + + // Over-cap transfer is rejected. + let other = Address::generate(&env); + client.buy_key(&creator, &other, &1000i128, &None); + assert_eq!( + client.try_transfer_keys(&creator, &buyer, &other, &1), + Err(Ok(ContractError::WalletCapExceeded)) + ); + + // Raising the cap allows the transfer. + client.set_holding_cap(&creator, &2); + client.transfer_keys(&creator, &buyer, &other, &1); + assert_eq!(client.get_key_balance(&creator, &other), 2); +} + +#[test] +fn test_early_access_mode_whitelist_and_permissions() { + let (env, client, admin, _treasury) = setup_test(); + let creator = Address::generate(&env); + register_creator(&env, &client, &creator); + client.set_circuit_breaker_threshold(&admin, &500u32); + + let wallet = Address::generate(&env); + let attacker = Address::generate(&env); + assert_eq!( + client.try_set_early_access_mode(&attacker, &creator, &true), + Err(Ok(ContractError::Unauthorized)) + ); + assert_eq!( + client.try_update_whitelist(&attacker, &creator, &wallet, &true), + Err(Ok(ContractError::Unauthorized)) + ); + + // Non-whitelisted wallet is rejected during early access. + client.set_early_access_mode(&creator, &creator, &true); + assert_eq!( + client.try_buy_key(&creator, &wallet, &1000i128, &None), + Err(Ok(ContractError::NotWhitelisted)) + ); + + // Admin can whitelist; the wallet can then buy. + assert!(!client.get_wallet_whitelist_status(&creator, &wallet)); + client.update_whitelist(&admin, &creator, &wallet, &true); + assert!(client.get_wallet_whitelist_status(&creator, &wallet)); + assert_eq!(client.buy_key(&creator, &wallet, &1000i128, &None), 1); + + // Removal blocks again; disabling the mode opens trading to everyone. + client.update_whitelist(&creator, &creator, &wallet, &false); + assert!(!client.get_wallet_whitelist_status(&creator, &wallet)); + assert_eq!( + client.try_buy_key(&creator, &wallet, &1000i128, &None), + Err(Ok(ContractError::NotWhitelisted)) + ); + client.set_early_access_mode(&admin, &creator, &false); + assert_eq!(client.buy_key(&creator, &wallet, &1000i128, &None), 2); +} + +#[test] +fn test_registered_referral_first_trade_fee_and_claim() { + let (env, client, admin, _treasury) = setup_test(); + let creator = Address::generate(&env); + register_creator(&env, &client, &creator); + client.set_circuit_breaker_threshold(&admin, &500u32); + + let referee = Address::generate(&env); + let referrer = Address::generate(&env); + + assert_eq!( + client.try_register_referral(&referee, &referee), + Err(Ok(ContractError::InvalidReferrer)) + ); + client.set_referral_fee_bps(&admin, &2000); + client.register_referral(&referee, &referrer); + assert_eq!(client.get_referrer(&referee), Some(referrer.clone())); + assert_eq!( + client.try_register_referral(&referee, &referrer), + Err(Ok(ContractError::AlreadyRegistered)) + ); + + // First trade: price 100, protocol fee 10, referrer gets 20% = 2. + client.buy_key(&creator, &referee, &1000i128, &None); + assert_eq!(client.get_referral_earnings(&referrer), 2); + + // Subsequent trades pay no referral fee. + client.buy_key(&creator, &referee, &1000i128, &None); + assert_eq!(client.get_referral_earnings(&referrer), 2); + + // Claim returns the accumulated amount and resets the balance. + assert_eq!(client.claim_referral_rewards(&referrer), 2); + assert_eq!(client.get_referral_earnings(&referrer), 0); + assert_eq!( + client.try_claim_referral_rewards(&referrer), + Err(Ok(ContractError::NotPositiveAmount)) + ); +}