From 2c3a1c036367cb903b57720f3ffcd1b090cb2edc Mon Sep 17 00:00:00 2001
From: Hongrui Fang <chfanghr@gmail.com>
Date: Fri, 25 Nov 2022 21:09:23 +0800
Subject: [PATCH] fix a bug that makes using delegated and own stakes together
 unreliable

---
 agora/Agora/Stake/Scripts.hs | 81 ++++++++++++++++++++++++++----------
 1 file changed, 60 insertions(+), 21 deletions(-)

diff --git a/agora/Agora/Stake/Scripts.hs b/agora/Agora/Stake/Scripts.hs
index aaa6582..7c3d9c7 100644
--- a/agora/Agora/Stake/Scripts.hs
+++ b/agora/Agora/Stake/Scripts.hs
@@ -226,6 +226,23 @@ stakePolicy =
 
 --------------------------------------------------------------------------------
 
+data PAuthCheckHelper (s :: S) = PAuthCheckHelper
+  { allHaveTheSameOwner :: Term s PBool
+  , allOwnByOrDelegatedToTheFirstOwner :: Term s PBool
+  , allOwnByOrDelegatedToTheFirstDelegatee :: Term s PBool
+  }
+  deriving stock
+    ( Generic
+    )
+  deriving anyclass
+    ( PlutusType
+    )
+
+instance DerivePlutusType PAuthCheckHelper where
+  type DPTStrat _ = PlutusTypeScott
+
+--------------------------------------------------------------------------------
+
 {- | Create a stake validator, given the implementation of stake redeemers.
 
      == Arguments
@@ -345,39 +362,61 @@ mkStakeValidator impl sstSymbol pstClass gtClass =
 
     authorizedBy <- pletC $ pauthorizedBy # authorizationContext txInfoF
 
-    PPair allHaveSameOwner allHaveSameOrOwnedByDelegatee <-
+    authCheckHelperF <-
       pmatchC $
         pfoldr
           # plam
-            ( \d p -> unTermCont $ do
-                dF <- pletFieldsC @'["owner", "delegatedTo"] d
+            ( \stake helper -> unTermCont $ do
+                stakeF <- pletFieldsC @'["owner", "delegatedTo"] stake
+                helperF <- pmatchC helper
 
-                pure $
-                  pmatch p $ \(PPair allHaveSameOwner allHaveSameDelegatee) ->
-                    let allHaveSameOwner' =
-                          allHaveSameOwner
-                            #&& dF.owner
-                            #== firstStakeInputDatumF.owner
-                        allHaveSameOrOwnedByDelegatee' =
-                          let delegated =
-                                dF.delegatedTo #== firstStakeInputDatumF.delegatedTo
-                              ownedByDelegatee =
-                                pdata (pdjust # dF.owner)
-                                  #== firstStakeInputDatumF.delegatedTo
-                           in allHaveSameDelegatee
-                                #&& (delegated #|| ownedByDelegatee)
-                     in pcon $ PPair allHaveSameOwner' allHaveSameOrOwnedByDelegatee'
+                haveTheSameOwnerAsTheFirstStake <-
+                  pletC $
+                    stakeF.owner #== firstStakeInputDatumF.owner
+
+                let ownerOfTheFirstStakeIsTheDelegatee =
+                      haveTheSameOwnerAsTheFirstStake
+                        #|| stakeF.delegatedTo
+                        #== pdata (pdjust # firstStakeInputDatumF.owner)
+
+                    delegateeOfTheFirstStakeIsTheDelegatee =
+                      pdata (pdjust # stakeF.owner)
+                        #== firstStakeInputDatumF.delegatedTo
+                        #|| stakeF.delegatedTo
+                        #== firstStakeInputDatumF.delegatedTo
+
+                    helper' =
+                      pcon $
+                        PAuthCheckHelper
+                          ( helperF.allHaveTheSameOwner
+                              #&& haveTheSameOwnerAsTheFirstStake
+                          )
+                          ( helperF.allOwnByOrDelegatedToTheFirstOwner
+                              #&& ownerOfTheFirstStakeIsTheDelegatee
+                          )
+                          ( helperF.allOwnByOrDelegatedToTheFirstDelegatee
+                              #&& delegateeOfTheFirstStakeIsTheDelegatee
+                          )
+
+                pure helper'
+            )
+          # pcon
+            ( PAuthCheckHelper
+                (pconstant True)
+                (pconstant True)
+                (pconstant True)
             )
-          # pcon (PPair (pconstant True) (pconstant True))
           # restOfStakeInputDatums
 
     let ownerSignsTransaction =
-          allHaveSameOwner
+          authCheckHelperF.allHaveTheSameOwner
             #&& authorizedBy
             # firstStakeInputDatumF.owner
 
         delegateSignsTransaction =
-          allHaveSameOrOwnedByDelegatee
+          ( authCheckHelperF.allOwnByOrDelegatedToTheFirstOwner
+              #|| authCheckHelperF.allOwnByOrDelegatedToTheFirstDelegatee
+          )
             #&& pmaybeData
             # pconstant False
             # plam ((authorizedBy #) . pfromData)