From 57fa61a01038103cec4ea0e438231d54fed6b98a Mon Sep 17 00:00:00 2001
From: Hongrui Fang <chfanghr@gmail.com>
Date: Thu, 13 Oct 2022 18:08:36 +0800
Subject: [PATCH] prevent sst from leaving stake validator

---
 agora/Agora/Stake/Scripts.hs | 48 +++++++++++++++++++++---------------
 agora/Agora/Utils.hs         | 10 ++++++++
 2 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/agora/Agora/Stake/Scripts.hs b/agora/Agora/Stake/Scripts.hs
index ef09ecf..67aa70c 100644
--- a/agora/Agora/Stake/Scripts.hs
+++ b/agora/Agora/Stake/Scripts.hs
@@ -61,7 +61,7 @@ import Agora.Stake.Redeemers (
   ppermitVote,
   pretractVote,
  )
-import Agora.Utils (pmapMaybe)
+import Agora.Utils (passert, pmapMaybe)
 import Data.Tagged (Tagged (Tagged))
 import Plutarch.Api.V1 (
   KeyGuarantees (Sorted),
@@ -93,6 +93,7 @@ import Plutarch.Extra.Maybe (
   pmaybeData,
   pnothing,
  )
+import Plutarch.Extra.Ord (POrdering (PEQ, PGT, PLT), pcompareBy, pfromOrd)
 import Plutarch.Extra.ScriptContext (
   pfindTxInByTxOutRef,
   pfromOutputDatum,
@@ -290,26 +291,33 @@ mkStakeValidator
       -- Returns stake datum if the given UTxO is a stake UTxO.
       getStakeDatum :: Term _ (PTxOut :--> PMaybe PStakeDatum) <-
         pletC $
-          plam $ \txOut -> unTermCont $ do
-            txOutF <- pletFieldsC @'["value", "datum", "address"] txOut
+          plam $
+            flip (pletFields @'["value", "datum", "address"]) $ \txOutF ->
+              pmatch
+                ( pcompareBy # pfromOrd
+                    # (sstValueOf # txOutF.value)
+                    # 1
+                )
+                $ \case
+                  -- > 1
+                  PGT -> ptraceError "More than one SST in one UTxO"
+                  -- 1
+                  PEQ ->
+                    let ownedByStakeValidator =
+                          txOutF.address #== stakeValidatorAddress
 
-            let isStakeUTxO =
-                  foldl1
-                    (#&&)
-                    [ ptraceIfFalse "Carries SST" $
-                        sstValueOf # txOutF.value #== 1
-                    , ptraceIfFalse "Owned by stake validator" $
-                        txOutF.address #== stakeValidatorAddress
-                    ]
-
-                datum =
-                  ptrace "Resolve stake datum" $
-                    pfromData $
-                      pfromOutputDatum @(PAsData PStakeDatum)
-                        # txOutF.datum
-                        # txInfoF.datums
-
-            pure $ pif isStakeUTxO (pjust # datum) pnothing
+                        datum =
+                          ptrace "Resolve stake datum" $
+                            pfromData $
+                              pfromOutputDatum @(PAsData PStakeDatum)
+                                # txOutF.datum
+                                # txInfoF.datums
+                     in passert
+                          "Should owned by stake validator"
+                          ownedByStakeValidator
+                          (pjust # datum)
+                  -- 0
+                  PLT -> pnothing
 
       --------------------------------------------------------------------------
 
diff --git a/agora/Agora/Utils.hs b/agora/Agora/Utils.hs
index 5ce3d2b..a877d14 100644
--- a/agora/Agora/Utils.hs
+++ b/agora/Agora/Utils.hs
@@ -33,6 +33,7 @@ module Agora.Utils (
   pltBy,
   pinsertUniqueBy,
   ptryFromRedeemer,
+  passert,
 ) where
 
 import Plutarch.Api.V1 (KeyGuarantees (Unsorted), PPOSIXTime, PRedeemer, PTokenName, PValidatorHash)
@@ -404,3 +405,12 @@ ptryFromRedeemer = phoistAcyclic $
     pfmap
       # plam (flip ptryFrom fst . pto)
       # (plookup # p # m)
+
+-- | @since 1.0.0
+passert ::
+  forall (a :: PType) (s :: S).
+  Term s PString ->
+  Term s PBool ->
+  Term s a ->
+  Term s a
+passert msg cond x = pif cond x $ ptraceError msg