From a51595cd1e561383e4959948e619e5a75b534927 Mon Sep 17 00:00:00 2001
From: Hongrui Fang <chfanghr@gmail.com>
Date: Thu, 20 Oct 2022 20:32:15 +0800
Subject: [PATCH] regression tests for privilege escalation while voting

---
 .../Sample/Proposal/PrivilegeEscalate.hs      | 242 ++++++++++++++++++
 agora-specs/Spec/Proposal.hs                  |  12 +
 agora.cabal                                   |   1 +
 3 files changed, 255 insertions(+)
 create mode 100644 agora-specs/Sample/Proposal/PrivilegeEscalate.hs

diff --git a/agora-specs/Sample/Proposal/PrivilegeEscalate.hs b/agora-specs/Sample/Proposal/PrivilegeEscalate.hs
new file mode 100644
index 0000000..24f4e30
--- /dev/null
+++ b/agora-specs/Sample/Proposal/PrivilegeEscalate.hs
@@ -0,0 +1,242 @@
+module Sample.Proposal.PrivilegeEscalate (
+  Operation (..),
+  privilegeEscalate,
+  Validity (..),
+  mkTestTree,
+) where
+
+import Agora.Proposal (
+  ProposalDatum (..),
+  ProposalId (ProposalId),
+  ProposalRedeemer (Unlock, Vote),
+  ProposalStatus (VotingReady),
+  ProposalVotes (ProposalVotes),
+  ResultTag (ResultTag),
+  emptyVotesFor,
+ )
+import Agora.Proposal.Time (
+  ProposalStartingTime (ProposalStartingTime),
+  ProposalTimingConfig (draftTime, votingTime),
+ )
+import Agora.SafeMoney (GTTag)
+import Agora.Stake (
+  ProposalLock (
+    Voted
+  ),
+  StakeDatum (..),
+  StakeRedeemer (PermitVote, RetractVotes),
+ )
+import Data.Default (Default (def))
+import Data.Map.Strict qualified as StrictMap
+import Data.Tagged (Tagged, untag)
+import Plutarch.Context (
+  input,
+  normalizeValue,
+  output,
+  script,
+  signedWith,
+  timeRange,
+  withDatum,
+  withRedeemer,
+  withRef,
+  withValue,
+ )
+import Plutarch.Extra.AssetClass (assetClassValue)
+import PlutusLedgerApi.V1 (Credential (PubKeyCredential))
+import PlutusLedgerApi.V2 (PubKeyHash, TxOutRef (TxOutRef))
+import Sample.Proposal.Shared (proposalTxRef, stakeTxRef)
+import Sample.Shared (
+  minAda,
+  proposalAssetClass,
+  proposalValidator,
+  proposalValidatorHash,
+  stakeAssetClass,
+  stakeValidator,
+  stakeValidatorHash,
+ )
+import Test.Specification (SpecificationTree, group, testValidator)
+import Test.Util (CombinableBuilder, closedBoundedInterval, mkSpending, pubKeyHashes)
+
+data Operation = Voting | RetractingVotes
+
+data Validity = Validity
+  { forStakeValidator :: Bool
+  , forProposalValidator :: Bool
+  }
+
+wrap :: forall x y. Operation -> (x -> x -> y) -> x -> x -> y
+wrap Voting = id
+wrap RetractingVotes = flip
+
+defStakeAmount :: Tagged GTTag Integer
+defStakeAmount = 100000
+
+defResultTag :: ResultTag
+defResultTag = ResultTag 0
+
+defProposalId :: ProposalId
+defProposalId = ProposalId 0
+
+mkProposalInputOutputDatum :: Operation -> (ProposalDatum, ProposalDatum)
+mkProposalInputOutputDatum op =
+  let effects = StrictMap.singleton defResultTag StrictMap.empty
+
+      proposal =
+        ProposalDatum
+          { proposalId = defProposalId
+          , effects = effects
+          , status = VotingReady
+          , cosigners = [] -- doesn't matter
+          , thresholds = def
+          , votes = emptyVotesFor effects
+          , timingConfig = def
+          , startingTime = ProposalStartingTime 0
+          }
+
+      proposalWithVotes =
+        proposal
+          { votes =
+              ProposalVotes $
+                StrictMap.singleton defResultTag (untag defStakeAmount)
+          }
+   in wrap op (,) proposal proposalWithVotes
+
+mkProposalRedeemer :: Operation -> ProposalRedeemer
+mkProposalRedeemer op = wrap op const (Vote defResultTag) Unlock
+
+proposalRef :: TxOutRef
+proposalRef = TxOutRef proposalTxRef 1
+
+attacker :: PubKeyHash
+attacker = head pubKeyHashes
+
+mkStakeInputOutputDatums :: Operation -> ([StakeDatum], [StakeDatum])
+mkStakeInputOutputDatums op =
+  let delegatee = pubKeyHashes !! 1
+
+      firstStake =
+        StakeDatum
+          { stakedAmount = defStakeAmount
+          , owner = PubKeyCredential attacker
+          , delegatedTo = Just $ PubKeyCredential delegatee
+          , lockedBy = []
+          }
+
+      otherStakes =
+        (\pkh -> firstStake {owner = PubKeyCredential pkh})
+          <$> drop 2 pubKeyHashes
+
+      allStakes = take 10 $ firstStake : otherStakes
+
+      stakeWithLock =
+        (\stake -> stake {lockedBy = [Voted defProposalId defResultTag]})
+          <$> allStakes
+   in wrap op (,) allStakes stakeWithLock
+
+mkStakeRedeemer :: Operation -> StakeRedeemer
+mkStakeRedeemer op = wrap op const PermitVote RetractVotes
+
+mkStakeRef :: Integer -> TxOutRef
+mkStakeRef o = TxOutRef stakeTxRef $ 1 + o
+
+privilegeEscalate :: forall b. CombinableBuilder b => Operation -> b
+privilegeEscalate op =
+  let sst = assetClassValue stakeAssetClass 1
+
+      stakeValue = normalizeValue $ minAda <> sst
+
+      (stakeInputDatums, stakeOutputDatums) = mkStakeInputOutputDatums op
+
+      stakeBuilder =
+        mconcat $
+          zipWith3
+            ( \index stakeInput stakeOutput ->
+                mconcat @b
+                  [ input $
+                      mconcat
+                        [ script stakeValidatorHash
+                        , withDatum stakeInput
+                        , withValue stakeValue
+                        , withRef $ mkStakeRef index
+                        , withRedeemer $ mkStakeRedeemer op
+                        ]
+                  , output $
+                      mconcat
+                        [ script stakeValidatorHash
+                        , withDatum stakeOutput
+                        , withValue stakeValue
+                        ]
+                  ]
+            )
+            [1 ..]
+            stakeInputDatums
+            stakeOutputDatums
+
+      ---
+
+      pst = assetClassValue proposalAssetClass 1
+
+      proposalValue = normalizeValue $ minAda <> pst
+
+      (proposalInput, proposalOutput) = mkProposalInputOutputDatum op
+
+      proposalBuilder =
+        mconcat @b
+          [ input $
+              mconcat
+                [ script proposalValidatorHash
+                , withDatum proposalInput
+                , withRedeemer $ mkProposalRedeemer op
+                , withValue proposalValue
+                , withRef proposalRef
+                ]
+          , output $
+              mconcat
+                [ script proposalValidatorHash
+                , withDatum proposalOutput
+                , withValue proposalValue
+                ]
+          ]
+
+      ---
+
+      validTimeRange =
+        closedBoundedInterval
+          ((def :: ProposalTimingConfig).draftTime + 1)
+          ((def :: ProposalTimingConfig).votingTime - 1)
+
+      miscBuilder =
+        mconcat @b
+          [ signedWith attacker
+          , timeRange validTimeRange
+          ]
+   in mconcat
+        [ miscBuilder
+        , stakeBuilder
+        , proposalBuilder
+        ]
+
+mkTestTree :: String -> Operation -> Validity -> SpecificationTree
+mkTestTree name op val = group name [proposal, stake]
+  where
+    spend = mkSpending privilegeEscalate op
+
+    proposal =
+      testValidator
+        val.forProposalValidator
+        "proposal"
+        proposalValidator
+        (fst $ mkProposalInputOutputDatum op)
+        (mkProposalRedeemer op)
+        (spend proposalRef)
+
+    stakeInputdDatum = head $ fst $ mkStakeInputOutputDatums op
+
+    stake =
+      testValidator
+        val.forStakeValidator
+        "stake"
+        stakeValidator
+        stakeInputdDatum
+        (mkStakeRedeemer op)
+        (spend $ mkStakeRef 1)
diff --git a/agora-specs/Spec/Proposal.hs b/agora-specs/Spec/Proposal.hs
index 61049e3..dd82c4f 100644
--- a/agora-specs/Spec/Proposal.hs
+++ b/agora-specs/Spec/Proposal.hs
@@ -10,6 +10,7 @@ module Spec.Proposal (specs) where
 import Sample.Proposal.Advance qualified as Advance
 import Sample.Proposal.Cosign qualified as Cosign
 import Sample.Proposal.Create qualified as Create
+import Sample.Proposal.PrivilegeEscalate qualified as PrivilegeEscalate
 import Sample.Proposal.Unlock qualified as Unlock
 import Sample.Proposal.Vote qualified as Vote
 
@@ -397,4 +398,15 @@ specs =
               illegalGroup = group "illegal" $ map mkIllegalGroup stakeCountCases
            in [legalGroup, illegalGroup]
       ]
+  , group
+      "privilege escalate"
+      [ PrivilegeEscalate.mkTestTree
+          "vote"
+          PrivilegeEscalate.Voting
+          (PrivilegeEscalate.Validity False False)
+      , PrivilegeEscalate.mkTestTree
+          "retract votes"
+          PrivilegeEscalate.RetractingVotes
+          (PrivilegeEscalate.Validity False False)
+      ]
   ]
diff --git a/agora.cabal b/agora.cabal
index 2ee025f..8eb2209 100644
--- a/agora.cabal
+++ b/agora.cabal
@@ -199,6 +199,7 @@ library agora-specs
     Sample.Proposal.Advance
     Sample.Proposal.Cosign
     Sample.Proposal.Create
+    Sample.Proposal.PrivilegeEscalate
     Sample.Proposal.Shared
     Sample.Proposal.Unlock
     Sample.Proposal.Vote