Skip to content

feat: Add mainnet security preparation features #110

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

Open
damianosakwe wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add mainnet security preparation features #110

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
170 changes: 170 additions & 0 deletions contracts/contracts/escrow/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ pub enum MilestoneStatus {
Released = 4, // Funds successfully transferred to freelancer
Refunded = 5, // Funds returned to client
Disputed = 6, // Under dispute, waiting for arbiter resolution
EmergencyWithdrawn = 7, // Funds withdrawn via emergency mechanism
}

#[derive(Clone, Debug, PartialEq, Eq)]
Expand All @@ -33,6 +34,11 @@ pub enum DataKey {
Token,
Milestones,
IsFunded,
Admin,
Paused,
EmergencyWithdrawTimelock,
MaxContractValue,
DisputeTimelock,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
Expand All @@ -46,6 +52,11 @@ pub enum Error {
InvalidMilestoneStatus = 6,
Unauthorized = 7,
ZeroAmount = 8,
ContractPaused = 9,
EmergencyWithdrawNotReady = 10,
MaxValueExceeded = 11,
DisputeTimelockNotExpired = 12,
InvalidAddress = 13,
}

#[contract]
Expand All @@ -61,6 +72,9 @@ impl EscrowContract {
arbiter: Address,
token: Address,
milestones: Vec<Milestone>,
admin: Address,
max_contract_value: i128,
dispute_timelock: u64,
) -> Result<(), Error> {
if env.storage().instance().has(&DataKey::Client) {
return Err(Error::AlreadyInitialized);
Expand All @@ -70,6 +84,11 @@ impl EscrowContract {
return Err(Error::MilestoneNotFound);
}

// Validate addresses are not zero
if client == Address::from(&env) || freelancer == Address::from(&env) || arbiter == Address::from(&env) || token == Address::from(&env) {
return Err(Error::InvalidAddress);
}

// Validate all milestone amounts are greater than zero
for i in 0..milestones.len() {
let milestone = milestones.get(i).unwrap();
Expand All @@ -78,18 +97,37 @@ impl EscrowContract {
}
}

// Validate max_contract_value is positive
if max_contract_value <= 0 {
return Err(Error::ZeroAmount);
}

env.storage().instance().set(&DataKey::Client, &client);
env.storage().instance().set(&DataKey::Freelancer, &freelancer);
env.storage().instance().set(&DataKey::Arbiter, &arbiter);
env.storage().instance().set(&DataKey::Token, &token);
env.storage().instance().set(&DataKey::Milestones, &milestones);
env.storage().instance().set(&DataKey::IsFunded, &false);
env.storage().instance().set(&DataKey::Admin, &admin);
env.storage().instance().set(&DataKey::Paused, &false);
env.storage().instance().set(&DataKey::MaxContractValue, &max_contract_value);
env.storage().instance().set(&DataKey::DisputeTimelock, &dispute_timelock);

Ok(())
}

/// Helper function to check if contract is paused
fn is_paused(env: &Env) -> bool {
env.storage().instance().get(&DataKey::Paused).unwrap_or(false)
}

/// The client locks the total funds for all milestones into the contract.
pub fn fund(env: Env) -> Result<(), Error> {
// Check if contract is paused
if Self::is_paused(&env) {
return Err(Error::ContractPaused);
}

let client: Address = env.storage().instance().get(&DataKey::Client).ok_or(Error::NotInitialized)?;
client.require_auth();

Expand All @@ -110,6 +148,12 @@ impl EscrowContract {
return Err(Error::ZeroAmount);
}

// Check against max contract value
let max_value: i128 = env.storage().instance().get(&DataKey::MaxContractValue).ok_or(Error::NotInitialized)?;
if total_amount > max_value {
return Err(Error::MaxValueExceeded);
}

// Transfer payment tokens from the client to this contract
let token_address: Address = env.storage().instance().get(&DataKey::Token).ok_or(Error::NotInitialized)?;
let token_client = token::Client::new(&env, &token_address);
Expand All @@ -133,6 +177,11 @@ impl EscrowContract {

/// Freelancer submits milestone progress for client review.
pub fn submit_milestone(env: Env, milestone_id: u32) -> Result<(), Error> {
// Check if contract is paused
if Self::is_paused(&env) {
return Err(Error::ContractPaused);
}

let freelancer: Address = env.storage().instance().get(&DataKey::Freelancer).ok_or(Error::NotInitialized)?;
freelancer.require_auth();

Expand Down Expand Up @@ -162,6 +211,11 @@ impl EscrowContract {

/// Client approves milestone completion.
pub fn approve(env: Env, milestone_id: u32) -> Result<(), Error> {
// Check if contract is paused
if Self::is_paused(&env) {
return Err(Error::ContractPaused);
}

let client: Address = env.storage().instance().get(&DataKey::Client).ok_or(Error::NotInitialized)?;
client.require_auth();

Expand Down Expand Up @@ -192,6 +246,11 @@ impl EscrowContract {
/// Transfers funds of an approved milestone to the freelancer.
/// Can be triggered by either client or freelancer.
pub fn release(env: Env, milestone_id: u32, caller: Address) -> Result<(), Error> {
// Check if contract is paused
if Self::is_paused(&env) {
return Err(Error::ContractPaused);
}

caller.require_auth();

let client: Address = env.storage().instance().get(&DataKey::Client).ok_or(Error::NotInitialized)?;
Expand Down Expand Up @@ -234,6 +293,11 @@ impl EscrowContract {

/// Freelancer voluntarily refunds locked funds back to the client.
pub fn refund(env: Env, milestone_id: u32, caller: Address) -> Result<(), Error> {
// Check if contract is paused
if Self::is_paused(&env) {
return Err(Error::ContractPaused);
}

caller.require_auth();

let freelancer: Address = env.storage().instance().get(&DataKey::Freelancer).ok_or(Error::NotInitialized)?;
Expand Down Expand Up @@ -277,6 +341,11 @@ impl EscrowContract {
/// Puts a milestone into dispute, halting regular flow and delegating resolution to the arbiter.
/// Can be raised by client or freelancer.
pub fn dispute(env: Env, milestone_id: u32, caller: Address) -> Result<(), Error> {
// Check if contract is paused
if Self::is_paused(&env) {
return Err(Error::ContractPaused);
}

caller.require_auth();

let client: Address = env.storage().instance().get(&DataKey::Client).ok_or(Error::NotInitialized)?;
Expand All @@ -298,6 +367,8 @@ impl EscrowContract {
return Err(Error::InvalidMilestoneStatus);
}
milestone.status = MilestoneStatus::Disputed;
// Store the dispute timestamp for timelock check
env.storage().instance().set(&DataKey::DisputeTimelock, &env.ledger().timestamp());
}
updated_milestones.push_back(milestone);
}
Expand Down Expand Up @@ -356,6 +427,89 @@ impl EscrowContract {
Ok(())
}

// --- Admin Functions ---

/// Pause the contract - only callable by admin
pub fn pause(env: Env) -> Result<(), Error> {
let admin: Address = env.storage().instance().get(&DataKey::Admin).ok_or(Error::NotInitialized)?;
admin.require_auth();

env.storage().instance().set(&DataKey::Paused, &true);
Ok(())
}

/// Unpause the contract - only callable by admin
pub fn unpause(env: Env) -> Result<(), Error> {
let admin: Address = env.storage().instance().get(&DataKey::Admin).ok_or(Error::NotInitialized)?;
admin.require_auth();

env.storage().instance().set(&DataKey::Paused, &false);
Ok(())
}

/// Initiate emergency withdrawal - sets timelock, only callable by admin
pub fn initiate_emergency_withdraw(env: Env) -> Result<(), Error> {
let admin: Address = env.storage().instance().get(&DataKey::Admin).ok_or(Error::NotInitialized)?;
admin.require_auth();

// Set timelock to 48 hours from now (in ledger timestamp)
let timelock = env.ledger().timestamp() + 172800; // 48 hours in seconds
env.storage().instance().set(&DataKey::EmergencyWithdrawTimelock, &timelock);

Ok(())
}

/// Execute emergency withdrawal - only callable by admin after timelock expires
pub fn emergency_withdraw(env: Env, recipient: Address) -> Result<(), Error> {
let admin: Address = env.storage().instance().get(&DataKey::Admin).ok_or(Error::NotInitialized)?;
admin.require_auth();

let timelock: u64 = env.storage().instance().get(&DataKey::EmergencyWithdrawTimelock).ok_or(Error::EmergencyWithdrawNotReady)?;
let current_timestamp = env.ledger().timestamp();

if current_timestamp < timelock {
return Err(Error::EmergencyWithdrawNotReady);
}

// Calculate total remaining balance
let milestones: Vec<Milestone> = env.storage().instance().get(&DataKey::Milestones).ok_or(Error::NotInitialized)?;
let mut total_amount: i128 = 0;

for i in 0..milestones.len() {
let milestone = milestones.get(i).unwrap();
// Only withdraw from milestones that haven't been released or refunded
if milestone.status == MilestoneStatus::Funded || milestone.status == MilestoneStatus::Submitted || milestone.status == MilestoneStatus::Approved {
total_amount += milestone.amount;
}
}

if total_amount <= 0 {
return Err(Error::ZeroAmount);
}

// Transfer all remaining funds to recipient
let token_address: Address = env.storage().instance().get(&DataKey::Token).ok_or(Error::NotInitialized)?;
let token_client = token::Client::new(&env, &token_address);
token_client.transfer(&env.current_contract_address(), &recipient, &total_amount);

// Update all affected milestones to EmergencyWithdrawn status
let mut updated_milestones = Vec::new(&env);
for i in 0..milestones.len() {
let mut milestone = milestones.get(i).unwrap();
if milestone.status == MilestoneStatus::Funded || milestone.status == MilestoneStatus::Submitted || milestone.status == MilestoneStatus::Approved {
milestone.status = MilestoneStatus::EmergencyWithdrawn;
}
updated_milestones.push_back(milestone);
}

env.storage().instance().set(&DataKey::Milestones, &updated_milestones);

// Clear the timelock
env.storage().instance().set(&DataKey::EmergencyWithdrawTimelock, &0u64);

Ok(())
}

// --- State Getters ---

pub fn get_client(env: Env) -> Result<Address, Error> {
Expand All @@ -381,6 +535,22 @@ impl EscrowContract {
pub fn is_funded(env: Env) -> bool {
env.storage().instance().get(&DataKey::IsFunded).unwrap_or(false)
}

pub fn get_admin(env: Env) -> Result<Address, Error> {
env.storage().instance().get(&DataKey::Admin).ok_or(Error::NotInitialized)
}

pub fn is_paused(env: Env) -> bool {
env.storage().instance().get(&DataKey::Paused).unwrap_or(false)
}

pub fn get_max_contract_value(env: Env) -> Result<i128, Error> {
env.storage().instance().get(&DataKey::MaxContractValue).ok_or(Error::NotInitialized)
}

pub fn get_emergency_withdraw_timelock(env: Env) -> Result<u64, Error> {
env.storage().instance().get(&DataKey::EmergencyWithdrawTimelock).ok_or(Error::NotInitialized)
}
}

mod test;
Loading
Loading