From b834b4ebad576b79d3f2e07caad326d0899c3bea Mon Sep 17 00:00:00 2001
From: Emily Martins <emi@haskell.fyi>
Date: Thu, 17 Feb 2022 20:52:35 +0100
Subject: [PATCH] correctly check output datum and value

---
 src/Agora/SafeMoney.hs | 39 ++++++++++++++++++++
 src/Agora/Stake.hs     | 81 +++++++++++++++++++++++++++++++++---------
 2 files changed, 104 insertions(+), 16 deletions(-)

diff --git a/src/Agora/SafeMoney.hs b/src/Agora/SafeMoney.hs
index daa9553..1f8faa4 100644
--- a/src/Agora/SafeMoney.hs
+++ b/src/Agora/SafeMoney.hs
@@ -31,6 +31,7 @@ import Plutus.V1.Ledger.Value qualified as Ledger
 --------------------------------------------------------------------------------
 
 import Plutarch.Api.V1
+import Plutarch.Builtin
 import Plutarch.Internal
 import Plutarch.Prelude
 
@@ -106,6 +107,17 @@ plookup =
         PNothing -> pcon PNothing
         PJust p -> pcon (PJust (psndBuiltin # p))
 
+-- This is quite silly.
+plookupTuple ::
+  (PEq a, PIsListLike list (PAsData (PTuple a b)), PIsData a, PIsData b) =>
+  Term s (a :--> list (PAsData (PTuple a b)) :--> PMaybe b)
+plookupTuple =
+  phoistAcyclic $
+    plam $ \k xs ->
+      pmatch (pfind' (\p -> (pfield @"_0" # pfromData p) #== k) # xs) $ \case
+        PNothing -> pcon PNothing
+        PJust p -> pcon (PJust (pfield @"_1" # pfromData p))
+
 matchMaybe :: Term s r -> Term s (PMaybe a) -> TermCont @r s (Term s a)
 matchMaybe r f = TermCont $ \k ->
   pmatch f $ \case
@@ -142,3 +154,30 @@ valueDiscrete = phoistAcyclic $
       passetClassValueOf # (pconstant $ fromString $ symbolVal $ Proxy @ac)
         # (pconstant $ fromString $ symbolVal $ Proxy @n)
         # f
+
+-- NOTE: discreteValue after valueDiscrete is loses information
+
+-- | Get a 'PValue' from a 'Discrete'
+discreteValue ::
+  forall (moneyClass :: MoneyClass) (ac :: Symbol) (n :: Symbol) (scale :: Nat) s.
+  ( KnownSymbol ac
+  , KnownSymbol n
+  , moneyClass ~ '(ac, n, scale)
+  ) =>
+  Term s (Discrete moneyClass :--> PValue)
+discreteValue = phoistAcyclic $
+  plam $ \f -> pmatch f $ \case
+    Discrete p ->
+      psingletonValue
+        # (pconstant $ fromString $ symbolVal $ Proxy @ac)
+        # (pconstant $ fromString $ symbolVal $ Proxy @n)
+        # p
+
+-- | Create a value with a single asset class
+psingletonValue :: forall s. Term s (PCurrencySymbol :--> PTokenName :--> PInteger :--> PValue)
+psingletonValue = phoistAcyclic $
+  plam $ \sym tok int ->
+    let innerTup = pcon $ PMap $ psingleton #$ ppairDataBuiltin # pdata tok # pdata int
+        outerTup = pcon $ PMap $ psingleton #$ ppairDataBuiltin # pdata sym # pdata innerTup
+        res = pcon $ PValue outerTup
+     in res
diff --git a/src/Agora/Stake.hs b/src/Agora/Stake.hs
index a5519a3..34d7230 100644
--- a/src/Agora/Stake.hs
+++ b/src/Agora/Stake.hs
@@ -12,6 +12,9 @@ module Agora.Stake (
 --------------------------------------------------------------------------------
 
 import GHC.Generics qualified as GHC
+import GHC.TypeLits (
+  KnownSymbol,
+ )
 import Generics.SOP (Generic, I (I))
 import Prelude
 
@@ -63,28 +66,74 @@ newtype StakeDatum (gt :: MoneyClass) (s :: S) = StakeDatum
     (PlutusType, PIsData, PDataFields)
     via (PIsDataReprInstances (StakeDatum gt))
 
+-- | Assert a particular bool, trace on falsehood. Use in monadic context
 passert :: Term s PString -> Term s PBool -> Term s k -> Term s k
 passert errorMessage check k = pif check k (ptraceError errorMessage)
 
+-- | Find a datum with the given hash.
+pfindDatum :: Term s (PDatumHash :--> PTxInfo :--> PMaybe PDatum)
+pfindDatum = phoistAcyclic $
+  plam $ \datumHash txInfo'' -> P.do
+    PTxInfo txInfo' <- pmatch txInfo''
+    plookupTuple # datumHash #$ pfield @"data" # txInfo'
+
+-- | Find a datum with the given hash. NOTE: this is unsafe in the sense that, if the data layout is wrong, this is UB.
+pfindDatum' :: PIsData a => Term s (PDatumHash :--> PTxInfo :--> PMaybe (PAsData a))
+pfindDatum' = phoistAcyclic $ plam $ \dh x -> punsafeCoerce $ pfindDatum # dh # x
+
+-- | Check if a PubKeyHash signs this transaction
+ptxSignedBy :: Term s (PTxInfo :--> PAsData PPubKeyHash :--> PBool)
+ptxSignedBy = phoistAcyclic $
+  plam $ \txInfo' pkh -> P.do
+    txInfo <- pletFields @'["signatories"] txInfo'
+    pelem @PBuiltinList # pkh # txInfo.signatories
+
+-- | Check if any output matches the predicate
+anyOutput ::
+  forall (datum :: PType) s.
+  ( PIsData datum
+  ) =>
+  Term s (PTxInfo :--> (PValue :--> datum :--> PBool) :--> PBool)
+anyOutput = phoistAcyclic $
+  plam $ \txInfo' predicate -> P.do
+    txInfo <- pletFields @'["outputs"] txInfo'
+    pany
+      # ( plam $ \txOut'' -> P.do
+            PTxOut txOut' <- pmatch (pfromData txOut'')
+            txOut <- pletFields @'["value", "datumHash"] txOut'
+            PDJust dh <- pmatch txOut.datumHash
+            pmatch (pfindDatum' @datum # (pfield @"_0" # dh) # txInfo') $ \case
+              PJust datum -> P.do
+                predicate # txOut.value # pfromData datum
+              PNothing -> pcon PFalse
+        )
+      # pfromData txInfo.outputs
+
 stakePolicy ::
-  forall (gt :: MoneyClass) s.
-  Stake gt ->
+  forall (gt :: MoneyClass) ac n scale s.
+  ( KnownSymbol ac
+  , KnownSymbol n
+  , gt ~ '(ac, n, scale)
+  ) =>
+  Stake
+    gt ->
   Term s (PData :--> PScriptContext :--> PUnit)
 stakePolicy _stake =
-  plam $ \_redeemer ctx -> P.do
-    PScriptContext ctx' <- pmatch ctx
-    ctx'' <- pletFields @'["txInfo", "purpose"] ctx'
-    PTxInfo txInfo <- pmatch $ pfromData (hrecField @"txInfo" ctx'')
-    txInfo' <- pletFields @'["signatories", "outputs"] txInfo
-    let outputs = hrecField @"outputs" txInfo'
+  plam $ \_redeemer ctx'' -> P.do
+    PScriptContext ctx' <- pmatch ctx''
+    ctx <- pletFields @'["txInfo", "purpose"] ctx'
 
-    passert "Created stake must be owned by a signer of this transaction" $
-      pany
-        # ( plam $ \txOut -> P.do
-              PTxOut txOut' <- pmatch (pfromData txOut)
-              _txOut'' <- pletFields @'["value", "datumHash"] txOut'
-              pcon PTrue
+    PMinting ownSymbol <- pmatch $ pfromData ctx.purpose
+    -- TODO: add this to 'valueCorrect'
+    let _stValue = psingletonValue # (pfield @"_0" # ownSymbol) # pconstant "ST" # 1
+
+    passert "A UTXO must exist with the correct output" $
+      anyOutput @(StakeDatum gt) # pfromData ctx.txInfo
+        # ( plam $ \value stakeDatum' -> P.do
+              stakeDatum <- pletFields @'["owner", "stakedAmount"] stakeDatum'
+              let ownerSignsTransaction = ptxSignedBy # ctx.txInfo # stakeDatum.owner
+              let valueCorrect = pdata value #== pdata (discreteValue # stakeDatum.stakedAmount)
+              ownerSignsTransaction #&& valueCorrect
           )
-        # outputs
 
-    pcon PUnit
+    pconstant ()