From dfe4bba15f66717e06a8f17131674b5ce15f3d24 Mon Sep 17 00:00:00 2001
From: fanghr <chfanghr@gmail.com>
Date: Thu, 19 May 2022 13:58:12 +0800
Subject: [PATCH] ensure the new proposal lock is placed on the stake

---
 agora/Agora/Proposal/Scripts.hs |  5 +++--
 agora/Agora/Stake/Scripts.hs    | 29 ++++++++++++++++++++++++++---
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/agora/Agora/Proposal/Scripts.hs b/agora/Agora/Proposal/Scripts.hs
index 73de60b..dc85b71 100644
--- a/agora/Agora/Proposal/Scripts.hs
+++ b/agora/Agora/Proposal/Scripts.hs
@@ -271,8 +271,8 @@ proposalValidator proposal =
           tcassert "Invalid output proposal" $ proposalOut #== expectedProposalOut
 
           -- We validate the output stake datum here as well: We need the vote option
-          --   to create a proper 'ProposalLock'. However the vote option is encoded
-          --   in the proposal redeemer, which is invisible for the stake validator.
+          -- to create a valid 'ProposalLock', however the vote option is encoded
+          -- in the proposal redeemer, which is invisible for the stake validator.
 
           let stakeOutput =
                 mustBePJust # "Stake output not found"
@@ -292,6 +292,7 @@ proposalValidator proposal =
                   ( #vote .= pdata voteFor
                       .& #proposalTag .= proposalF.proposalId
                   )
+              -- Prepend the new lock to existing locks
               expectedProposalLocks =
                 pcons
                   # pdata newProposalLock
diff --git a/agora/Agora/Stake/Scripts.hs b/agora/Agora/Stake/Scripts.hs
index 74494e8..aaf114b 100644
--- a/agora/Agora/Stake/Scripts.hs
+++ b/agora/Agora/Stake/Scripts.hs
@@ -7,6 +7,7 @@ Plutus Scripts for Stakes.
 -}
 module Agora.Stake.Scripts (stakePolicy, stakeValidator) where
 
+import Agora.Record (mkRecordConstr, (.&), (.=))
 import Agora.SafeMoney (GTTag)
 import Agora.Stake
 import Agora.Utils (
@@ -222,7 +223,7 @@ stakeValidator stake =
     -- TODO: Use PTryFrom
     let stakeDatum' :: Term _ PStakeDatum
         stakeDatum' = pfromData $ punsafeCoerce datum
-    stakeDatum <- tcont $ pletFields @'["owner", "stakedAmount"] stakeDatum'
+    stakeDatum <- tcont $ pletFields @'["owner", "stakedAmount", "lockedBy"] stakeDatum'
 
     PSpending txOutRef <- tcmatch $ pfromData ctx.purpose
 
@@ -291,7 +292,7 @@ stakeValidator stake =
 
           pure $ popaque (pconstant ())
         --------------------------------------------------------------------------
-        PPermitVote _ -> unTermCont $ do
+        PPermitVote l -> unTermCont $ do
           tcassert
             "Owner signs this transaction"
             ownerSignsTransaction
@@ -301,18 +302,40 @@ stakeValidator stake =
           tcassert "Proposal ST spent" $
             spentProposalST #== 1
 
+          -- Update the stake datum, but only the 'lockedBy' field.
+
+          let -- We actually don't know whether the given lock is valid or not.
+              -- This is checked in the proposal validator.
+              newLock = pfield @"lock" # l
+              -- Prepend the new lock to the existing locks.
+              expectedLocks = pcons # newLock # stakeDatum.lockedBy
+
+          expectedDatum <-
+            tclet $
+              pdata $
+                mkRecordConstr
+                  PStakeDatum
+                  ( #stakedAmount .= stakeDatum.stakedAmount
+                      .& #owner .= stakeDatum.owner
+                      .& #lockedBy .= pdata expectedLocks
+                  )
+
           tcassert "A UTXO must exist with the correct output" $
+            -- FIXME: no need to pass the whole txInfo to 'anyOutput'.
             anyOutput @PStakeDatum # txInfo
               #$ plam
               $ \value address newStakeDatum' ->
                 let isScriptAddress = pdata address #== ownAddress
-                    _correctOutputDatum = pdata newStakeDatum' #== pdata stakeDatum'
+                    correctOutputDatum = pdata newStakeDatum' #== expectedDatum
+                    -- TODO: Is this correct? I think We only need to ensure
+                    --       correct amount of GT/SST in the continuing output.
                     valueCorrect = pdata continuingValue #== pdata value
                  in pif
                       isScriptAddress
                       ( foldl1
                           (#&&)
                           [ ptraceIfFalse "valueCorrect" valueCorrect
+                          , ptraceIfFalse "datumCorrect" correctOutputDatum
                           ]
                       )
                       (pcon PFalse)