chore: add unit tests for multi-asset seize

Author: colll78

src/examples/regulated-stablecoin/lib/Wst/Offchain/BuildTx/ProgrammableLogic.hs

@@ -11,18 +11,22 @@ module Wst.Offchain.BuildTx.ProgrammableLogic
 where
 
 import Cardano.Api qualified as C
+import Cardano.Api.Internal.Tx.Body qualified as C
 import Cardano.Api.Shelley qualified as C
 import Control.Lens ((^.))
+import Control.Lens qualified as L
+import Control.Monad (forM_)
 import Control.Monad.Reader (MonadReader, asks)
-import Convex.BuildTx (MonadBuildTx, addReference, addWithdrawalWithTxBody,
-                       buildScriptWitness, findIndexReference,
-                       findIndexSpending, prependTxOut, spendPlutusInlineDatum)
+import Convex.BuildTx (MonadBuildTx, addOutput, addReference,
+                       addWithdrawalWithTxBody, buildScriptWitness,
+                       findIndexReference, findIndexSpending,
+                       spendPlutusInlineDatum)
 import Convex.CardanoApi.Lenses as L
 import Convex.Class (MonadBlockchain (queryNetworkId))
 import Convex.PlutusLedger.V1 (transPolicyId)
 import Convex.Utils qualified as Utils
 import Data.Foldable (find)
-import Data.List (partition)
+import Data.List (findIndex, partition)
 import Data.Maybe (fromJust)
 import GHC.Exts (IsList (..))
 import PlutusLedgerApi.V3 (CurrencySymbol (..))
@@ -44,65 +48,84 @@ import Wst.Offchain.Query (UTxODat (..))
    NOTE: Seems the issuer is only able to seize 1 UTxO at a time.
    In the future we should allow multiple UTxOs in 1 Tx.
 -}
-seizeProgrammableToken :: forall a env era m. (MonadReader env m, Env.HasDirectoryEnv env, C.IsBabbageBasedEra era, MonadBlockchain era m, C.HasScriptLanguageInEra C.PlutusScriptV3 era, MonadBuildTx era m) => UTxODat era ProgrammableLogicGlobalParams -> UTxODat era a -> C.PolicyId -> [UTxODat era DirectorySetNode] -> m ()
-seizeProgrammableToken UTxODat{uIn = paramsTxIn} UTxODat{uIn = seizingTxIn, uOut = seizingTxOut} seizingTokenPolicyId directoryList = Utils.inBabbage @era $ do
+seizeProgrammableToken ::
+  forall a env era m.
+  ( MonadReader env m,
+    Env.HasDirectoryEnv env,
+    C.IsBabbageBasedEra era,
+    MonadBlockchain era m,
+    C.HasScriptLanguageInEra C.PlutusScriptV3 era,
+    MonadBuildTx era m
+  ) =>
+  UTxODat era ProgrammableLogicGlobalParams ->
+  [UTxODat era a] ->
+  C.PolicyId ->
+  [UTxODat era DirectorySetNode] ->
+  m ()
+seizeProgrammableToken UTxODat{uIn = paramsTxIn} seizingUTxOs seizingTokenPolicyId directoryList = Utils.inBabbage @era $ do
   nid <- queryNetworkId
   globalStakeScript <- asks (Env.dsProgrammableLogicGlobalScript . Env.directoryEnv)
   baseSpendingScript <- asks (Env.dsProgrammableLogicBaseScript . Env.directoryEnv)
 
   let globalStakeCred = C.StakeCredentialByScript $ C.hashScript $ C.PlutusScript C.PlutusScriptV3 globalStakeScript
+      programmableLogicBaseCredential = C.PaymentCredentialByScript $ C.hashScript $ C.PlutusScript C.PlutusScriptV3 baseSpendingScript
 
   -- Finds the directory node entry that references the programmable token symbol
   dirNodeRef <-
     maybe (error "Cannot seize non-programmable token. Entry does not exist in directoryList") (pure . uIn) $
       find (isNodeWithProgrammableSymbol (transPolicyId seizingTokenPolicyId)) directoryList
 
   -- destStakeCred <- either (error . ("Could not unTrans credential: " <>) . show) pure $ unTransStakeCredential $ transCredential seizeDestinationCred
-  let
-      -- issuerDestinationAddress = C.makeShelleyAddressInEra C.shelleyBasedEra nid progLogicBaseCred (C.StakeAddressByValue destStakeCred)
-
-      (seizedAddr, remainingValue) = case seizingTxOut of
-        (C.TxOut a v _ _) ->
-          let (_seized, other) =
-                  partition
-                    ( \case
-                        (C.AdaAssetId, _q) -> False
-                        (C.AssetId a_ _, _q) -> a_ == seizingTokenPolicyId
-                    )
-                    $ toList $ C.txOutValueToValue v
-          in (a, fromList other)
-
-      remainingTxOutValue = C.TxOutValueShelleyBased C.shelleyBasedEra $ C.toLedgerValue @era C.maryBasedEra remainingValue
 
-      seizedOutput = C.TxOut seizedAddr remainingTxOutValue C.TxOutDatumNone C.ReferenceScriptNone
+  forM_ seizingUTxOs $ \UTxODat{uIn = seizingTxIn, uOut = seizingTxOut} -> do
+    spendPlutusInlineDatum seizingTxIn baseSpendingScript ()
+    let (seizedAddr, remainingValue, seizedDatum, referenceScript) = case seizingTxOut of
+          (C.TxOut a v dat refScript) ->
+            let (_seized, other) =
+                    partition
+                      ( \case
+                          (C.AdaAssetId, _q) -> False
+                          (C.AssetId a_ _, _q) -> a_ == seizingTokenPolicyId
+                      )
+                      $ toList $ C.txOutValueToValue v
+            in (a, fromList other, dat, refScript)
+        remainingTxOutValue = C.TxOutValueShelleyBased C.shelleyBasedEra $ C.toLedgerValue @era C.maryBasedEra remainingValue
+        seizedOutput = C.TxOut seizedAddr remainingTxOutValue seizedDatum referenceScript
+    addOutput (C.fromCtxUTxOTxOut seizedOutput)
 
+  let
       -- Finds the index of the directory node reference in the transaction ref
       -- inputs
       directoryNodeReferenceIndex txBody =
         fromIntegral @Int @Integer $ findIndexReference dirNodeRef txBody
 
       -- Finds the index of the issuer input in the transaction body
       seizingInputIndex txBody =
-        fromIntegral @Int @Integer $ findIndexSpending seizingTxIn txBody
-
-      -- Finds the index of the issuer seized output in the transaction body
-      seizingOutputIndex txBody =
-        fromIntegral @Int @Integer $ fst $ fromJust (find ((== seizedOutput) . snd ) $ zip [0 ..] $ txBody ^. L.txOuts)
+        map (\UTxODat{uIn = seizingTxIn} -> fromIntegral @Int @Integer $ findIndexSpending seizingTxIn txBody) seizingUTxOs
+
+      -- Finds the index of the first output to the programmable logic base credential
+      firstSeizeContinuationOutputIndex txBody =
+        fromIntegral @Int @Integer $
+          fromJust $
+            findIndex
+              ( maybe False ((== programmableLogicBaseCredential) . C.fromShelleyPaymentCredential)
+                  . L.preview (L._TxOut . L._1 . L._AddressInEra . L._Address . L._2)
+              )
+              (txBody ^. L.txOuts)
 
       -- The seizing redeemer for the global script
       programmableLogicGlobalRedeemer txBody =
         SeizeAct
-          { plgrSeizeInputIdx = seizingInputIndex txBody,
-            plgrSeizeOutputIdx = seizingOutputIndex txBody,
-            plgrDirectoryNodeIdx = directoryNodeReferenceIndex txBody
+          { plgrDirectoryNodeIdx = directoryNodeReferenceIndex txBody,
+            plgrInputIdxs = seizingInputIndex txBody,
+            plgrOutputsStartIdx = firstSeizeContinuationOutputIndex txBody,
+            plgrLengthInputIdxs = fromIntegral @Int @Integer $ length seizingUTxOs
           }
 
       programmableGlobalWitness txBody = buildScriptWitness globalStakeScript C.NoScriptDatumForStake (programmableLogicGlobalRedeemer txBody)
 
-  prependTxOut seizedOutput
   addReference paramsTxIn -- Protocol Params TxIn
   addReference dirNodeRef -- Directory Node TxIn
-  spendPlutusInlineDatum seizingTxIn baseSpendingScript () -- Redeemer is ignored in programmableLogicBase
   addWithdrawalWithTxBody -- Add the global script witness to the transaction
     (C.makeStakeAddress nid globalStakeCred)
     (C.Quantity 0)

src/examples/regulated-stablecoin/lib/Wst/Offchain/BuildTx/TransferLogic.hs

@@ -9,6 +9,7 @@ module Wst.Offchain.BuildTx.TransferLogic
     issueSmartTokens,
     SeizeReason(..),
     seizeSmartTokens,
+    multiSeizeSmartTokens,
     initBlacklist,
     BlacklistReason(..),
     insertBlacklistNode,
@@ -20,15 +21,17 @@ where
 import Cardano.Api qualified as C
 import Cardano.Api.Shelley qualified as C
 import Control.Lens (at, over, set, (&), (?~), (^.))
+import Control.Lens qualified as L
 import Control.Monad (when)
 import Control.Monad.Error.Lens (throwing_)
 import Control.Monad.Except (MonadError)
 import Control.Monad.Reader (MonadReader, asks)
 import Convex.BuildTx (MonadBuildTx (addTxBuilder), TxBuilder (TxBuilder),
-                       addBtx, addRequiredSignature, addScriptWithdrawal,
-                       addWithdrawalWithTxBody, buildScriptWitness,
-                       findIndexReference, mintPlutus, payToAddress,
-                       prependTxOut, spendPlutusInlineDatum)
+                       addBtx, addOutput, addRequiredSignature,
+                       addScriptWithdrawal, addWithdrawalWithTxBody,
+                       buildScriptWitness, findIndexReference, mintPlutus,
+                       payToAddress, payToAddressTxOut, prependTxOut,
+                       spendPlutusInlineDatum)
 import Convex.CardanoApi.Lenses qualified as L
 import Convex.Class (MonadBlockchain (queryNetworkId))
 import Convex.PlutusLedger.V1 (transCredential, transPolicyId,
@@ -282,7 +285,7 @@ seizeSmartTokens reason paramsTxIn seizingTxo destinationCred directoryList = Ut
             (toList $ C.txOutValueToValue v)
 
   (progTokenPolId, an, q) <- maybe (error "No programmable token found in seizing transaction") pure maybeProgAsset
-  seizeProgrammableToken paramsTxIn seizingTxo progTokenPolId directoryList
+  seizeProgrammableToken paramsTxIn [seizingTxo] progTokenPolId directoryList
   addSeizeWitness
 
   progLogicBaseCred <- asks (Env.programmableLogicBaseCredential . Env.directoryEnv)
@@ -295,7 +298,36 @@ seizeSmartTokens reason paramsTxIn seizingTxo destinationCred directoryList = Ut
 
   addSeizeReason reason
   -- Send seized funds to destinationCred
-  payToAddress destinationAddress seizedVal
+  addOutput $ payToAddressTxOut destinationAddress seizedVal
+
+-- This function should probably accept the programmable token policy which we want to seize as a parameter.
+-- As of now, it just assumes that the first non-ada token in the seizing inputs is the token we want to seize.
+multiSeizeSmartTokens :: forall env era a m. (MonadReader env m, Env.HasOperatorEnv era env, Env.HasTransferLogicEnv env, Env.HasDirectoryEnv env, C.IsBabbageBasedEra era, MonadBlockchain era m, C.HasScriptLanguageInEra C.PlutusScriptV3 era, MonadBuildTx era m) => SeizeReason -> UTxODat era ProgrammableLogicGlobalParams -> C.PolicyId -> [UTxODat era a] -> C.PaymentCredential -> [UTxODat era DirectorySetNode] -> m ()
+multiSeizeSmartTokens reason paramsTxIn toSeizePolicyId seizingTxos destinationCred directoryList = Utils.inBabbage @era $ do
+  nid <- queryNetworkId
+
+  let seizedValue = C.policyAssetsToValue toSeizePolicyId $
+        foldl (\acc (uOut -> utxoDat) ->
+          let filteredAssets = foldMap ( \(pid, assets) ->
+                if pid == toSeizePolicyId then assets else mempty
+                )
+                (toList $ C.valueToPolicyAssets (L.view (L._TxOut . L._2 . L._TxOutValue) utxoDat))
+          in acc <> filteredAssets
+          )
+          (mempty :: C.PolicyAssets)
+          seizingTxos
+
+  seizeProgrammableToken paramsTxIn seizingTxos toSeizePolicyId directoryList
+  addSeizeWitness
+
+  progLogicBaseCred <- asks (Env.programmableLogicBaseCredential . Env.directoryEnv)
+  destStakeCred <- either (error . ("Could not unTrans credential: " <>) . show) pure $ unTransStakeCredential $ transCredential destinationCred
+  let
+      destinationAddress = C.makeShelleyAddressInEra C.shelleyBasedEra nid progLogicBaseCred (C.StakeAddressByValue destStakeCred)
+
+  addSeizeReason reason
+  -- Send seized funds to destinationCred
+  addOutput $ payToAddressTxOut destinationAddress seizedValue
 
 addIssueWitness :: forall era env m. (MonadReader env m, Env.HasOperatorEnv era env, Env.HasTransferLogicEnv env, C.IsBabbageBasedEra era, MonadBlockchain era m, C.HasScriptLanguageInEra C.PlutusScriptV3 era, MonadBuildTx era m) => m ()
 addIssueWitness = Utils.inBabbage @era $ do

src/examples/regulated-stablecoin/lib/Wst/Offchain/Endpoints/Deployment.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE NamedFieldPuns #-}
 {-| Deploy the directory and global params
 -}
 module Wst.Offchain.Endpoints.Deployment(
@@ -9,6 +8,7 @@ module Wst.Offchain.Endpoints.Deployment(
   insertBlacklistNodeTx,
   removeBlacklistNodeTx,
   seizeCredentialAssetsTx,
+  seizeMultiCredentialAssetsTx
 ) where
 
 import Cardano.Api (Quantity)
@@ -25,7 +25,9 @@ import Data.Function (on)
 import GHC.IsList (IsList (..))
 import ProgrammableTokens.OffChain.BuildTx qualified as BuildTx
 import ProgrammableTokens.OffChain.Env.Operator qualified as Env
+import ProgrammableTokens.OffChain.Env.TransferLogic (programmableTokenPolicyId)
 import ProgrammableTokens.OffChain.Error (AsProgrammableTokensError (..))
+import ProgrammableTokens.OffChain.Query (utxoHasPolicyId)
 import ProgrammableTokens.OffChain.Query qualified as Query
 import SmartTokens.Core.Scripts (ScriptTarget (..))
 import Wst.AppError (AsRegulatedStablecoinError (..))
@@ -175,3 +177,36 @@ seizeCredentialAssetsTx reason sanctionedCred = do
   (tx, _) <- Env.balanceTxEnv_ $ do
     BuildTx.seizeSmartTokens reason paramsTxIn seizeTxo (C.PaymentCredentialByKey opPkh) directory
   pure (Convex.CoinSelection.signBalancedTxBody [] tx)
+
+seizeMultiCredentialAssetsTx :: forall era env err m.
+  ( MonadReader env m
+  , Env.HasOperatorEnv era env
+  , Env.HasTransferLogicEnv env
+  , Env.HasDirectoryEnv env
+  , MonadBlockchain era m
+  , MonadError err m
+  , C.IsBabbageBasedEra era
+  , C.HasScriptLanguageInEra C.PlutusScriptV3 era
+  , MonadUtxoQuery m
+  , AsProgrammableTokensError err
+  , AsBalancingError err era
+  , AsCoinSelectionError err
+  , AsRegulatedStablecoinError err
+  )
+  => BuildTx.SeizeReason
+  -> Int
+  -> C.PaymentCredential -- ^ Source/User credential
+  -> m (C.Tx era)
+seizeMultiCredentialAssetsTx reason numUTxOsToSeize sanctionedCred = do
+  toSeizePolicyId <- asks programmableTokenPolicyId
+  opPkh <- asks (fst . Env.bteOperator . Env.operatorEnv @era)
+  directory <- Query.registryNodes @era
+
+  utxosToSeize <- take numUTxOsToSeize . filter (utxoHasPolicyId toSeizePolicyId) <$> Query.userProgrammableOutputs sanctionedCred
+
+  when (null utxosToSeize) $
+    throwing_ _NoTokensToSeize
+  paramsTxIn <- Query.globalParamsNode @era
+  (tx, _) <- Env.balanceTxEnv_ $ do
+    BuildTx.multiSeizeSmartTokens reason paramsTxIn toSeizePolicyId utxosToSeize (C.PaymentCredentialByKey opPkh) directory
+  pure (Convex.CoinSelection.signBalancedTxBody [] tx)

src/examples/regulated-stablecoin/test/unit/Wst/Test/UnitTest.hs

@@ -5,7 +5,7 @@ module Wst.Test.UnitTest(
 
 import Cardano.Api qualified as C
 import Cardano.Api.Shelley qualified as C
-import Control.Monad (void)
+import Control.Monad (void, when)
 import Control.Monad.Except (MonadError)
 import Control.Monad.Reader (MonadReader (ask), ReaderT (runReaderT), asks)
 import Convex.BuildTx qualified as BuildTx
@@ -22,6 +22,7 @@ import Data.List (isPrefixOf)
 import Data.String (IsString (..))
 import GHC.Exception (SomeException, throw)
 import ProgrammableTokens.OffChain.Endpoints qualified as Endpoints
+import ProgrammableTokens.OffChain.Env (programmableTokenPolicyId)
 import ProgrammableTokens.OffChain.Env.Operator qualified as Env
 import ProgrammableTokens.OffChain.Query qualified as Query
 import ProgrammableTokens.Test (deployDirectorySet)
@@ -55,6 +56,7 @@ scriptTargetTests target =
         , testCase "blacklisted transfer" (mockchainFails (blacklistTransfer DontSubmitFailingTx) assertBlacklistedAddressException)
         , testCase "blacklisted transfer (failing tx)" (Test.mockchainSucceedsWithTarget @(AppError C.ConwayEra) target (blacklistTransfer SubmitFailingTx >>= Test.assertFailingTx))
         , testCase "seize user output" (Test.mockchainSucceedsWithTarget @(AppError C.ConwayEra) target $ deployDirectorySet admin >>= seizeUserOutput)
+        , testCase "seize multi user outputs" (Test.mockchainSucceedsWithTarget @(AppError C.ConwayEra) target $ deployDirectorySet admin >>= seizeMultiUserOutputs)
         , testCase "deploy all" (Test.mockchainSucceedsWithTarget @(AppError C.ConwayEra) target deployAll)
         ]
     ]
@@ -224,8 +226,48 @@ seizeUserOutput scriptRoot = Env.withEnv $ do
     Query.userProgrammableOutputs (C.PaymentCredentialByKey userPkh)
       >>= void . Test.expectN 1 "user programmable outputs"
     Query.userProgrammableOutputs (C.PaymentCredentialByKey opPkh)
+      >>= void . Test.expectN 2 "operator programmable outputs"
+
+seizeMultiUserOutputs ::  (MonadUtxoQuery m, MonadFail m, MonadMockchain C.ConwayEra m, MonadError (AppError C.ConwayEra) m) => DirectoryScriptRoot -> m ()
+seizeMultiUserOutputs scriptRoot = Env.withEnv $ do
+  userPkh <- asWallet @C.ConwayEra Wallet.w2 $ asks (fst . Env.bteOperator . Env.operatorEnv @C.ConwayEra)
+  let userPaymentCred = C.PaymentCredentialByKey userPkh
+
+  aid <- issueTransferLogicProgrammableToken scriptRoot
+
+  asAdmin @C.ConwayEra $ Env.withDirectoryFor scriptRoot $ Env.withTransferFromOperator @C.ConwayEra $ Endpoints.deployBlacklistTx
+    >>= void . sendTx . signTxOperator admin
+
+  asAdmin @C.ConwayEra $ Env.withDirectoryFor scriptRoot $ Env.withTransferFromOperator @C.ConwayEra $ do
+    Endpoints.transferSmartTokensTx DontSubmitFailingTx aid 50 (C.PaymentCredentialByKey userPkh)
+      >>= void . sendTx . signTxOperator admin
+    Query.programmableLogicOutputs @C.ConwayEra
+      >>= void . Test.expectN 2 "programmable logic outputs"
+    Query.userProgrammableOutputs (C.PaymentCredentialByKey userPkh)
+      >>= void . Test.expectN 1 "user programmable outputs"
+
+  asAdmin @C.ConwayEra $ Env.withDirectoryFor scriptRoot $ Env.withTransferFromOperator @C.ConwayEra $ do
+    Endpoints.transferSmartTokensTx DontSubmitFailingTx aid 50 (C.PaymentCredentialByKey userPkh)
+      >>= void . sendTx . signTxOperator admin
+    Query.programmableLogicOutputs @C.ConwayEra
+      >>= void . Test.expectN 3 "programmable logic outputs"
+    Query.userProgrammableOutputs (C.PaymentCredentialByKey userPkh)
       >>= void . Test.expectN 2 "user programmable outputs"
 
+  asAdmin @C.ConwayEra $ Env.withDirectoryFor scriptRoot $ Env.withTransferFromOperator @C.ConwayEra $ do
+    opPkh <- asks (fst . Env.bteOperator . Env.operatorEnv @C.ConwayEra)
+    toSeizePolicyId <- asks programmableTokenPolicyId
+    Endpoints.seizeMultiCredentialAssetsTx mempty 2 userPaymentCred
+      >>= void . sendTx . signTxOperator admin
+    Query.programmableLogicOutputs @C.ConwayEra
+      >>= void . Test.expectN 4 "programmable logic outputs"
+    userOutputs <- Query.userProgrammableOutputs (C.PaymentCredentialByKey userPkh)
+    Test.expectN 2 "user programmable outputs" userOutputs
+    mapM_ (\utxo -> when (Query.utxoHasPolicyId toSeizePolicyId utxo) $ fail "User should not have any UTxOs with the programmable token policy ID") userOutputs
+
+    Query.userProgrammableOutputs (C.PaymentCredentialByKey opPkh)
+      >>= void . Test.expectN 2 "operator programmable outputs"
+
 -- TODO: registration to be moved to the endpoints
 registerTransferScripts :: (MonadFail m, MonadError (AppError C.ConwayEra) m, MonadReader env m, Env.HasTransferLogicEnv env, MonadMockchain C.ConwayEra m) => C.Hash C.PaymentKey -> m C.TxId
 registerTransferScripts pkh = do

src/programmable-tokens-offchain/lib/ProgrammableTokens/OffChain/Query.hs

@@ -13,8 +13,11 @@ module ProgrammableTokens.OffChain.Query(
   issuanceCborHexUTxO,
   globalParamsNode,
   programmableLogicOutputs,
-  selectProgammableOutputsFor
+  selectProgammableOutputsFor,
 
+  utxoHasPolicyId,
+  hasPolicyId,
+  extractValue,
 ) where
 
 import Cardano.Api qualified as C