use crate::error::{Result, TokenError};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use tenzro_types::primitives::Address;
use tenzro_types::asset::AssetId;
use tenzro_storage::kv::{KvStore, CF_ACCOUNTS};
use std::sync::Arc;
use tracing::{debug, info, warn};
pub const TNZO_DECIMALS: u8 = 18;
const CIRCUIT_BREAKER_MAX_OUTFLOW_PERCENT: u128 = 10;
const CIRCUIT_BREAKER_WINDOW_SECS: u64 = 3600;
const CIRCUIT_BREAKER_COOLDOWN_SECS: u64 = 1800;
pub const ONE_TNZO: u128 = 1_000_000_000_000_000_000;
pub const MAX_SUPPLY: u128 = 1_000_000_000 * ONE_TNZO;
pub struct CircuitBreaker {
max_outflow_per_window: u128,
window_seconds: u64,
current_outflow: parking_lot::RwLock<u128>,
window_start: parking_lot::RwLock<u64>,
tripped: parking_lot::RwLock<bool>,
cooldown_seconds: u64,
tripped_at: parking_lot::RwLock<Option<u64>>,
}
impl CircuitBreaker {
pub fn new(max_outflow_per_window: u128, window_seconds: u64, cooldown_seconds: u64) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
Self {
max_outflow_per_window,
window_seconds,
current_outflow: parking_lot::RwLock::new(0),
window_start: parking_lot::RwLock::new(now),
tripped: parking_lot::RwLock::new(false),
cooldown_seconds,
tripped_at: parking_lot::RwLock::new(None),
}
}
pub fn default_for_tnzo() -> Self {
let max_outflow = MAX_SUPPLY / 100 * CIRCUIT_BREAKER_MAX_OUTFLOW_PERCENT;
Self::new(max_outflow, CIRCUIT_BREAKER_WINDOW_SECS, CIRCUIT_BREAKER_COOLDOWN_SECS)
}
fn now_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
pub fn is_tripped(&self) -> bool {
let tripped = *self.tripped.read();
if !tripped {
return false;
}
if let Some(trip_time) = *self.tripped_at.read() {
let now = Self::now_secs();
if now >= trip_time + self.cooldown_seconds {
self.reset_internal();
return false;
}
}
true
}
pub fn check_outflow(&self, amount: u128) -> Result<()> {
if self.is_tripped() {
let remaining = self.tripped_at.read()
.map(|t| {
let elapsed = Self::now_secs().saturating_sub(t);
self.cooldown_seconds.saturating_sub(elapsed)
})
.unwrap_or(self.cooldown_seconds);
return Err(TokenError::Unauthorized {
reason: format!(
"Circuit breaker tripped: outflow limit exceeded. \
Auto-reset in {} seconds.",
remaining
),
});
}
let now = Self::now_secs();
let mut window_start = self.window_start.write();
let mut current_outflow = self.current_outflow.write();
if now >= *window_start + self.window_seconds {
*window_start = now;
*current_outflow = 0;
}
let new_outflow = current_outflow
.checked_add(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "circuit breaker outflow check".to_string(),
})?;
if new_outflow > self.max_outflow_per_window {
drop(window_start);
drop(current_outflow);
self.trip();
return Err(TokenError::Unauthorized {
reason: format!(
"Circuit breaker tripped: transfer of {} would exceed \
window limit of {} ({} already used). \
Cooldown: {} seconds.",
amount,
self.max_outflow_per_window,
new_outflow - amount,
self.cooldown_seconds,
),
});
}
Ok(())
}
pub fn record_outflow(&self, amount: u128) {
let mut current_outflow = self.current_outflow.write();
*current_outflow = current_outflow.saturating_add(amount);
}
pub fn reset(&self) {
self.reset_internal();
info!("Circuit breaker manually reset");
}
fn reset_internal(&self) {
*self.tripped.write() = false;
*self.tripped_at.write() = None;
*self.current_outflow.write() = 0;
*self.window_start.write() = Self::now_secs();
}
fn trip(&self) {
*self.tripped.write() = true;
*self.tripped_at.write() = Some(Self::now_secs());
warn!(
"Circuit breaker TRIPPED: outflow limit of {} exceeded in {}-second window. \
Transfers paused for {} seconds.",
self.max_outflow_per_window,
self.window_seconds,
self.cooldown_seconds,
);
}
pub fn current_outflow(&self) -> u128 {
*self.current_outflow.read()
}
pub fn max_outflow(&self) -> u128 {
self.max_outflow_per_window
}
}
impl std::fmt::Debug for CircuitBreaker {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CircuitBreaker")
.field("max_outflow_per_window", &self.max_outflow_per_window)
.field("window_seconds", &self.window_seconds)
.field("cooldown_seconds", &self.cooldown_seconds)
.field("tripped", &*self.tripped.read())
.field("current_outflow", &*self.current_outflow.read())
.finish()
}
}
pub trait StorageBackend: Send + Sync {
fn get_balance(&self, address: &Address) -> Result<Option<u128>>;
fn set_balance(&self, address: &Address, balance: u128) -> Result<()>;
fn get_total_supply(&self) -> Result<u128>;
fn set_total_supply(&self, supply: u128) -> Result<()>;
}
pub struct RocksDbBackend {
store: Arc<dyn KvStore>,
}
impl RocksDbBackend {
pub fn new(store: Arc<dyn KvStore>) -> Self {
Self { store }
}
fn balance_key(address: &Address) -> Vec<u8> {
let mut key = b"balance:".to_vec();
key.extend_from_slice(address.as_bytes());
key
}
fn supply_key() -> Vec<u8> {
b"total_supply".to_vec()
}
}
impl StorageBackend for RocksDbBackend {
fn get_balance(&self, address: &Address) -> Result<Option<u128>> {
let key = Self::balance_key(address);
match self.store.get(CF_ACCOUNTS, &key)? {
Some(bytes) => {
if bytes.len() == 16 {
let array: [u8; 16] = bytes.try_into().unwrap();
Ok(Some(u128::from_le_bytes(array)))
} else {
warn!("Invalid balance bytes length for {}: {}", address, bytes.len());
Ok(None)
}
}
None => Ok(None),
}
}
fn set_balance(&self, address: &Address, balance: u128) -> Result<()> {
let key = Self::balance_key(address);
let value = balance.to_le_bytes();
self.store.put(CF_ACCOUNTS, &key, &value)?;
Ok(())
}
fn get_total_supply(&self) -> Result<u128> {
let key = Self::supply_key();
match self.store.get(CF_ACCOUNTS, &key)? {
Some(bytes) => {
if bytes.len() == 16 {
let array: [u8; 16] = bytes.try_into().unwrap();
Ok(u128::from_le_bytes(array))
} else {
warn!("Invalid supply bytes length: {}", bytes.len());
Ok(0)
}
}
None => Ok(0),
}
}
fn set_total_supply(&self, supply: u128) -> Result<()> {
let key = Self::supply_key();
let value = supply.to_le_bytes();
self.store.put(CF_ACCOUNTS, &key, &value)?;
Ok(())
}
}
pub struct TnzoToken {
pub asset_id: AssetId,
balances: DashMap<Address, u128>,
total_supply: parking_lot::RwLock<u128>,
total_burned: parking_lot::RwLock<u128>,
treasury_address: parking_lot::RwLock<Option<Address>>,
storage: Option<Arc<dyn StorageBackend>>,
circuit_breaker: CircuitBreaker,
}
impl std::fmt::Debug for TnzoToken {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TnzoToken")
.field("asset_id", &self.asset_id)
.field("total_supply", &self.total_supply)
.field("total_burned", &self.total_burned)
.field("treasury_address", &self.treasury_address)
.field("storage", &self.storage.as_ref().map(|_| "Some(...)"))
.field("circuit_breaker", &self.circuit_breaker)
.finish()
}
}
impl TnzoToken {
pub fn new() -> Self {
Self {
asset_id: AssetId::tnzo(),
balances: DashMap::new(),
total_supply: parking_lot::RwLock::new(0),
total_burned: parking_lot::RwLock::new(0),
treasury_address: parking_lot::RwLock::new(None),
storage: None,
circuit_breaker: CircuitBreaker::default_for_tnzo(),
}
}
pub fn with_storage(storage: Arc<dyn StorageBackend>) -> Result<Self> {
let total_supply = storage.get_total_supply()?;
Ok(Self {
asset_id: AssetId::tnzo(),
balances: DashMap::new(),
total_supply: parking_lot::RwLock::new(total_supply),
total_burned: parking_lot::RwLock::new(0),
treasury_address: parking_lot::RwLock::new(None),
storage: Some(storage),
circuit_breaker: CircuitBreaker::default_for_tnzo(),
})
}
pub fn set_treasury_address(&self, address: Address) {
*self.treasury_address.write() = Some(address);
info!("Treasury address set to: {}", address);
}
pub fn treasury_address_ref(&self) -> Option<Address> {
*self.treasury_address.read()
}
pub fn balance_of(&self, address: &Address) -> u128 {
if let Some(storage) = &self.storage {
return match storage.get_balance(address) {
Ok(Some(balance)) => balance,
Ok(None) => 0,
Err(e) => {
warn!("Failed to load balance from storage: {}", e);
0
}
};
}
self.balances.get(address).map(|v| *v).unwrap_or(0)
}
fn persist_balance(&self, address: &Address, balance: u128) -> Result<()> {
if let Some(storage) = &self.storage {
storage.set_balance(address, balance)?;
}
Ok(())
}
fn persist_supply(&self, supply: u128) -> Result<()> {
if let Some(storage) = &self.storage {
storage.set_total_supply(supply)?;
}
Ok(())
}
pub fn transfer(&self, from: &Address, to: &Address, amount: u128) -> Result<()> {
if amount == 0 {
return Err(TokenError::InvalidAmount("Amount must be greater than zero".to_string()));
}
self.circuit_breaker.check_outflow(amount)?;
let from_balance = self.balance_of(from);
if from_balance < amount {
return Err(TokenError::InsufficientBalance {
required: amount,
available: from_balance,
});
}
let new_from_balance = from_balance.checked_sub(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "transfer subtraction".to_string(),
})?;
let to_balance = self.balance_of(to);
let new_to_balance = to_balance.checked_add(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "transfer addition".to_string(),
})?;
self.balances.insert(*from, new_from_balance);
self.balances.insert(*to, new_to_balance);
self.persist_balance(from, new_from_balance)?;
self.persist_balance(to, new_to_balance)?;
self.circuit_breaker.record_outflow(amount);
debug!("Transferred {} TNZO from {} to {}", amount, from, to);
Ok(())
}
pub fn mint(&self, to: &Address, amount: u128, caller: &Address) -> Result<()> {
let treasury = self.treasury_address.read();
if treasury.as_ref() != Some(caller) {
return Err(TokenError::Unauthorized {
reason: "Only treasury can mint tokens".to_string(),
});
}
if amount == 0 {
return Err(TokenError::InvalidAmount("Amount must be greater than zero".to_string()));
}
let current_supply = *self.total_supply.read();
let new_supply = current_supply.checked_add(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "mint supply increase".to_string(),
})?;
if new_supply > MAX_SUPPLY {
return Err(TokenError::InvalidAmount(
format!("Minting would exceed max supply: {} > {}", new_supply, MAX_SUPPLY)
));
}
let current_balance = self.balance_of(to);
let new_balance = current_balance.checked_add(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "mint balance increase".to_string(),
})?;
self.balances.insert(*to, new_balance);
*self.total_supply.write() = new_supply;
self.persist_balance(to, new_balance)?;
self.persist_supply(new_supply)?;
info!("Minted {} TNZO to {}", amount, to);
Ok(())
}
pub fn burn(&self, from: &Address, amount: u128) -> Result<()> {
if amount == 0 {
return Err(TokenError::InvalidAmount("Amount must be greater than zero".to_string()));
}
let balance = self.balance_of(from);
if balance < amount {
return Err(TokenError::InsufficientBalance {
required: amount,
available: balance,
});
}
let new_balance = balance.checked_sub(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "burn balance decrease".to_string(),
})?;
self.balances.insert(*from, new_balance);
let mut supply = self.total_supply.write();
*supply = supply.checked_sub(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "burn supply decrease".to_string(),
})?;
let new_supply = *supply;
drop(supply);
let mut burned = self.total_burned.write();
*burned = burned.checked_add(amount)
.ok_or_else(|| TokenError::ArithmeticOverflow {
operation: "burn total increase".to_string(),
})?;
self.persist_balance(from, new_balance)?;
self.persist_supply(new_supply)?;
info!("Burned {} TNZO from {}", amount, from);
Ok(())
}
pub fn total_supply(&self) -> u128 {
*self.total_supply.read()
}
pub fn circulating_supply(&self) -> u128 {
self.total_supply()
}
pub fn total_burned(&self) -> u128 {
*self.total_burned.read()
}
pub fn stats(&self) -> TokenStats {
TokenStats {
total_supply: self.total_supply(),
circulating_supply: self.circulating_supply(),
total_burned: self.total_burned(),
total_accounts: self.balances.len() as u64,
}
}
pub fn circuit_breaker(&self) -> &CircuitBreaker {
&self.circuit_breaker
}
pub fn get_all_balances(&self) -> Vec<(Address, u128)> {
self.balances
.iter()
.map(|entry| (*entry.key(), *entry.value()))
.collect()
}
}
impl Default for TnzoToken {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenStats {
pub total_supply: u128,
pub circulating_supply: u128,
pub total_burned: u128,
pub total_accounts: u64,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_token_creation() {
let token = TnzoToken::new();
assert_eq!(token.total_supply(), 0);
assert_eq!(token.total_burned(), 0);
}
#[test]
fn test_transfer() {
let token = TnzoToken::new();
let from = Address::new([1u8; 32]);
let to = Address::new([2u8; 32]);
token.balances.insert(from, 1000 * ONE_TNZO);
token.transfer(&from, &to, 100 * ONE_TNZO).unwrap();
assert_eq!(token.balance_of(&from), 900 * ONE_TNZO);
assert_eq!(token.balance_of(&to), 100 * ONE_TNZO);
}
#[test]
fn test_mint() {
let token = TnzoToken::new();
let treasury = Address::new([1u8; 32]);
let recipient = Address::new([2u8; 32]);
token.set_treasury_address(treasury);
token.mint(&recipient, 1000 * ONE_TNZO, &treasury).unwrap();
assert_eq!(token.balance_of(&recipient), 1000 * ONE_TNZO);
assert_eq!(token.total_supply(), 1000 * ONE_TNZO);
}
#[test]
fn test_burn() {
let token = TnzoToken::new();
let address = Address::new([1u8; 32]);
token.balances.insert(address, 1000 * ONE_TNZO);
*token.total_supply.write() = 1000 * ONE_TNZO;
token.burn(&address, 100 * ONE_TNZO).unwrap();
assert_eq!(token.balance_of(&address), 900 * ONE_TNZO);
assert_eq!(token.total_supply(), 900 * ONE_TNZO);
assert_eq!(token.total_burned(), 100 * ONE_TNZO);
}
#[test]
fn test_circuit_breaker_allows_normal_transfer() {
let cb = CircuitBreaker::new(1000 * ONE_TNZO, 3600, 1800);
assert!(!cb.is_tripped());
assert!(cb.check_outflow(500 * ONE_TNZO).is_ok());
cb.record_outflow(500 * ONE_TNZO);
assert_eq!(cb.current_outflow(), 500 * ONE_TNZO);
}
#[test]
fn test_circuit_breaker_trips_on_excess() {
let cb = CircuitBreaker::new(1000 * ONE_TNZO, 3600, 1800);
cb.record_outflow(900 * ONE_TNZO);
let result = cb.check_outflow(200 * ONE_TNZO);
assert!(result.is_err());
assert!(cb.is_tripped());
}
#[test]
fn test_circuit_breaker_blocks_after_trip() {
let cb = CircuitBreaker::new(1000 * ONE_TNZO, 3600, 1800);
cb.record_outflow(900 * ONE_TNZO);
let _ = cb.check_outflow(200 * ONE_TNZO); assert!(cb.is_tripped());
let result = cb.check_outflow(1);
assert!(result.is_err());
}
#[test]
fn test_circuit_breaker_manual_reset() {
let cb = CircuitBreaker::new(1000 * ONE_TNZO, 3600, 1800);
cb.record_outflow(900 * ONE_TNZO);
let _ = cb.check_outflow(200 * ONE_TNZO); assert!(cb.is_tripped());
cb.reset();
assert!(!cb.is_tripped());
assert_eq!(cb.current_outflow(), 0);
assert!(cb.check_outflow(500 * ONE_TNZO).is_ok());
}
#[test]
fn test_circuit_breaker_wired_into_transfer() {
let token = TnzoToken::new();
let from = Address::new([1u8; 32]);
let to = Address::new([2u8; 32]);
token.balances.insert(from, MAX_SUPPLY);
*token.total_supply.write() = MAX_SUPPLY;
let limit = MAX_SUPPLY / 100 * CIRCUIT_BREAKER_MAX_OUTFLOW_PERCENT;
let small_amount = limit / 2;
assert!(token.transfer(&from, &to, small_amount).is_ok());
let over_amount = limit;
let result = token.transfer(&from, &to, over_amount);
assert!(result.is_err());
assert!(token.circuit_breaker().is_tripped());
token.circuit_breaker().reset();
assert!(token.transfer(&from, &to, small_amount).is_ok());
}
#[test]
fn test_treasury_address_ref() {
let token = TnzoToken::new();
assert!(token.treasury_address_ref().is_none());
let treasury = Address::new([0xFF; 32]);
token.set_treasury_address(treasury);
assert_eq!(token.treasury_address_ref(), Some(treasury));
}
#[test]
fn test_balance_of_reflects_external_storage_writes() {
use std::sync::Mutex as StdMutex;
struct MockBackend {
balances: StdMutex<std::collections::HashMap<Address, u128>>,
supply: StdMutex<u128>,
}
impl StorageBackend for MockBackend {
fn get_balance(&self, address: &Address) -> Result<Option<u128>> {
Ok(self.balances.lock().unwrap().get(address).copied())
}
fn set_balance(&self, address: &Address, balance: u128) -> Result<()> {
self.balances.lock().unwrap().insert(*address, balance);
Ok(())
}
fn get_total_supply(&self) -> Result<u128> {
Ok(*self.supply.lock().unwrap())
}
fn set_total_supply(&self, s: u128) -> Result<()> {
*self.supply.lock().unwrap() = s;
Ok(())
}
}
let backend = Arc::new(MockBackend {
balances: StdMutex::new(Default::default()),
supply: StdMutex::new(0),
});
let token = TnzoToken::with_storage(backend.clone() as Arc<dyn StorageBackend>).unwrap();
let recipient = Address::new([0xAB; 32]);
assert_eq!(token.balance_of(&recipient), 0);
backend.balances.lock().unwrap().insert(recipient, 5_000 * ONE_TNZO);
assert_eq!(token.balance_of(&recipient), 5_000 * ONE_TNZO);
backend.balances.lock().unwrap().insert(recipient, 1_000 * ONE_TNZO);
assert_eq!(token.balance_of(&recipient), 1_000 * ONE_TNZO);
}
}