From 3f845ff30fa660f4bdd29e8370d3cdf30a9a4330 Mon Sep 17 00:00:00 2001
From: Pawel Jakubas <pawel.jakubas@iohk.io>
Date: Tue, 16 May 2023 16:23:47 +0200
Subject: [PATCH] deal with withdrawal script wits

---
 .../src/Cardano/Wallet/Shelley/Transaction.hs | 46 +++++++++++++++----
 1 file changed, 36 insertions(+), 10 deletions(-)

diff --git a/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs b/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs
index cb0703a8504..fbb88124121 100644
--- a/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs
+++ b/lib/wallet/src/Cardano/Wallet/Shelley/Transaction.hs
@@ -81,8 +81,10 @@ import Cardano.Address.Script
     , KeyHash (..)
     , KeyRole (..)
     , Script (..)
+    , ScriptHash (..)
     , ScriptTemplate (..)
     , foldScript
+    , toScriptHash
     )
 import Cardano.Api
     ( AnyCardanoEra (..)
@@ -312,6 +314,7 @@ import qualified Data.Sequence.Strict as StrictSeq
 import qualified Data.Set as Set
 import qualified Data.Text as T
 
+
 -- | Type encapsulating what we need to know to add things -- payloads,
 -- certificates -- to a transaction.
 --
@@ -664,10 +667,24 @@ newTransactionLayer keyF networkId = TransactionLayer
         case view #txDelegationAction ctx of
             Nothing -> do
                 let md = view #txMetadata ctx
-                constructUnsignedTx networkId (md, []) ttl wdrl
-                    selection delta assetsToBeMinted assetsToBeBurned inpsScripts
-                    Nothing
-                    (Write.shelleyBasedEraFromRecentEra Write.recentEra)
+                let ourRewardAcctM = FromScriptHash . unScriptHash . toScriptHash <$> stakingScriptM
+                case wdrl of
+                    WithdrawalSelf rewardAcct _ _ ->
+                        if ourRewardAcctM == Just rewardAcct then
+                            constructUnsignedTx networkId (md, []) ttl wdrl
+                            selection delta assetsToBeMinted assetsToBeBurned inpsScripts
+                            stakingScriptM
+                            (Write.shelleyBasedEraFromRecentEra Write.recentEra)
+                        else
+                            constructUnsignedTx networkId (md, []) ttl wdrl
+                            selection delta assetsToBeMinted assetsToBeBurned inpsScripts
+                            Nothing
+                            (Write.shelleyBasedEraFromRecentEra Write.recentEra)
+                    _ ->
+                        constructUnsignedTx networkId (md, []) ttl wdrl
+                        selection delta assetsToBeMinted assetsToBeBurned inpsScripts
+                        Nothing
+                        (Write.shelleyBasedEraFromRecentEra Write.recentEra)
             Just action -> do
                 let certs = case stakeCred of
                         Left xpub ->
@@ -2259,12 +2276,21 @@ mkUnsignedTx
     , txInsReference = Cardano.TxInsReferenceNone
 
     , Cardano.txOuts = map (toCardanoTxOut era) outs
-    , Cardano.txWithdrawals =
-        let wit = Cardano.BuildTxWith
-                $ Cardano.KeyWitness Cardano.KeyWitnessForStakeAddr
-        in
-        Cardano.TxWithdrawals wdrlsSupported
-            (map (\(key, coin) -> (key, coin, wit)) wdrls)
+    , Cardano.txWithdrawals = case stakingScriptM of
+        Nothing ->
+            let ctx = Cardano.BuildTxWith
+                    $ Cardano.KeyWitness Cardano.KeyWitnessForStakeAddr
+            in
+                Cardano.TxWithdrawals wdrlsSupported
+                (map (\(key, coin) -> (key, coin, ctx)) wdrls)
+        Just stakingScript ->
+            let
+                buildVal = Cardano.ScriptWitness Cardano.ScriptWitnessForStakeAddr
+                    (toScriptWitness stakingScript)
+                ctx = Cardano.BuildTxWith buildVal
+            in
+                Cardano.TxWithdrawals wdrlsSupported
+                (map (\(key, coin) -> (key, coin, ctx)) wdrls)
 
     -- @mkUnsignedTx@ is never used with Plutus scripts, and so we never have to
     -- care about collateral or PParams (for script integrity hash) here.