Skip to content

feat: Treasury Storage, Fee Collection & Allocation Logic #325

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Olowodarey merged 1 commit into from
Mar 24, 2026
Merged

feat: Treasury Storage, Fee Collection & Allocation Logic #325

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 60 additions & 16 deletions contracts/src/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,10 @@ const MAX_FEE_BPS: u32 = 10_000;
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Config {
pub admin: Address,
pub treasury: Address,
pub protocol_fee_bps: u32,
pub treasury: Address, // Treasury Address
pub deposit_fee_bps: u32,
pub withdrawal_fee_bps: u32,
pub performance_fee_bps: u32,
pub paused: bool,
}

Expand Down Expand Up @@ -58,7 +60,9 @@ pub fn initialize_config(
env: &Env,
admin: Address,
treasury: Address,
protocol_fee_bps: u32,
deposit_fee_bps: u32,
withdrawal_fee_bps: u32,
performance_fee_bps: u32,
) -> Result<(), SavingsError> {
// Prevent re-initialization
let already_init: bool = env
Expand All @@ -74,21 +78,35 @@ pub fn initialize_config(
require_admin(env, &admin)?;

// Validate fee bounds
if protocol_fee_bps > MAX_FEE_BPS {
if deposit_fee_bps > MAX_FEE_BPS
|| withdrawal_fee_bps > MAX_FEE_BPS
|| performance_fee_bps > MAX_FEE_BPS
{
return Err(SavingsError::InvalidFeeBps);
}

// Store config values
env.storage().instance().set(&DataKey::Treasury, &treasury);
env.storage()
.instance()
.set(&DataKey::ProtocolFeeBps, &protocol_fee_bps);
.set(&DataKey::TreasuryAddress, &treasury);
env.storage()
.instance()
.set(&DataKey::DepositFeeBps, &deposit_fee_bps);
env.storage()
.instance()
.set(&DataKey::WithdrawalFeeBps, &withdrawal_fee_bps);
env.storage()
.instance()
.set(&DataKey::PerformanceFeeBps, &performance_fee_bps);
env.storage()
.instance()
.set(&DataKey::ConfigInitialized, &true);

// Initialize the treasury struct with default zero values
crate::treasury::initialize_treasury(env);

env.events()
.publish((symbol_short!("cfg_init"),), protocol_fee_bps);
.publish((symbol_short!("cfg_init"),), performance_fee_bps);

Ok(())
}
Expand All @@ -114,13 +132,25 @@ pub fn get_config(env: &Env) -> Result<Config, SavingsError> {
let treasury: Address = env
.storage()
.instance()
.get(&DataKey::Treasury)
.get(&DataKey::TreasuryAddress)
.unwrap_or(admin.clone());

let protocol_fee_bps: u32 = env
let deposit_fee_bps: u32 = env
.storage()
.instance()
.get(&DataKey::DepositFeeBps)
.unwrap_or(0);

let withdrawal_fee_bps: u32 = env
.storage()
.instance()
.get(&DataKey::ProtocolFeeBps)
.get(&DataKey::WithdrawalFeeBps)
.unwrap_or(0);

let performance_fee_bps: u32 = env
.storage()
.instance()
.get(&DataKey::PerformanceFeeBps)
.unwrap_or(0);

let paused: bool = env
Expand All @@ -132,7 +162,9 @@ pub fn get_config(env: &Env) -> Result<Config, SavingsError> {
Ok(Config {
admin,
treasury,
protocol_fee_bps,
deposit_fee_bps,
withdrawal_fee_bps,
performance_fee_bps,
paused,
})
}
Expand All @@ -151,7 +183,7 @@ pub fn set_treasury(env: &Env, admin: Address, new_treasury: Address) -> Result<

env.storage()
.instance()
.set(&DataKey::Treasury, &new_treasury);
.set(&DataKey::TreasuryAddress, &new_treasury);

env.events()
.publish((symbol_short!("set_trs"),), new_treasury);
Expand All @@ -169,19 +201,31 @@ pub fn set_treasury(env: &Env, admin: Address, new_treasury: Address) -> Result<
/// # Errors
/// * `SavingsError::Unauthorized` - If caller is not the admin
/// * `SavingsError::InvalidFeeBps` - If fee exceeds 10000 bps
pub fn set_protocol_fee(env: &Env, admin: Address, new_fee_bps: u32) -> Result<(), SavingsError> {
pub fn set_fees(
env: &Env,
admin: Address,
deposit_fee: u32,
withdrawal_fee: u32,
performance_fee: u32,
) -> Result<(), SavingsError> {
require_admin(env, &admin)?;

if new_fee_bps > MAX_FEE_BPS {
if deposit_fee > MAX_FEE_BPS || withdrawal_fee > MAX_FEE_BPS || performance_fee > MAX_FEE_BPS {
return Err(SavingsError::InvalidFeeBps);
}

env.storage()
.instance()
.set(&DataKey::ProtocolFeeBps, &new_fee_bps);
.set(&DataKey::DepositFeeBps, &deposit_fee);
env.storage()
.instance()
.set(&DataKey::WithdrawalFeeBps, &withdrawal_fee);
env.storage()
.instance()
.set(&DataKey::PerformanceFeeBps, &performance_fee);

env.events()
.publish((symbol_short!("set_fee"),), new_fee_bps);
.publish((symbol_short!("set_fee"),), performance_fee);

Ok(())
}
Expand Down
74 changes: 40 additions & 34 deletions contracts/src/config_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,14 @@ fn test_initialize_config_succeeds() {
let treasury = Address::generate(&env);

env.mock_all_auths();
let result = client.try_initialize_config(&admin, &treasury, &100);
let result = client.try_initialize_config(&admin, &treasury, &100, &100, &100);
assert!(result.is_ok(), "initialize_config should succeed");

// Verify stored values
let config = client.get_config();
assert_eq!(config.admin, admin);
assert_eq!(config.treasury, treasury);
assert_eq!(config.protocol_fee_bps, 100);
assert_eq!(config.performance_fee_bps, 100);
assert!(!config.paused);
}

Expand All @@ -47,11 +47,11 @@ fn test_initialize_config_zero_fee() {
let treasury = Address::generate(&env);

env.mock_all_auths();
let result = client.try_initialize_config(&admin, &treasury, &0);
let result = client.try_initialize_config(&admin, &treasury, &0, &0, &0);
assert!(result.is_ok(), "zero fee should be valid");

let config = client.get_config();
assert_eq!(config.protocol_fee_bps, 0);
assert_eq!(config.performance_fee_bps, 0);
}

#[test]
Expand All @@ -60,11 +60,11 @@ fn test_initialize_config_max_fee() {
let treasury = Address::generate(&env);

env.mock_all_auths();
let result = client.try_initialize_config(&admin, &treasury, &10_000);
let result = client.try_initialize_config(&admin, &treasury, &10_000, &10_000, &10_000);
assert!(result.is_ok(), "max fee (10000 bps = 100%) should be valid");

let config = client.get_config();
assert_eq!(config.protocol_fee_bps, 10_000);
assert_eq!(config.performance_fee_bps, 10_000);
}

#[test]
Expand All @@ -74,14 +74,14 @@ fn test_reinitialize_config_fails() {

env.mock_all_auths();
assert!(client
.try_initialize_config(&admin, &treasury, &100)
.try_initialize_config(&admin, &treasury, &100, &100, &100)
.is_ok());

// Second initialization should fail
let treasury2 = Address::generate(&env);
assert_savings_error(
client
.try_initialize_config(&admin, &treasury2, &200)
.try_initialize_config(&admin, &treasury2, &200, &200, &200)
.unwrap_err(),
SavingsError::ConfigAlreadyInitialized,
);
Expand All @@ -95,7 +95,7 @@ fn test_initialize_config_fee_too_high() {
env.mock_all_auths();
assert_savings_error(
client
.try_initialize_config(&admin, &treasury, &10_001)
.try_initialize_config(&admin, &treasury, &10_001, &10_001, &10_001)
.unwrap_err(),
SavingsError::InvalidFeeBps,
);
Expand All @@ -110,7 +110,7 @@ fn test_non_admin_cannot_initialize_config() {
env.mock_all_auths();
assert_savings_error(
client
.try_initialize_config(&non_admin, &treasury, &100)
.try_initialize_config(&non_admin, &treasury, &100, &100, &100)
.unwrap_err(),
SavingsError::Unauthorized,
);
Expand All @@ -126,7 +126,7 @@ fn test_get_config_before_config_init() {
// Config should still be retrievable with defaults even without initialize_config
let config = client.get_config();
assert_eq!(config.admin, admin);
assert_eq!(config.protocol_fee_bps, 0); // default
assert_eq!(config.performance_fee_bps, 0); // default
assert!(!config.paused); // default
}

Expand All @@ -137,17 +137,17 @@ fn test_get_config_reflects_updates() {
let new_treasury = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

// Update treasury
client.set_treasury(&admin, &new_treasury);

// Update fee
client.set_protocol_fee(&admin, &500);
client.set_fees(&admin, &500, &500, &500);

let config = client.get_config();
assert_eq!(config.treasury, new_treasury);
assert_eq!(config.protocol_fee_bps, 500);
assert_eq!(config.performance_fee_bps, 500);
}

// ========== set_treasury Tests ==========
Expand All @@ -159,7 +159,7 @@ fn test_set_treasury_succeeds() {
let new_treasury = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

let result = client.try_set_treasury(&admin, &new_treasury);
assert!(result.is_ok(), "admin should be able to update treasury");
Expand All @@ -176,7 +176,7 @@ fn test_non_admin_cannot_set_treasury() {
let new_treasury = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

assert_savings_error(
client
Expand All @@ -194,13 +194,13 @@ fn test_set_protocol_fee_succeeds() {
let treasury = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

let result = client.try_set_protocol_fee(&admin, &500);
let result = client.try_set_fees(&admin, &500, &500, &500);
assert!(result.is_ok(), "admin should be able to update fee");

let config = client.get_config();
assert_eq!(config.protocol_fee_bps, 500);
assert_eq!(config.performance_fee_bps, 500);
}

#[test]
Expand All @@ -209,13 +209,13 @@ fn test_set_protocol_fee_to_zero() {
let treasury = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

let result = client.try_set_protocol_fee(&admin, &0);
let result = client.try_set_fees(&admin, &0, &0, &0);
assert!(result.is_ok(), "setting fee to 0 should work");

let config = client.get_config();
assert_eq!(config.protocol_fee_bps, 0);
assert_eq!(config.performance_fee_bps, 0);
}

#[test]
Expand All @@ -224,9 +224,9 @@ fn test_set_protocol_fee_to_max() {
let treasury = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

let result = client.try_set_protocol_fee(&admin, &10_000);
let result = client.try_set_fees(&admin, &10_000, &10_000, &10_000);
assert!(result.is_ok(), "setting fee to 10000 should work");
}

Expand All @@ -236,10 +236,12 @@ fn test_set_protocol_fee_exceeds_max() {
let treasury = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

assert_savings_error(
client.try_set_protocol_fee(&admin, &10_001).unwrap_err(),
client
.try_set_fees(&admin, &10_001, &10_001, &10_001)
.unwrap_err(),
SavingsError::InvalidFeeBps,
);
}
Expand All @@ -251,10 +253,12 @@ fn test_non_admin_cannot_set_protocol_fee() {
let non_admin = Address::generate(&env);

env.mock_all_auths();
client.initialize_config(&admin, &treasury, &100);
client.initialize_config(&admin, &treasury, &100, &100, &100);

assert_savings_error(
client.try_set_protocol_fee(&non_admin, &500).unwrap_err(),
client
.try_set_fees(&non_admin, &500, &500, &500)
.unwrap_err(),
SavingsError::Unauthorized,
);
}
Expand Down Expand Up @@ -414,11 +418,11 @@ fn test_full_config_lifecycle() {
env.mock_all_auths();

// 1. Initialize config
client.initialize_config(&admin, &treasury1, &250);
client.initialize_config(&admin, &treasury1, &250, &250, &250);

let config = client.get_config();
assert_eq!(config.treasury, treasury1);
assert_eq!(config.protocol_fee_bps, 250);
assert_eq!(config.performance_fee_bps, 250);
assert!(!config.paused);

// 2. Update treasury
Expand All @@ -427,17 +431,19 @@ fn test_full_config_lifecycle() {
assert_eq!(config.treasury, treasury2);

// 3. Update fee
client.set_protocol_fee(&admin, &500);
client.set_fees(&admin, &500, &500, &500);
let config = client.get_config();
assert_eq!(config.protocol_fee_bps, 500);
assert_eq!(config.deposit_fee_bps, 500);
assert_eq!(config.withdrawal_fee_bps, 500);
assert_eq!(config.performance_fee_bps, 500);

// 4. Pause
client.pause_contract(&admin);
assert!(client.get_config().paused);

// 5. Admin can still update config while paused
client.set_protocol_fee(&admin, &300);
assert_eq!(client.get_config().protocol_fee_bps, 300);
client.set_fees(&admin, &300, &300, &300);
assert_eq!(client.get_config().performance_fee_bps, 300);

// 6. Unpause
client.unpause_contract(&admin);
Expand Down
Loading
Loading