From 6680882cd5b83ca244d734474b849c161d9b5447 Mon Sep 17 00:00:00 2001
From: Sebastian Nagel <sebastian.nagel@ncoding.at>
Date: Tue, 3 May 2022 12:20:39 +0200
Subject: [PATCH] Implement checkHeadOutputDatum

---
 hydra-plutus/src/Hydra/Contract/Head.hs | 40 ++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 7 deletions(-)

diff --git a/hydra-plutus/src/Hydra/Contract/Head.hs b/hydra-plutus/src/Hydra/Contract/Head.hs
index bf133fdf4f5..3288acf3a8f 100644
--- a/hydra-plutus/src/Hydra/Contract/Head.hs
+++ b/hydra-plutus/src/Hydra/Contract/Head.hs
@@ -46,7 +46,7 @@ import Plutus.V1.Ledger.Api (
   adaToken,
   mkValidatorScript,
  )
-import Plutus.V1.Ledger.Contexts (findDatum, findOwnInput)
+import Plutus.V1.Ledger.Contexts (findDatum, findDatumHash, findOwnInput, getContinuingOutputs)
 import Plutus.V1.Ledger.Crypto (Signature (getSignature))
 import Plutus.V1.Ledger.Value (valueOf)
 import PlutusTx (CompiledCode)
@@ -269,6 +269,14 @@ checkCollectCom context@ScriptContext{scriptContextTxInfo = txInfo} headContext
         traceError "fromBuiltinData failed"
 {-# INLINEABLE checkCollectCom #-}
 
+-- | The close validator must verify that:
+--
+--   * The closing snapshot number and signature is correctly signed
+--
+--   * The resulting closed state is consistent with the open state or the
+--     closing snapshot, depending on snapshot number
+--
+--   * The transaction is performed (i.e. signed) by one of the head participants
 checkClose ::
   ScriptContext ->
   HeadContext ->
@@ -278,19 +286,37 @@ checkClose ::
   BuiltinByteString ->
   [Signature] ->
   Bool
-checkClose context headContext parties initialUtxoHash snapshotNumber closedUtxoHash sig =
-  checkSnapshot && mustBeSignedByParticipant context headContext
+checkClose ctx headContext parties initialUtxoHash snapshotNumber closedUtxoHash sig =
+  checkSnapshot && mustBeSignedByParticipant ctx headContext
  where
   checkSnapshot
-    | snapshotNumber == 0 = checkHeadOutputDatum (Closed 0 initialUtxoHash)
+    | snapshotNumber == 0 = checkHeadOutputDatum ctx (Closed 0 initialUtxoHash)
     | snapshotNumber > 0 =
       verifySnapshotSignature parties snapshotNumber closedUtxoHash sig
-        && checkHeadOutputDatum (Closed snapshotNumber closedUtxoHash)
+        && checkHeadOutputDatum ctx (Closed snapshotNumber closedUtxoHash)
     | otherwise = traceError "negative snapshot number"
-
-  checkHeadOutputDatum _datum = traceError "not implemented"
 {-# INLINEABLE checkClose #-}
 
+checkHeadOutputDatum :: ToData a => ScriptContext -> a -> Bool
+checkHeadOutputDatum ctx d =
+  case (ownDatumHash, expectedDatumHash) of
+    (Just actual, Just expected) ->
+      traceIfFalse "output datum hash mismatch" $ actual == expected
+    (Nothing, _) ->
+      traceError "no head output datum"
+    (_, Nothing) ->
+      traceError "expected datum hash not found"
+ where
+  expectedDatumHash = findDatumHash (Datum $ toBuiltinData d) txInfo
+
+  ownDatumHash =
+    case getContinuingOutputs ctx of
+      [o] -> txOutDatumHash o
+      _ -> traceError "expected only one head output"
+
+  ScriptContext{scriptContextTxInfo = txInfo} = ctx
+{-# INLINEABLE checkHeadOutputDatum #-}
+
 txOutAdaValue :: TxOut -> Integer
 txOutAdaValue o = valueOf (txOutValue o) adaSymbol adaToken
 {-# INLINEABLE txOutAdaValue #-}