feat: add delegation with storage griefing guard to governance

contracts/predictify-hybrid/src/governance.rs

@@ -23,12 +23,22 @@ pub struct GovernanceProposal {
 enum StorageKey {
     Proposal(Symbol),
     ProposalList,          // Vec<Symbol>
-    Vote(Symbol, Address), // proposal id + voter -> u8 (0 none, 1 for, 2 against)
+    Vote(Symbol, Address), // proposal id + voter -> i32 (0 none, 1 for, 2 against)
     VotingPeriod,          // u64
     QuorumVotes,           // u128 minimum FOR votes required
     Admin,                 // Address
+    /// Maps delegator -> delegate (Address). At most one per delegator (griefing guard).
+    Delegate(Address),
+    /// Tracks how many delegators are currently delegating to a given address.
+    /// Capped at MAX_INCOMING_DELEGATIONS to bound storage.
+    DelegateCount(Address),
 }
 
+/// Maximum number of delegators that may point to a single delegate address.
+/// Limits griefing: prevents an attacker from forcing the contract to walk an
+/// unbounded list when tallying delegated votes.
+const MAX_INCOMING_DELEGATIONS: u32 = 50;
+
 /// Simple errors for the contract
 #[contracttype]
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -44,6 +54,16 @@ pub enum GovernanceError {
     NotAdmin,
     NotInitialized,
     InvalidParams,
+    /// Caller tried to delegate to themselves.
+    SelfDelegation,
+    /// Caller already has an active delegation; must unset first.
+    DelegationAlreadySet,
+    /// The target delegate has reached the incoming-delegation cap (griefing guard).
+    DelegateLimitReached,
+    /// No active delegation found for the caller.
+    NoDelegationSet,
+    /// Delegation would create a cycle (A→B→A).
+    DelegationCycle,
 }
 
 /// ---------- CONTRACT ----------
@@ -154,7 +174,7 @@ impl GovernanceContract {
     }
 
     /// Vote on a proposal. `support = true` means FOR, false means AGAINST.
-    /// One address one vote (no weighting).
+    /// Each address counts as 1 vote plus 1 for each address that has delegated to it.
     pub fn vote(
         env: Env,
         voter: Address,
@@ -193,13 +213,21 @@ impl GovernanceContract {
             return Err(GovernanceError::AlreadyVoted);
         }
 
+        // Voting weight = 1 (own vote) + number of addresses delegating to this voter.
+        let delegated: u128 = env
+            .storage()
+            .persistent()
+            .get::<StorageKey, u32>(&StorageKey::DelegateCount(voter.clone()))
+            .unwrap_or(0) as u128;
+        let weight: u128 = 1 + delegated;
+
         if support {
-            p.for_votes += 1;
+            p.for_votes += weight;
             env.storage()
                 .persistent()
                 .set(&StorageKey::Vote(proposal_id.clone(), voter.clone()), &1i32);
         } else {
-            p.against_votes += 1;
+            p.against_votes += weight;
             env.storage()
                 .persistent()
                 .set(&StorageKey::Vote(proposal_id.clone(), voter.clone()), &2i32);
@@ -216,6 +244,101 @@ impl GovernanceContract {
         Ok(())
     }
 
+    /// Delegate the caller's vote to `delegate`.
+    ///
+    /// Storage griefing guard:
+    /// - A delegator may hold at most **one** active delegation at a time.
+    /// - A delegate may receive at most `MAX_INCOMING_DELEGATIONS` incoming delegations.
+    /// - Self-delegation and two-hop cycles (A→B while B→A exists) are rejected.
+    pub fn set_delegate(
+        env: Env,
+        delegator: Address,
+        delegate: Address,
+    ) -> Result<(), GovernanceError> {
+        delegator.require_auth();
+
+        // No self-delegation
+        if delegator == delegate {
+            return Err(GovernanceError::SelfDelegation);
+        }
+
+        // Detect two-hop cycle: reject if delegate has already delegated back to delegator
+        if let Some(delegates_delegate) = env
+            .storage()
+            .persistent()
+            .get::<StorageKey, Address>(&StorageKey::Delegate(delegate.clone()))
+        {
+            if delegates_delegate == delegator {
+                return Err(GovernanceError::DelegationCycle);
+            }
+        }
+
+        // Enforce max-1 active delegation per delegator (griefing guard)
+        if env
+            .storage()
+            .persistent()
+            .has(&StorageKey::Delegate(delegator.clone()))
+        {
+            return Err(GovernanceError::DelegationAlreadySet);
+        }
+
+        // Enforce incoming-delegation cap on the delegate (griefing guard)
+        let incoming: u32 = env
+            .storage()
+            .persistent()
+            .get::<StorageKey, u32>(&StorageKey::DelegateCount(delegate.clone()))
+            .unwrap_or(0);
+        if incoming >= MAX_INCOMING_DELEGATIONS {
+            return Err(GovernanceError::DelegateLimitReached);
+        }
+
+        // Persist delegation and bump counter
+        env.storage()
+            .persistent()
+            .set(&StorageKey::Delegate(delegator.clone()), &delegate);
+        env.storage()
+            .persistent()
+            .set(&StorageKey::DelegateCount(delegate.clone()), &(incoming + 1));
+
+        Ok(())
+    }
+
+    /// Remove the caller's active delegation.
+    pub fn unset_delegate(env: Env, delegator: Address) -> Result<(), GovernanceError> {
+        delegator.require_auth();
+
+        let delegate: Address = env
+            .storage()
+            .persistent()
+            .get::<StorageKey, Address>(&StorageKey::Delegate(delegator.clone()))
+            .ok_or(GovernanceError::NoDelegationSet)?;
+
+        // Decrement incoming count on the previously-pointed-at delegate
+        let incoming: u32 = env
+            .storage()
+            .persistent()
+            .get::<StorageKey, u32>(&StorageKey::DelegateCount(delegate.clone()))
+            .unwrap_or(0);
+        if incoming > 0 {
+            env.storage()
+                .persistent()
+                .set(&StorageKey::DelegateCount(delegate.clone()), &(incoming - 1));
+        }
+
+        env.storage()
+            .persistent()
+            .remove(&StorageKey::Delegate(delegator.clone()));
+
+        Ok(())
+    }
+
+    /// Return the delegate for `delegator`, if any.
+    pub fn get_delegate(env: Env, delegator: Address) -> Option<Address> {
+        env.storage()
+            .persistent()
+            .get::<StorageKey, Address>(&StorageKey::Delegate(delegator))
+    }
+
     /// Validate governance votes for a proposal. Returns (passed: bool, reason: String)
     pub fn validate_proposal(
         env: Env,

contracts/predictify-hybrid/src/governance_tests.rs

@@ -109,6 +109,24 @@ impl GovernanceFixture {
             .ledger()
             .with_mut(|li| li.timestamp = li.timestamp.saturating_add(seconds));
     }
+
+    fn set_delegate(&self, delegator: Address, delegate: Address) -> Result<(), GovernanceError> {
+        self.env.as_contract(&self.contract_id, || {
+            GovernanceContract::set_delegate(self.env.clone(), delegator, delegate)
+        })
+    }
+
+    fn unset_delegate(&self, delegator: Address) -> Result<(), GovernanceError> {
+        self.env.as_contract(&self.contract_id, || {
+            GovernanceContract::unset_delegate(self.env.clone(), delegator)
+        })
+    }
+
+    fn get_delegate(&self, delegator: Address) -> Option<Address> {
+        self.env.as_contract(&self.contract_id, || {
+            GovernanceContract::get_delegate(self.env.clone(), delegator)
+        })
+    }
 }
 
 #[test]
@@ -328,3 +346,180 @@ fn governance_set_voting_period_is_applied_to_new_proposals() {
 
     assert_eq!(proposal.end_time - proposal.start_time, 250);
 }
+
+// ===== DELEGATION TESTS =====
+
+#[test]
+fn delegation_set_and_get_round_trip() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let delegator = Address::generate(&fixture.env);
+    let delegate = Address::generate(&fixture.env);
+
+    assert_eq!(fixture.get_delegate(delegator.clone()), None);
+
+    fixture.set_delegate(delegator.clone(), delegate.clone()).unwrap();
+    assert_eq!(fixture.get_delegate(delegator.clone()), Some(delegate));
+}
+
+#[test]
+fn delegation_unset_removes_entry() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let delegator = Address::generate(&fixture.env);
+    let delegate = Address::generate(&fixture.env);
+
+    fixture.set_delegate(delegator.clone(), delegate.clone()).unwrap();
+    fixture.unset_delegate(delegator.clone()).unwrap();
+
+    assert_eq!(fixture.get_delegate(delegator), None);
+}
+
+#[test]
+fn delegation_rejects_self_delegation() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let addr = Address::generate(&fixture.env);
+
+    let result = fixture.set_delegate(addr.clone(), addr.clone());
+    assert_eq!(result, Err(GovernanceError::SelfDelegation));
+}
+
+#[test]
+fn delegation_rejects_cycle() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let a = Address::generate(&fixture.env);
+    let b = Address::generate(&fixture.env);
+
+    // A delegates to B; then B tries to delegate back to A
+    fixture.set_delegate(a.clone(), b.clone()).unwrap();
+    let result = fixture.set_delegate(b.clone(), a.clone());
+    assert_eq!(result, Err(GovernanceError::DelegationCycle));
+}
+
+#[test]
+fn delegation_rejects_double_set_without_unset() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let delegator = Address::generate(&fixture.env);
+    let d1 = Address::generate(&fixture.env);
+    let d2 = Address::generate(&fixture.env);
+
+    fixture.set_delegate(delegator.clone(), d1.clone()).unwrap();
+    let result = fixture.set_delegate(delegator.clone(), d2.clone());
+    assert_eq!(result, Err(GovernanceError::DelegationAlreadySet));
+}
+
+#[test]
+fn delegation_allows_reset_after_unset() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let delegator = Address::generate(&fixture.env);
+    let d1 = Address::generate(&fixture.env);
+    let d2 = Address::generate(&fixture.env);
+
+    fixture.set_delegate(delegator.clone(), d1.clone()).unwrap();
+    fixture.unset_delegate(delegator.clone()).unwrap();
+    fixture.set_delegate(delegator.clone(), d2.clone()).unwrap();
+
+    assert_eq!(fixture.get_delegate(delegator), Some(d2));
+}
+
+#[test]
+fn delegation_unset_without_active_delegation_errors() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let delegator = Address::generate(&fixture.env);
+
+    let result = fixture.unset_delegate(delegator);
+    assert_eq!(result, Err(GovernanceError::NoDelegationSet));
+}
+
+#[test]
+fn delegation_griefing_guard_rejects_beyond_cap() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let delegate = Address::generate(&fixture.env);
+
+    // Fill delegate up to the cap (MAX_INCOMING_DELEGATIONS = 50)
+    for _ in 0..50 {
+        let d = Address::generate(&fixture.env);
+        fixture.set_delegate(d, delegate.clone()).unwrap();
+    }
+
+    // The 51st should be rejected
+    let overflow = Address::generate(&fixture.env);
+    let result = fixture.set_delegate(overflow, delegate.clone());
+    assert_eq!(result, Err(GovernanceError::DelegateLimitReached));
+}
+
+#[test]
+fn delegation_weight_counted_in_vote() {
+    // voter_one has two delegators; their single vote should count as 3
+    let fixture = GovernanceFixture::new(100, 3);
+    let proposal_id = fixture.create_noop_proposal("gov_deleg_1");
+
+    let delegator_a = Address::generate(&fixture.env);
+    let delegator_b = Address::generate(&fixture.env);
+
+    fixture
+        .set_delegate(delegator_a.clone(), fixture.voter_one.clone())
+        .unwrap();
+    fixture
+        .set_delegate(delegator_b.clone(), fixture.voter_one.clone())
+        .unwrap();
+
+    // voter_one votes FOR (weight = 1 own + 2 delegated = 3 → meets quorum of 3)
+    fixture
+        .vote(fixture.voter_one.clone(), proposal_id.clone(), true)
+        .unwrap();
+
+    fixture.advance_time(100);
+    let (passed, _) = fixture.validate(proposal_id).unwrap();
+    assert!(passed, "proposal should pass with delegated weight");
+}
+
+#[test]
+fn delegation_against_weight_counted_in_vote() {
+    // voter_one has 1 delegator; their against vote counts as 2
+    let fixture = GovernanceFixture::new(100, 1);
+    let proposal_id = fixture.create_noop_proposal("gov_deleg_2");
+
+    let delegator = Address::generate(&fixture.env);
+    fixture
+        .set_delegate(delegator.clone(), fixture.voter_one.clone())
+        .unwrap();
+
+    // voter_two: 1 FOR; voter_one: 2 AGAINST → proposal should not pass
+    fixture
+        .vote(fixture.voter_two.clone(), proposal_id.clone(), true)
+        .unwrap();
+    fixture
+        .vote(fixture.voter_one.clone(), proposal_id.clone(), false)
+        .unwrap();
+
+    fixture.advance_time(100);
+    let (passed, _) = fixture.validate(proposal_id).unwrap();
+    assert!(!passed, "proposal should fail when against weight exceeds for");
+}
+
+#[test]
+fn unset_delegation_restores_delegate_capacity() {
+    let fixture = GovernanceFixture::new(100, 1);
+    let delegate = Address::generate(&fixture.env);
+
+    // Fill to cap
+    let mut delegators = soroban_sdk::Vec::new(&fixture.env);
+    for _ in 0..50 {
+        let d = Address::generate(&fixture.env);
+        fixture.set_delegate(d.clone(), delegate.clone()).unwrap();
+        delegators.push_back(d);
+    }
+
+    // Overflow still rejected
+    let overflow = Address::generate(&fixture.env);
+    assert_eq!(
+        fixture.set_delegate(overflow.clone(), delegate.clone()),
+        Err(GovernanceError::DelegateLimitReached)
+    );
+
+    // Remove one delegation
+    let first = delegators.get(0).unwrap();
+    fixture.unset_delegate(first).unwrap();
+
+    // Now the overflow address can delegate
+    fixture.set_delegate(overflow, delegate).unwrap();
+}

contracts/predictify-hybrid/src/lib.rs

@@ -125,7 +125,7 @@ mod upgrade_manager_tests;
 // mod event_management_tests;
 
 // #[cfg(test)]
-// mod governance_tests;
+mod governance_tests;
 
 #[cfg(any())]
 mod category_tags_tests;