Fix chain index query filtering (#153)

* Adds handling of missing cases when collateral should not be returned by chain-index responses
* Adds node query effect like UtxosAt but with filtering - UtxosAtExcluding
* Enables other collateral tests (they were not included to suite for some reason)
* Fixes some collateral-affected tests

src/BotPlutusInterface/Balance.hs

@@ -13,9 +13,9 @@ module BotPlutusInterface.Balance (
 
 import BotPlutusInterface.BodyBuilder qualified as BodyBuilder
 import BotPlutusInterface.CardanoCLI qualified as CardanoCLI
-import BotPlutusInterface.CardanoNode.Effects (NodeQuery (UtxosAt))
+import BotPlutusInterface.CardanoNode.Effects (NodeQuery (UtxosAt, UtxosAtExcluding))
 import BotPlutusInterface.CoinSelection (selectTxIns)
-import BotPlutusInterface.Collateral (removeCollateralFromMap)
+
 import BotPlutusInterface.Effects (
   PABEffect,
   createDirectoryIfMissingCLI,
@@ -247,8 +247,14 @@ utxosAndCollateralAtAddress ::
   Eff effs (Either Text (Map TxOutRef Tx.ChainIndexTxOut, Maybe CollateralUtxo))
 utxosAndCollateralAtAddress balanceCfg _pabConf changeAddr =
   runEitherT $ do
-    utxos <- firstEitherT (Text.pack . show) $ newEitherT $ queryNode @w (UtxosAt changeAddr)
     inMemCollateral <- lift $ getInMemCollateral @w
+    let nodeQuery =
+          maybe
+            (UtxosAt changeAddr)
+            (UtxosAtExcluding changeAddr . Set.singleton . collateralTxOutRef)
+            inMemCollateral
+
+    utxos <- firstEitherT (Text.pack . show) $ newEitherT $ queryNode @w nodeQuery
 
     -- check if `bcHasScripts` is true, if this is the case then we search of
     -- collateral UTxO in the environment, if such collateral is not present we throw Error.
@@ -259,9 +265,9 @@ utxosAndCollateralAtAddress balanceCfg _pabConf changeAddr =
               "The given transaction uses script, but there's no collateral provided."
                 <> "This usually means that, we failed to create Tx and update our ContractEnvironment."
           )
-          (const $ pure (removeCollateralFromMap inMemCollateral utxos, inMemCollateral))
+          (const $ pure (utxos, inMemCollateral))
           inMemCollateral
-      else pure (removeCollateralFromMap inMemCollateral utxos, Nothing)
+      else pure (utxos, Nothing)
 
 hasChangeUTxO :: Address -> Tx -> Bool
 hasChangeUTxO changeAddr tx =

src/BotPlutusInterface/CardanoNode/Effects.hs

@@ -41,9 +41,10 @@ import Control.Monad.Trans.Class (lift)
 import Control.Monad.Trans.Either (firstEitherT, hoistEither, newEitherT, runEitherT)
 import Data.Map (Map)
 import Data.Map qualified as Map
+import Data.Set (Set)
 import Data.Set qualified as Set
 import Ledger.Address (Address)
-import Ledger.Tx (ChainIndexTxOut (..))
+import Ledger.Tx (ChainIndexTxOut (..), TxOutRef)
 import Ledger.Tx.CardanoAPI qualified as TxApi
 import Plutus.V2.Ledger.Tx qualified as V2
 import Prelude
@@ -54,6 +55,9 @@ import Prelude
 data NodeQuery a where
   -- | 'UtxosAt' queries local node to get all the utxos at particular address.
   UtxosAt :: Address -> NodeQuery (Either NodeQueryError (Map V2.TxOutRef ChainIndexTxOut))
+  -- | 'UtxosAtExcluding' queries local node to get all the utxos at particular address
+  -- excluding `TxOutRefs`'s specified in `Set`.
+  UtxosAtExcluding :: Address -> Set TxOutRef -> NodeQuery (Either NodeQueryError (Map V2.TxOutRef ChainIndexTxOut))
   -- | 'PParams' queries local node to get it's 'ProtocolParameters'.
   PParams :: NodeQuery (Either NodeQueryError CApi.S.ProtocolParameters)
 
@@ -78,6 +82,9 @@ handleNodeQuery =
   interpret $ \case
     UtxosAt addr -> handleUtxosAt addr
     PParams -> queryBabbageEra CApi.QueryProtocolParameters
+    UtxosAtExcluding addr excluded ->
+      let filterOuts = Map.filterWithKey (\oref _ -> not $ oref `Set.member` excluded)
+       in fmap filterOuts <$> handleUtxosAt addr
 
 handleUtxosAt ::
   forall effs.

src/BotPlutusInterface/ChainIndex.hs

@@ -4,11 +4,9 @@ module BotPlutusInterface.ChainIndex (
   handleChainIndexReq,
 ) where
 
-import BotPlutusInterface.Collateral (removeCollateralFromPage)
 import BotPlutusInterface.Types (
   ContractEnvironment (ContractEnvironment, cePABConfig),
   PABConfig,
-  readCollateralUtxo,
  )
 import Data.Kind (Type)
 import Network.HTTP.Client (
@@ -21,10 +19,10 @@ import Network.HTTP.Types (Status (statusCode))
 import Plutus.ChainIndex.Api (
   QueryAtAddressRequest (QueryAtAddressRequest),
   TxoAtAddressRequest (TxoAtAddressRequest),
-  TxosResponse (TxosResponse),
+  TxosResponse,
   UtxoAtAddressRequest (UtxoAtAddressRequest),
   UtxoWithCurrencyRequest (UtxoWithCurrencyRequest),
-  UtxosResponse (UtxosResponse),
+  UtxosResponse,
  )
 import Plutus.ChainIndex.Client qualified as ChainIndexClient
 import Plutus.Contract.Effects (ChainIndexQuery (..), ChainIndexResponse (..))
@@ -73,7 +71,7 @@ handleChainIndexReq contractEnv@ContractEnvironment {cePABConfig} =
           contractEnv
           (ChainIndexClient.getUtxoSetAtAddress (UtxoAtAddressRequest (Just page) credential))
     UtxoSetWithCurrency page assetClass ->
-      UtxoSetAtResponse
+      UtxoSetWithCurrencyResponse
         <$> chainIndexUtxoQuery
           contractEnv
           (ChainIndexClient.getUtxoSetWithCurrency (UtxoWithCurrencyRequest (Just page) assetClass))
@@ -105,24 +103,16 @@ chainIndexQueryOne pabConf endpoint = do
       | otherwise -> error (show failureResp)
     Left failureResp -> error (show failureResp)
 
--- | Query for utxo's and filter collateral utxo from result.
+-- | Query for utxo's.
 chainIndexUtxoQuery :: forall (w :: Type). ContractEnvironment w -> ClientM UtxosResponse -> IO UtxosResponse
 chainIndexUtxoQuery contractEnv query = do
-  collateralUtxo <- readCollateralUtxo contractEnv
-  let removeCollateral :: UtxosResponse -> UtxosResponse
-      removeCollateral (UtxosResponse tip page) = UtxosResponse tip (removeCollateralFromPage collateralUtxo page)
-  removeCollateral
-    <$> chainIndexQueryMany
-      contractEnv.cePABConfig
-      query
+  chainIndexQueryMany
+    contractEnv.cePABConfig
+    query
 
--- | Query for txo's and filter collateral txo from result.
+-- | Query for txo's.
 chainIndexTxoQuery :: forall (w :: Type). ContractEnvironment w -> ClientM TxosResponse -> IO TxosResponse
 chainIndexTxoQuery contractEnv query = do
-  collateralUtxo <- readCollateralUtxo contractEnv
-  let removeCollateral :: TxosResponse -> TxosResponse
-      removeCollateral (TxosResponse page) = TxosResponse (removeCollateralFromPage collateralUtxo page)
-  removeCollateral
-    <$> chainIndexQueryMany
-      contractEnv.cePABConfig
-      query
+  chainIndexQueryMany
+    contractEnv.cePABConfig
+    query

src/BotPlutusInterface/Collateral.hs

@@ -1,10 +1,8 @@
 module BotPlutusInterface.Collateral (
   getInMemCollateral,
   setInMemCollateral,
-  filterCollateral,
   mkCollateralTx,
-  removeCollateralFromPage,
-  removeCollateralFromMap,
+  withCollateralHandling,
 ) where
 
 import BotPlutusInterface.Types (
@@ -16,12 +14,32 @@ import BotPlutusInterface.Types (
  )
 import Cardano.Prelude (Void)
 import Control.Concurrent.STM (atomically, readTVarIO, writeTVar)
+import Control.Monad (unless)
 import Data.Kind (Type)
-import Data.Map (Map)
-import Data.Map qualified as Map
-import Ledger (ChainIndexTxOut, PaymentPubKeyHash (PaymentPubKeyHash), TxOutRef)
+import Ledger (PaymentPubKeyHash (PaymentPubKeyHash), TxOutRef)
 import Ledger.Constraints qualified as Constraints
 import Plutus.ChainIndex (Page (pageItems))
+import Plutus.ChainIndex.Api (
+  IsUtxoResponse (IsUtxoResponse),
+  QueryResponse (QueryResponse),
+  TxosResponse (paget),
+  UtxosResponse (page),
+ )
+import Plutus.Contract.Effects (
+  ChainIndexQuery (..),
+  ChainIndexResponse (
+    TxOutRefResponse,
+    TxoSetAtResponse,
+    UnspentTxOutResponse,
+    UnspentTxOutsAtResponse,
+    UtxoSetAtResponse,
+    UtxoSetMembershipResponse,
+    UtxoSetWithCurrencyResponse
+  ),
+  PABReq (ChainIndexQueryReq),
+  PABResp (ChainIndexQueryResp),
+  matches,
+ )
 import Prelude
 
 getInMemCollateral :: forall (w :: Type). ContractEnvironment w -> IO (Maybe CollateralUtxo)
@@ -38,16 +56,68 @@ mkCollateralTx pabConf = Constraints.mkTx @Void mempty txc
     txc :: Constraints.TxConstraints Void Void
     txc = Constraints.mustPayToPubKey (PaymentPubKeyHash $ pcOwnPubKeyHash pabConf) (collateralValue pabConf)
 
-filterCollateral :: CollateralUtxo -> [TxOutRef] -> [TxOutRef]
-filterCollateral (CollateralUtxo collateralTxOutRef) = filter (/= collateralTxOutRef)
+-- | Middleware to run `chain-index` queries and filter out collateral output from response.
+withCollateralHandling ::
+  Monad m =>
+  Maybe CollateralUtxo ->
+  (ChainIndexQuery -> m ChainIndexResponse) ->
+  ChainIndexQuery ->
+  m ChainIndexResponse
+withCollateralHandling mCollateral runChainIndexQuery = \query -> do
+  response <-
+    adjustChainIndexResponse mCollateral query
+      <$> runChainIndexQuery query
+  ensureMatches query response
+  pure response
+  where
+    ensureMatches query result =
+      unless (matches (ChainIndexQueryReq query) (ChainIndexQueryResp result)) $
+        error $
+          mconcat
+            [ "Chain-index request doesn't match response."
+            , "\nRequest: " ++ show query
+            , "\nResponse:" ++ show result
+            ]
+
+adjustChainIndexResponse :: Maybe CollateralUtxo -> ChainIndexQuery -> ChainIndexResponse -> ChainIndexResponse
+adjustChainIndexResponse mc ciQuery ciResponse =
+  case mc of
+    Nothing -> ciResponse
+    Just (CollateralUtxo collateralOref) -> case (ciQuery, ciResponse) of
+      -- adjustment based on response
+      (_, UtxoSetAtResponse utxosResp) ->
+        let newPage = removeCollateralFromPage mc (page utxosResp)
+         in UtxoSetAtResponse $ utxosResp {page = newPage}
+      (_, TxoSetAtResponse txosResp) ->
+        let newPaget = removeCollateralFromPage mc (paget txosResp)
+         in TxoSetAtResponse $ txosResp {paget = newPaget}
+      (_, UnspentTxOutsAtResponse (QueryResponse refsAndOuts nq)) ->
+        let filtered = filter (\v -> fst v /= collateralOref) refsAndOuts
+         in UnspentTxOutsAtResponse $ QueryResponse filtered nq
+      (_, UtxoSetWithCurrencyResponse utxosResp) ->
+        let newPage = removeCollateralFromPage mc (page utxosResp)
+         in UtxoSetWithCurrencyResponse $ utxosResp {page = newPage}
+      -- adjustment based on request
+      (UtxoSetMembership oref, UtxoSetMembershipResponse (IsUtxoResponse ct isU)) ->
+        UtxoSetMembershipResponse $
+          IsUtxoResponse ct $
+            oref /= collateralOref && isU
+      (TxOutFromRef oref, TxOutRefResponse _) ->
+        if collateralOref == oref
+          then TxOutRefResponse Nothing
+          else ciResponse
+      (UnspentTxOutFromRef oref, UnspentTxOutResponse _) ->
+        if collateralOref == oref
+          then UnspentTxOutResponse Nothing
+          else ciResponse
+      -- all other cases
+      (_, rest) -> rest
 
 -- | Removes collateral utxo from the UtxoResponse page. Receives `Nothing` if Collateral uninitialized.
 removeCollateralFromPage :: Maybe CollateralUtxo -> Page TxOutRef -> Page TxOutRef
 removeCollateralFromPage = \case
   Nothing -> id
-  Just txOutRef -> \page -> page {pageItems = filterCollateral txOutRef (pageItems page)}
+  Just txOutRef -> \page' -> page' {pageItems = filterCollateral txOutRef (pageItems page')}
 
-removeCollateralFromMap :: Maybe CollateralUtxo -> Map TxOutRef ChainIndexTxOut -> Map TxOutRef ChainIndexTxOut
-removeCollateralFromMap = \case
-  Nothing -> id
-  Just (CollateralUtxo collateral) -> Map.filterWithKey (\oref _ -> collateral /= oref)
+filterCollateral :: CollateralUtxo -> [TxOutRef] -> [TxOutRef]
+filterCollateral (CollateralUtxo collateralTxOutRef) = filter (/= collateralTxOutRef)

src/BotPlutusInterface/Effects.hs

@@ -37,6 +37,7 @@ module BotPlutusInterface.Effects (
 
 import BotPlutusInterface.CardanoNode.Effects (NodeQuery, runNodeQuery)
 import BotPlutusInterface.ChainIndex (handleChainIndexReq)
+import BotPlutusInterface.Collateral (withCollateralHandling)
 import BotPlutusInterface.Collateral qualified as Collateral
 import BotPlutusInterface.ExBudget qualified as ExBudget
 import BotPlutusInterface.TimeSlot qualified as TimeSlot
@@ -193,8 +194,12 @@ handlePABEffect contractEnv =
             Local -> pure ()
             Remote ipAddr ->
               void $ readProcess "scp" ["-r", Text.unpack dir, Text.unpack $ ipAddr <> ":$HOME"] ""
-        QueryChainIndex query ->
-          handleChainIndexReq contractEnv query
+        QueryChainIndex query -> do
+          collateralUtxo <- Collateral.getInMemCollateral contractEnv
+          withCollateralHandling
+            collateralUtxo
+            (handleChainIndexReq contractEnv)
+            query
         QueryNode query -> runNodeQuery contractEnv.cePABConfig (send query)
         EstimateBudget txPath ->
           ExBudget.estimateBudget contractEnv.cePABConfig txPath

test/Spec.hs

@@ -3,6 +3,7 @@ module Main (main) where
 import Spec.BotPlutusInterface.AdjustUnbalanced qualified
 import Spec.BotPlutusInterface.Balance qualified
 import Spec.BotPlutusInterface.CoinSelection qualified
+import Spec.BotPlutusInterface.Collateral qualified
 import Spec.BotPlutusInterface.Contract qualified
 import Spec.BotPlutusInterface.ContractStats qualified
 import Spec.BotPlutusInterface.Server qualified
@@ -30,4 +31,5 @@ tests =
     , Spec.BotPlutusInterface.ContractStats.tests
     , Spec.BotPlutusInterface.TxStatusChange.tests
     , Spec.BotPlutusInterface.AdjustUnbalanced.tests
+    , Spec.BotPlutusInterface.Collateral.tests
     ]