Split mkRewardAccoundBuilder into smaller builders

lib/wallet/api/http/Cardano/Wallet/Api/Http/Shelley/Server.hs

@@ -638,6 +638,8 @@ import qualified Cardano.Wallet.Primitive.Types.UTxOSelection as UTxOSelection
 import qualified Cardano.Wallet.Registry as Registry
 import qualified Cardano.Wallet.Write.Tx as WriteTx
 import qualified Control.Concurrent.Concierge as Concierge
+import Control.Monad.Error.Class
+    ( throwError )
 import qualified Data.ByteString as BS
 import qualified Data.Foldable as F
 import qualified Data.List as L
@@ -1626,7 +1628,7 @@ selectCoins ctx@ApiLayer {..} genChange (ApiT wid) body = do
         wdrl <- case body ^. #withdrawal of
             Nothing -> pure NoWithdrawal
             Just apiWdrl ->
-                unsafeShelleyMkWithdrawal @s @k @n @'CredFromKeyK
+                shelleyOnlyMkWithdrawal @s @k @n @'CredFromKeyK
                     netLayer txLayer db wid era apiWdrl
         let outs = addressAmountToTxOut <$> body ^. #payments
         let txCtx = defaultTransactionCtx
@@ -1729,7 +1731,7 @@ selectCoinsForQuit ctx@ApiLayer{..} (ApiT wid) =
         let db = wrk ^. typed @(DBLayer IO s k)
         era <- liftIO $ NW.currentNodeEra netLayer
         wdrl <- liftHandler $ ExceptT
-            $ W.unsafeShelleyMkSelfWithdrawal @s @k @_ @_ @n
+            $ W.shelleyOnlyMkSelfWithdrawal @s @k @_ @_ @n
                 netLayer txLayer era db wid
         action <- liftIO $ W.validatedQuitStakePoolAction db wid wdrl
 
@@ -1758,7 +1760,7 @@ selectCoinsForQuit ctx@ApiLayer{..} (ApiT wid) =
             liftHandler $ W.selectAssets @_ @_ @s @k @'CredFromKeyK
                 wrk era pp selectAssetsParams transform
         (_, _, path) <- liftHandler
-            $ W.unsafeShelleyReadRewardAccount @s @k @n db wid
+            $ W.shelleyOnlyReadRewardAccount @s @k @n db wid
         let refund = W.stakeKeyDeposit pp
         pure $ mkApiCoinSelection [] [refund] (Just (action, path)) Nothing utx
 
@@ -1995,16 +1997,17 @@ postTransactionOld ctx@ApiLayer{..} genChange (ApiT wid) body = do
     let outs = addressAmountToTxOut <$> body ^. #payments
     let md = body ^? #metadata . traverse . #txMetadataWithSchema_metadata
     let mTTL = body ^? #timeToLive . traverse . #getQuantity
-
-    mkRwdAcct <- mkRewardAccountBuilder @s @_ @n (body ^. #withdrawal)
+    mkRwdAcct <- case body ^. #withdrawal of
+        Nothing -> pure selfRewardAccountBuilder
+        Just w -> either liftE pure $ shelleyOnlyRewardAccountBuilder @s @_ @n w
     withWorkerCtx ctx wid liftE liftE $ \wrk -> do
         let db = wrk ^. dbLayer
         era <- liftIO $ NW.currentNodeEra netLayer
         ttl <- liftIO $ W.transactionExpirySlot ti mTTL
         wdrl <- case body ^. #withdrawal of
             Nothing -> pure NoWithdrawal
             Just apiWdrl ->
-                unsafeShelleyMkWithdrawal @s @k @n
+                shelleyOnlyMkWithdrawal @s @k @n
                     netLayer txLayer db wid era apiWdrl
         let txCtx = defaultTransactionCtx
                 { txWithdrawal = wdrl
@@ -2194,7 +2197,7 @@ postTransactionFeeOld ctx@ApiLayer{..} (ApiT wid) body =
         wdrl <- case body ^. #withdrawal of
             Nothing -> pure NoWithdrawal
             Just apiWdrl ->
-                unsafeShelleyMkWithdrawal @s @k @n @'CredFromKeyK
+                shelleyOnlyMkWithdrawal @s @k @n @'CredFromKeyK
                     netLayer txLayer db wid era apiWdrl
         let txCtx = defaultTransactionCtx
                 { txWithdrawal = wdrl
@@ -2328,7 +2331,7 @@ constructTransaction ctx genChange knownPools getPoolStatus (ApiT wid) body = do
         era <- liftIO $ NW.currentNodeEra netLayer
         wdrl <- case body ^. #withdrawal of
             Just SelfWithdraw -> liftHandler
-                $ ExceptT $ W.unsafeShelleyMkSelfWithdrawal
+                $ ExceptT $ W.shelleyOnlyMkSelfWithdrawal
                     @s @k @'CredFromKeyK @_ @n netLayer txLayer era db wid
             _ -> pure NoWithdrawal
         (deposit, refund, txCtx) <- case body ^. #delegations of
@@ -2865,7 +2868,7 @@ decodeTransaction
     withWorkerCtx ctx wid liftE liftE $ \wrk -> do
         let db = wrk ^. dbLayer
         (acct, _, acctPath) <-
-            liftHandler $ W.unsafeShelleyReadRewardAccount @s @k @n db wid
+            liftHandler $ W.shelleyOnlyReadRewardAccount @s @k @n db wid
         inputPaths <-
             liftHandler $ W.lookupTxIns @_ @s @k wrk wid $
             fst <$> resolvedInputs
@@ -3001,7 +3004,7 @@ submitTransaction ctx apiw@(ApiT wid) apitx = do
 
     _ <- withWorkerCtx ctx wid liftE liftE $ \wrk -> do
         let db = wrk ^. dbLayer
-        (acct, _, path) <- liftHandler $ W.unsafeShelleyReadRewardAccount @s @k @n db wid
+        (acct, _, path) <- liftHandler $ W.shelleyOnlyReadRewardAccount @s @k @n db wid
         let wdrl = getOurWdrl acct path apiDecoded
         let txCtx = defaultTransactionCtx
                 { -- TODO: [ADP-1193]
@@ -3196,11 +3199,10 @@ joinStakePool ctx knownPools getPoolStatus apiPool (ApiT wid) body = do
                 wrk era pp selectAssetsParams (const Prelude.id)
         sel' <- liftHandler
             $ W.assignChangeAddressesAndUpdateDb wrk wid genChange sel
-        mkRwdAcct <- mkRewardAccountBuilder @s @_ @n Nothing
         (tx, txMeta, txTime, sealedTx) <- liftHandler $ do
             let pwd = coerce $ getApiT $ body ^. #passphrase
             W.buildAndSignTransaction @_ @s @k
-                wrk wid era mkRwdAcct pwd txCtx sel'
+                wrk wid era selfRewardAccountBuilder pwd txCtx sel'
         liftHandler $ W.submitTx @_ @s @k wrk wid (tx, txMeta, sealedTx)
         mkApiTransaction
             (timeInterpreter (ctx ^. networkLayer))
@@ -3287,7 +3289,6 @@ quitStakePool
     -> ApiWalletPassphrase
     -> Handler (ApiTransaction n)
 quitStakePool ctx@ApiLayer{..} (ApiT walletId) body = do
-    mkRwdAcct <- mkRewardAccountBuilder @s @_ @n (Just SelfWithdrawal)
     withWorkerCtx ctx walletId liftE liftE $ \wrk -> do
         let db = wrk ^. typed @(DBLayer IO s k)
             notShelleyWallet =
@@ -3323,7 +3324,7 @@ quitStakePool ctx@ApiLayer{..} (ApiT walletId) body = do
         (tx, txMeta, txTime, sealedTx) <- do
             let pwd = coerce $ getApiT $ body ^. #passphrase
             liftHandler $ W.buildAndSignTransaction @_ @s @k
-                wrk walletId era mkRwdAcct pwd txCtx sel'
+                wrk walletId era selfRewardAccountBuilder pwd txCtx sel'
         liftHandler
             $ W.submitTx @_ @s @k wrk walletId (tx, txMeta, sealedTx)
         mkApiTransaction ti wrk walletId #pendingSince
@@ -3459,7 +3460,7 @@ createMigrationPlan ctx@ApiLayer{..} withdrawalType (ApiT wid) postData =
         era <- liftIO $ NW.currentNodeEra netLayer
         rewardWithdrawal <- case withdrawalType of
             Nothing -> pure NoWithdrawal
-            Just pd -> unsafeShelleyMkWithdrawal @s @k @n @'CredFromKeyK
+            Just pd -> shelleyOnlyMkWithdrawal @s @k @n @'CredFromKeyK
                 netLayer txLayer db wid era pd
         (wallet, _, _) <- liftHandler
             $ withExceptT ErrCreateMigrationPlanNoSuchWallet
@@ -3551,13 +3552,17 @@ migrateWallet
     -> ApiWalletMigrationPostData n p
     -> Handler (NonEmpty (ApiTransaction n))
 migrateWallet ctx@ApiLayer{..} withdrawalType (ApiT wid) postData = do
-    mkRewardAccount <- mkRewardAccountBuilder @s @_ @n withdrawalType
+    mkRewardAccount <-
+        case withdrawalType of
+            Nothing -> pure selfRewardAccountBuilder
+            Just w ->
+                either liftE pure $ shelleyOnlyRewardAccountBuilder @s @_ @n w
     withWorkerCtx ctx wid liftE liftE $ \wrk -> do
         let db = wrk ^. dbLayer
         era <- liftIO $ NW.currentNodeEra netLayer
         rewardWithdrawal <- case withdrawalType of
             Nothing -> pure NoWithdrawal
-            Just pd -> unsafeShelleyMkWithdrawal @s @k @n
+            Just pd -> shelleyOnlyMkWithdrawal @s @k @n
                 netLayer txLayer db wid era pd
         plan <- liftHandler $ W.createMigrationPlan wrk era wid rewardWithdrawal
         ttl <- liftIO $ W.transactionExpirySlot ti Nothing
@@ -3918,7 +3923,7 @@ mkWithdrawal netLayer txLayer db wallet era = \case
 
 -- | Unsafe version of `mkWithdrawal` that throws runtime error
 -- when applied to a non-shelley or non-sequential wallet state.
-unsafeShelleyMkWithdrawal
+shelleyOnlyMkWithdrawal
     :: forall s k (n :: NetworkDiscriminant) ktype tx
      . (Typeable n, Typeable s, Typeable k)
     => NetworkLayer IO Block
@@ -3928,7 +3933,7 @@ unsafeShelleyMkWithdrawal
     -> AnyCardanoEra
     -> ApiWithdrawalPostData
     -> Handler Withdrawal
-unsafeShelleyMkWithdrawal netLayer txLayer db wallet era postData =
+shelleyOnlyMkWithdrawal netLayer txLayer db wallet era postData =
     case testEquality (typeRep @s) (typeRep @(SeqState n k)) of
         Nothing -> notShelleyWallet
         Just Refl -> case testEquality (typeRep @k) (typeRep @ShelleyKey) of
@@ -3938,28 +3943,34 @@ unsafeShelleyMkWithdrawal netLayer txLayer db wallet era postData =
     notShelleyWallet =
         liftHandler $ throwE ErrReadRewardAccountNotAShelleyWallet
 
-mkRewardAccountBuilder
+shelleyOnlyRewardAccountBuilder
     :: forall s k (n :: NetworkDiscriminant)
      . ( HardDerivation k
        , Bounded (Index (AddressIndexDerivationType k) (AddressCredential k))
        , WalletKey k
        , Typeable s
        , Typeable n
        )
-    => Maybe ApiWithdrawalPostData
-    -> Handler (RewardAccountBuilder k)
-mkRewardAccountBuilder withdrawal = do
-    let selfRewardCredentials (rootK, pwdP) =
-            (getRawKey (deriveRewardAccount @k pwdP rootK), pwdP)
+    => ApiWithdrawalPostData
+    -> Either ErrReadRewardAccount (RewardAccountBuilder k)
+shelleyOnlyRewardAccountBuilder w =
     case testEquality (typeRep @s) (typeRep @(SeqState n ShelleyKey)) of
-        Nothing -> liftHandler $ throwE ErrReadRewardAccountNotAShelleyWallet
-        Just Refl -> case withdrawal of
-            Nothing -> pure selfRewardCredentials
-            Just w -> case w of
-                SelfWithdrawal -> pure selfRewardCredentials
-                ExternalWithdrawal (ApiMnemonicT m) -> do
-                    let (xprv, _acct, _path) = W.someRewardAccount @ShelleyKey m
-                    pure (const (xprv, mempty))
+        Nothing -> throwError ErrReadRewardAccountNotAShelleyWallet
+        Just Refl -> case w of
+            SelfWithdrawal -> pure selfRewardAccountBuilder
+            ExternalWithdrawal (ApiMnemonicT m) -> do
+                let (xprv, _acct, _path) = W.someRewardAccount @ShelleyKey m
+                pure (const (xprv, mempty))
+
+selfRewardAccountBuilder
+    :: forall k
+     . ( HardDerivation k
+       , Bounded (Index (AddressIndexDerivationType k) (AddressCredential k))
+       , WalletKey k
+       )
+    => RewardAccountBuilder k
+selfRewardAccountBuilder (rootK, pwdP) =
+    (getRawKey (deriveRewardAccount @k pwdP rootK), pwdP)
 
 -- | Makes an 'ApiCoinSelection' from the given 'UnsignedTx'.
 mkApiCoinSelection
@@ -4161,7 +4172,7 @@ mkApiTransaction timeInterpreter wrk wid timeRefLens tx = do
     -- using additional context from the 'WorkerCtx'.
     getApiAnyCertificates db ParsedTxCBOR{certificates} = do
         (rewardAccount, _, derivPath) <- liftHandler
-            $ W.unsafeShelleyReadRewardAccount @s @k @n db wid
+            $ W.shelleyOnlyReadRewardAccount @s @k @n db wid
         pure $ mkApiAnyCertificate rewardAccount derivPath <$> certificates
 
     depositIfAny :: Natural

lib/wallet/api/http/Cardano/Wallet/Shelley.hs

@@ -168,8 +168,6 @@ import Network.Wai.Handler.Warp
     ( setBeforeMainLoop )
 import Ouroboros.Network.Client.Wallet
     ( PipeliningStrategy )
-import Servant.Server
-    ( ServerError )
 import System.Exit
     ( ExitCode (..) )
 import System.IOManager

lib/wallet/src/Cardano/Wallet.hs

@@ -81,9 +81,9 @@ module Cardano.Wallet
     , checkWalletIntegrity
     , mkExternalWithdrawal
     , mkSelfWithdrawal
-    , unsafeShelleyMkSelfWithdrawal
+    , shelleyOnlyMkSelfWithdrawal
     , readRewardAccount
-    , unsafeShelleyReadRewardAccount
+    , shelleyOnlyReadRewardAccount
     , someRewardAccount
     , readPolicyPublicKey
     , writePolicyPublicKey
@@ -1238,7 +1238,7 @@ mkSelfWithdrawal netLayer txLayer era db wallet = do
 
 -- | Unsafe version of the `mkSelfWithdrawal` function that throws an exception
 -- when applied to a non-shelley or a non-sequential wallet.
-unsafeShelleyMkSelfWithdrawal
+shelleyOnlyMkSelfWithdrawal
     :: forall s k ktype tx (n :: NetworkDiscriminant)
      . (Typeable s, Typeable k, Typeable n)
     => NetworkLayer IO Block
@@ -1247,7 +1247,7 @@ unsafeShelleyMkSelfWithdrawal
     -> DBLayer IO s k
     -> WalletId
     -> IO (Either ErrWithdrawalNotWorth Withdrawal)
-unsafeShelleyMkSelfWithdrawal netLayer txLayer era db wallet =
+shelleyOnlyMkSelfWithdrawal netLayer txLayer era db wallet =
     case testEquality (typeRep @s) (typeRep @(SeqState n k)) of
         Nothing -> notShelleyWallet
         Just Refl -> case testEquality (typeRep @k) (typeRep @ShelleyKey) of
@@ -1301,14 +1301,14 @@ readRewardAccount db wid = do
 -- | Unsafe version of the `readRewardAccount` function
 -- that throws error when applied to a non-sequential
 -- or a non-shelley wallet state.
-unsafeShelleyReadRewardAccount
+shelleyOnlyReadRewardAccount
     :: forall s k (n :: NetworkDiscriminant)
      . (Typeable s, Typeable n, Typeable k)
     => DBLayer IO s k
     -> WalletId
     -> ExceptT ErrReadRewardAccount IO
         (RewardAccount, XPub, NonEmpty DerivationIndex)
-unsafeShelleyReadRewardAccount db wid =
+shelleyOnlyReadRewardAccount db wid =
     case testEquality (typeRep @s) (typeRep @(SeqState n k)) of
         Nothing -> throwE ErrReadRewardAccountNotAShelleyWallet
         Just Refl ->
@@ -2529,7 +2529,7 @@ constructTransaction
     -> ExceptT ErrConstructTx IO SealedTx
 constructTransaction ctx wid era txCtx sel = db & \DBLayer{..} -> do
     (_, xpub, _) <- withExceptT ErrConstructTxReadRewardAccount $
-        unsafeShelleyReadRewardAccount @s @k @n db wid
+        shelleyOnlyReadRewardAccount @s @k @n db wid
     mapExceptT atomically $ do
         pp <- liftIO $ currentProtocolParameters nl
         withExceptT ErrConstructTxBody $ ExceptT $ pure $