fix bugs and vulnerabilities in premoveLocks

agora/Agora/Stake/Redeemers.hs

@@ -50,7 +50,6 @@ import Agora.Stake (
   ),
   pstakeLocked,
  )
-import Data.Functor ((<&>))
 import Plutarch.Api.V1.Address (PCredential)
 import Plutarch.Api.V2 (PMaybeData, PPOSIXTime)
 import Plutarch.Extra.Bool (passert)
@@ -60,7 +59,7 @@ import "liqwid-plutarch-extra" Plutarch.Extra.List (
   ptryDeleteFirstBy,
   ptryFromSingleton,
  )
-import Plutarch.Extra.Maybe (pdjust, pdnothing, pjust, pmaybe, pmaybeData, pnothing)
+import Plutarch.Extra.Maybe (pdjust, pdnothing, pmaybeData)
 import Plutarch.Extra.Record (mkRecordConstr, (.&), (.=))
 import "liqwid-plutarch-extra" Plutarch.Extra.TermCont (pguardC, pletC, pletFieldsC, pmatchC)
 import Plutarch.Extra.Time (PCurrentTime (PCurrentTime))
@@ -269,14 +268,14 @@ instance DerivePlutusType PRemoveLocksMode where
 {- | Remove stake locks with the proposal id given the list of existing locks.
      The first parameter controls whether to remove creator locks or not. If
      one of the locks performed voting action, the unlock cooldown will be
-     checked if it's given.
+     checked.
 -}
 premoveLocks ::
   forall (s :: S).
   Term
     s
     ( PProposalId
-        :--> PMaybe PPOSIXTime
+        :--> PPOSIXTime
         :--> PProposalTime
         :--> PRemoveLocksMode
         :--> PBuiltinList (PAsData PProposalLock)
@@ -293,31 +292,33 @@ premoveLocks =
             ( (pfield @"createdAt" #) ->
                 createdAt
               ) =
-              let notInCooldown =
+              let notInCooldown = createdAt + unlockCooldown #<= lowerBound
-                    pmaybe
-                      # pconstant True
-                      # plam (\c -> createdAt + c #<= lowerBound)
-                      # unlockCooldown
                in foldl1
                     (#||)
                     [ shouldRemoveAllLocks
-                    , ptraceIfFalse "Stake lock in cooldown" notInCooldown
+                    , -- Fail the transaction if a voter lock is in cooldown.∏
+                      passert
+                        "Voter lock shouldn't be in cooldown"
+                        notInCooldown
+                        (pconstant True)
                     ]
 
-          handleLock =
+          shouldRemoveLock =
-            plam $
+            flip
-              flip
+              pletAll
-                pletAll
+              ( \lockF ->
-                ( \lockF ->
+                  foldl1
-                    foldl1
+                    (#&&)
-                      (#&&)
+                    [ proposalId #== lockF.proposalId
-                      [ proposalId #== lockF.proposalId
+                    , pmatch lockF.action $ \case
-                      , pmatch lockF.action $ \case
+                        PVoted r -> handleVoter r
-                          PVoted r -> handleVoter r
+                        _ -> shouldRemoveAllLocks
-                          _ -> shouldRemoveAllLocks
+                    ]
-                      ]
+              )
-                )
+              . pfromData
-                . pfromData
+
+          -- Return true, given a lock that should be kept.
+          handleLock = plam $ (pnot #) . shouldRemoveLock
 
       pure $ pfilter # handleLock
 
@@ -340,18 +341,13 @@ pretractVote = pvoteHelper #$ phoistAcyclic $
                   ]
                 proposal
 
-            (mode, unlockCooldown) <-
+            let unlockCooldown =
-              pmatchC (proposalF.status #== pconstant Finished) <&> \case
+                  pfield @"minStakeVotingTime"
-                PTrue ->
+                    # proposalF.timingConfig
-                  ( pcon PRemoveAllLocks
+
-                  , pnothing
+                mode = pmatch (proposalF.status #== pconstant Finished) $ \case
-                  )
+                  PTrue -> pcon PRemoveAllLocks
-                _ ->
+                  _ -> pcon PRemoveVoterLockOnly
-                  ( pcon PRemoveVoterLockOnly
-                  , pjust
-                      #$ pfield @"minStakeVotingTime"
-                      # proposalF.timingConfig
-                  )
 
             pguardC "Authorized by either opwner or delegatee" $
               pisSignedBy # pconstant True # ctx