WitnessStake: accept multiple stakes at input

agora/Agora/Stake/Scripts.hs

@@ -8,7 +8,18 @@ Plutus Scripts for Stakes.
 module Agora.Stake.Scripts (stakePolicy, stakeValidator) where
 
 import Agora.SafeMoney (GTTag)
-import Agora.Stake
+import Agora.Stake (
+  PStakeDatum (PStakeDatum),
+  PStakeRedeemer (
+    PDepositWithdraw,
+    PDestroy,
+    PPermitVote,
+    PRetractVotes
+  ),
+  Stake (gtClassRef, proposalSTClass),
+  StakeRedeemer (WitnessStake),
+  stakeLocked,
+ )
 import Agora.Utils (
   mustBePJust,
   mustFindDatum',
@@ -18,18 +29,22 @@ import Data.Tagged (Tagged (..), untag)
 import Plutarch.Api.V1 (
   AmountGuarantees (Positive),
   PCredential (PPubKeyCredential, PScriptCredential),
+  PDatumHash,
   PMintingPolicy,
   PScriptPurpose (PMinting, PSpending),
   PTokenName,
   PTxInfo,
+  PTxOut,
   PValidator,
   PValue,
   mintingPolicySymbol,
   mkMintingPolicy,
  )
 import Plutarch.Api.V1.AssetClass (passetClass, passetClassValueOf, pvalueOf)
-import Plutarch.Api.V1.ScriptContext (pfindTxInByTxOutRef, pisTokenSpent, ptxSignedBy, pvalueSpent)
+import Plutarch.Api.V1.ScriptContext (pfindTxInByTxOutRef, ptxSignedBy, pvalueSpent)
 import "liqwid-plutarch-extra" Plutarch.Api.V1.Value (pgeqByClass', pgeqBySymbol, psymbolValueOf)
+import Plutarch.Extra.List (pmapMaybe, pmsortBy)
+import Plutarch.Extra.Maybe (pfromDJust)
 import Plutarch.Extra.Record (mkRecordConstr, (.&), (.=))
 import Plutarch.Extra.TermCont (pguardC, pletC, pletFieldsC, pmatchC, ptryFromC)
 import Plutarch.Internal (punsafeCoerce)
@@ -208,7 +223,15 @@ stakeValidator stake =
   plam $ \datum redeemer ctx' -> unTermCont $ do
     ctx <- pletFieldsC @'["txInfo", "purpose"] ctx'
     txInfo <- pletC $ pfromData ctx.txInfo
-    txInfoF <- pletFieldsC @'["mint", "inputs", "outputs", "signatories", "datums"] txInfo
+    txInfoF <-
+      pletFieldsC
+        @'[ "mint"
+          , "inputs"
+          , "outputs"
+          , "signatories"
+          , "datums"
+          ]
+        txInfo
 
     (pfromData -> stakeRedeemer, _) <- ptryFromC redeemer
 
@@ -219,23 +242,25 @@ stakeValidator stake =
 
     PSpending txOutRef <- pmatchC $ pfromData ctx.purpose
 
-    PJust txInInfo <- pmatchC $ pfindTxInByTxOutRef # (pfield @"_0" # txOutRef) # txInfoF.inputs
+    PJust ((pfield @"resolved" #) -> resolved) <-
-    ownAddress <- pletC $ pfield @"address" #$ pfield @"resolved" # txInInfo
+      pmatchC $
-    let continuingValue :: Term _ (PValue _ _)
+        pfindTxInByTxOutRef
-        continuingValue = pfield @"value" #$ pfield @"resolved" # txInInfo
+          # (pfield @"_0" # txOutRef)
+          # txInfoF.inputs
+    resolvedF <- pletFieldsC @'["address", "value", "datumHash"] resolved
 
     -- Whether the owner signs this transaction or not.
     ownerSignsTransaction <- pletC $ ptxSignedBy # txInfoF.signatories # stakeDatum.owner
 
-    stCurrencySymbol <- pletC $ pconstant $ mintingPolicySymbol $ mkMintingPolicy (stakePolicy stake.gtClassRef)
+    stCurrencySymbol <-
+      pletC $
+        pconstant $
+          mintingPolicySymbol $
+            mkMintingPolicy (stakePolicy stake.gtClassRef)
     mintedST <- pletC $ psymbolValueOf # stCurrencySymbol # txInfoF.mint
     valueSpent <- pletC $ pvalueSpent # txInfoF.inputs
     spentST <- pletC $ psymbolValueOf # stCurrencySymbol #$ valueSpent
 
-    let AssetClass (propCs, propTn) = stake.proposalSTClass
-        proposalSTClass = passetClass # pconstant propCs # pconstant propTn
-    spentProposalST <- pletC $ passetClassValueOf # valueSpent # proposalSTClass
-
     -- Is the stake currently locked?
     stakeIsLocked <- pletC $ stakeLocked # stakeDatum'
 
@@ -253,22 +278,92 @@ stakeValidator stake =
           pguardC "Owner signs this transaction" ownerSignsTransaction
 
           pure $ popaque (pconstant ())
-        --------------------------------------------------------------------------
+        ------------------------------------------------------------------------
         -- Handle redeemers that require own stake output.
+
         _ -> unTermCont $ do
-          -- Filter out own output with own address and PST.
+          let AssetClass (propCs, propTn) = stake.proposalSTClass
-          ownOutput <-
+              proposalSTClass = passetClass # pconstant propCs # pconstant propTn
+              spentProposalST = passetClassValueOf # valueSpent # proposalSTClass
+
+          proposalTokenMoved <- pletC $ spentProposalST #== 1
+
+          -- Filter out own outputs using own address and ST.
+          ownOutputs <-
             pletC $
-              mustBePJust # "Own output should be present" #$ pfind
+              pfilter
                 # plam
-                  ( \input -> unTermCont $ do
+                  ( \output -> unTermCont $ do
-                      inputF <- pletFieldsC @'["address", "value"] input
+                      outputF <- pletFieldsC @'["address", "value"] output
+
                       pure $
-                        inputF.address #== ownAddress
+                        outputF.address #== resolvedF.address
-                          #&& psymbolValueOf # stCurrencySymbol # inputF.value #== 1
+                          #&& psymbolValueOf # stCurrencySymbol # outputF.value #== 1
                   )
                 # pfromData txInfoF.outputs
 
+          let witnessStake = unTermCont $ do
+                pguardC "Either owner signs the transaction or propsoal token moved" $
+                  ownerSignsTransaction #|| proposalTokenMoved
+
+                -- FIXME: refactor this with reference input, once it's supported by plutarch.
+                --
+                -- Our goal here is to allow multiple input stakes, and also ensure that every the input stakes has a
+                --   corresponding output stake, which carries the same value and the same datum as the input stake.
+                --
+                -- Validation strategy I have tried/considered so far:
+                -- 1. Check that the number of input stakes equals to the number of output stakes, and verify
+                --     that every input stake has an output stake with the exact same value and datum hash.
+                --     However this approach has a fatal vulnerability: let's say we have two totally identical stakes,
+                --     a malicious user can comsume these two stakes and remove GTs from one of them.
+                -- 2. Perform the same checks as the last approch does, while also checking that every output stake is
+                --     valid(stakedAmount == actual value). However this requires that all the output stake datum are
+                --     included in the transaction, and we have to find and go through them one by one to access the
+                --     'stakedAmount' fields, meaning that computationally this approach is *very* expensive.
+                -- 3. The one implemented below. Find all the continuous input/output, sort them by 'datumHash', and
+                --     ensure that the two sorted lists are equal.
+                let ownInputs =
+                      pmapMaybe
+                        # plam
+                          ( \input -> plet (pfield @"resolved" # input) $ \resolvedInput ->
+                              let value = pfield @"value" # resolvedInput
+                               in pif
+                                    (psymbolValueOf # stCurrencySymbol # value #== 1)
+                                    (pcon $ PJust resolvedInput)
+                                    (pcon PNothing)
+                          )
+                        # pfromData txInfoF.inputs
+
+                    sortTxOuts :: Term _ (PBuiltinList (PAsData PTxOut) :--> PBuiltinList (PAsData PTxOut))
+                    sortTxOuts =
+                      plam
+                        ( pmsortBy
+                            # plam
+                              ( \((getDatumHash #) -> dhX)
+                                 ((getDatumHash #) -> dhY) -> dhX #< dhY
+                              )
+                            #
+                        )
+                      where
+                        getDatumHash :: Term _ (PAsData PTxOut :--> PDatumHash)
+                        getDatumHash = phoistAcyclic $ plam ((pfromDJust #) . pfromData . (pfield @"datumHash" #))
+
+                    sortedOwnInputs = sortTxOuts # ownInputs
+                    sortedOwnOutputs = sortTxOuts # ownOutputs
+
+                pguardC "Every stake inputs has a corresponding unchanged output" $
+                  plistEquals # sortedOwnInputs # sortedOwnOutputs
+
+                pure $ popaque $ pconstant ()
+
+          ----------------------------------------------------------------------
+
+          let onlyAcceptOneStake = unTermCont $ do
+                pguardC "ST at inputs must be 1" $
+                  spentST #== 1
+
+                ownOutput <- pletC $ pfromData $ phead # ownOutputs
+
                 stakeOut <-
                   pletC $
                     mustFindDatum' @PStakeDatum
@@ -281,11 +376,7 @@ stakeValidator stake =
 
                 ownOutputValueUnchanged <-
                   pletC $
-              pdata continuingValue #== pdata ownOutputValue
+                    pdata resolvedF.value #== pdata ownOutputValue
-
-          stakeOutUnchanged <-
-            pletC $
-              pdata stakeOut #== pdata stakeDatum'
 
                 pure $
                   pmatch stakeRedeemer $ \case
@@ -294,13 +385,9 @@ stakeValidator stake =
                         "Owner signs this transaction"
                         ownerSignsTransaction
 
-                pguardC "ST at inputs must be 1" $
-                  spentST #== 1
-
                       -- This puts trust into the Proposal. The Proposal must necessarily check
                       -- that this is not abused.
-                pguardC "Proposal ST spent" $
+                      pguardC "Proposal ST spent" proposalTokenMoved
-                  spentProposalST #== 1
 
                       pguardC "A UTXO must exist with the correct output" $
                         let expectedLocks = pfield @"locks" # l
@@ -322,7 +409,9 @@ stakeValidator stake =
                               ]
 
                       pure $ popaque (pconstant ())
-              --------------------------------------------------------------------------
+
+                    ------------------------------------------------------------
+
                     PPermitVote l -> unTermCont $ do
                       pguardC
                         "Owner signs this transaction"
@@ -330,8 +419,7 @@ stakeValidator stake =
 
                       -- This puts trust into the Proposal. The Proposal must necessarily check
                       -- that this is not abused.
-                pguardC "Proposal ST spent" $
+                      pguardC "Proposal ST spent" proposalTokenMoved
-                  spentProposalST #== 1
 
                       -- Update the stake datum, but only the 'lockedBy' field.
 
@@ -360,38 +448,10 @@ stakeValidator stake =
                               ]
 
                       pure $ popaque (pconstant ())
-              --------------------------------------------------------------------------
-              PWitnessStake _ -> unTermCont $ do
-                pguardC "ST at inputs must be 1" $
-                  spentST #== 1
 
-                let AssetClass (propCs, propTn) = stake.proposalSTClass
+                    ------------------------------------------------------------
-                    propAssetClass = passetClass # pconstant propCs # pconstant propTn
-                    proposalTokenMoved =
-                      pisTokenSpent
-                        # propAssetClass
-                        # txInfoF.inputs
 
-                -- In order for cosignature to be witnessed, it must be possible for a
-                -- proposal to allow this transaction to happen. This puts trust into the Proposal.
-                -- The Proposal must necessarily check that this is not abused.
-                pguardC
-                  "Owner signs this transaction OR proposal token is spent"
-                  (ownerSignsTransaction #|| proposalTokenMoved)
-
-                pguardC "A UTXO must exist with the correct output" $
-                  let correctOutputDatum = stakeOutUnchanged
-                      valueCorrect = ownOutputValueUnchanged
-                   in foldl1
-                        (#&&)
-                        [ ptraceIfFalse "valueCorrect" valueCorrect
-                        , ptraceIfFalse "correctOutputDatum" correctOutputDatum
-                        ]
-                pure $ popaque (pconstant ())
-              --------------------------------------------------------------------------
                     PDepositWithdraw r -> unTermCont $ do
-                pguardC "ST at inputs must be 1" $
-                  spentST #== 1
                       pguardC "Stake unlocked" $
                         pnot #$ stakeIsLocked
                       pguardC
@@ -420,7 +480,7 @@ stakeValidator stake =
                               valueDelta = pdiscreteValue' stake.gtClassRef # delta
 
                               expectedValue =
-                          continuingValue <> valueDelta
+                                resolvedF.value <> valueDelta
 
                               valueCorrect =
                                 foldr1
@@ -446,3 +506,9 @@ stakeValidator stake =
                       --
                       pure $ popaque (pconstant ())
                     _ -> popaque (pconstant ())
+
+          pure $
+            pif
+              (pdata stakeRedeemer #== pconstantData WitnessStake)
+              witnessStake
+              onlyAcceptOneStake

bench.csv

@@ -5,13 +5,13 @@ Agora/Effects/Treasury Withdrawal Effect/effect/Mixed Assets,456007605,1104500,3
 Agora/Effects/Governor Mutation Effect/validator/valid new governor datum/governor validator should pass,87839169,243032,8561
 Agora/Effects/Governor Mutation Effect/validator/valid new governor datum/effect validator should pass,106082031,292993,3609
 Agora/Stake/policy/stakeCreation,50939580,148729,2387
-Agora/Stake/validator/stakeDepositWithdraw deposit,181581435,493259,4413
+Agora/Stake/validator/stakeDepositWithdraw deposit,180222751,492217,5003
-Agora/Stake/validator/stakeDepositWithdraw withdraw,181581435,493259,4401
+Agora/Stake/validator/stakeDepositWithdraw withdraw,180222751,492217,4991
 Agora/Proposal/policy/proposalCreation,23140177,69194,1515
 Agora/Proposal/validator/cosignature/proposal,240482868,674626,8525
-Agora/Proposal/validator/cosignature/stake,125315872,312659,4942
+Agora/Proposal/validator/cosignature/stake,136781411,336612,5528
 Agora/Proposal/validator/voting/proposal,243946100,678901,8443
-Agora/Proposal/validator/voting/stake,120122971,320497,4899
+Agora/Proposal/validator/voting/stake,128972262,348186,5489
 Agora/Proposal/validator/advancing/successfully advance to next state/Draft -> VotringReady,219342631,620576,8350
 Agora/Proposal/validator/advancing/successfully advance to next state/VotingReady -> Locked,247748475,699343,8359
 Agora/Proposal/validator/advancing/successfully advance to next state/Locked -> Finished,236509366,666512,8359