fix(testing-interface): detect signing wallet instead of hardcoding w1

bogdan-manole · bogdan-manole · commit 7954019a1c79 · 2026-03-19T12:51:01.000+02:00
diff --git a/src/testing-interface/lib/Convex/TestingInterface.hs b/src/testing-interface/lib/Convex/TestingInterface.hs
@@ -76,7 +76,7 @@ import Convex.MockChain (MockChainState (..), MockchainT, initialStateFor, runMo
 import Convex.MockChain.Defaults qualified as Defaults
 import Convex.MonadLog (MonadLog)
 import Convex.NodeParams (NodeParams (..))
-import Convex.ThreatModel (ExceptT, ThreatModel, ThreatModelOutcome (..), getThreatModelName, runExceptT, runThreatModelCheck, threatModelEnvs)
+import Convex.ThreatModel (ExceptT, SigningWallet (AutoSign), ThreatModel, ThreatModelOutcome (..), getThreatModelName, runExceptT, runThreatModelCheck, threatModelEnvs)
 import Convex.Wallet.MockWallet qualified as Wallet
 import Data.Aeson (ToJSON (..), (.=))
 import Data.Aeson qualified as Aeson
@@ -424,9 +424,36 @@ positiveTest opts mGetTmResultsRef tms evs = monadicIO $ do
     tmResultsWithCov <- liftIO $ forM allToRun $ \tm -> do
       let name = fromMaybe "Unnamed" (getThreatModelName tm)
       (outcome, tmFinalState) <-
-        runMockchainIO (runThreatModelCheck Wallet.w1 tm envs) params state0
+        runMockchainIO (runThreatModelCheck AutoSign tm envs) params state0
       pure (name, outcome, mcsCoverageData tmFinalState)
 
+    --    tmResultsWithCov <- case (lastTx, lastUtxoBefore, lastMockChainState) of
+    --      (Just tx, Just utxo, Just mcState) -> do
+    --        let pparams' = params ^. ledgerProtocolParameters
+    --            env = ThreatModelEnv tx utxo pparams'
+    --        -- Check which threat models have already failed (from previous QuickCheck iterations)
+    --        existingResults <- case mGetTmResultsRef of
+    --          Just getTmRef -> liftIO $ do
+    --            tmRef <- getTmRef
+    --            readIORef tmRef
+    --          Nothing -> pure Map.empty
+    --        let isTMFailed (TMFailed _) = True
+    --            isTMFailed _ = False
+    --            alreadyFailed name = any isTMFailed (fromMaybe [] (Map.lookup name existingResults))
+    --            -- Only filter threat models (tms) for early-stop; expected vulnerabilities (evs) always run
+    --            tmsToRun = filter (not . alreadyFailed . fromMaybe "Unnamed" . getThreatModelName) tms
+    --            allToRun = tmsToRun <> evs -- evs always run, no filtering
+    --            -- Run each threat model in an isolated MockchainT context
+    --        liftIO $ forM allToRun $ \tm -> do
+    --          let name = fromMaybe "Unnamed" (getThreatModelName tm)
+    --          case detectSigningWallet tx of
+    --            Left err -> pure (name, TMError err, mempty)
+    --            Right wallet -> do
+    --              (outcome, tmFinalState) <-
+    --                runMockchainIO (runThreatModelCheck wallet tm [env]) params mcState
+    --              pure (name, outcome, mcsCoverageData tmFinalState)
+    --      _ -> pure []
+
     -- Extract just the (name, outcome) pairs for downstream processing
     let tmResults = [(n, o) | (n, o, _) <- tmResultsWithCov]
         -- Aggregate coverage from all threat model runs
diff --git a/src/testing-interface/lib/Convex/ThreatModel.hs b/src/testing-interface/lib/Convex/ThreatModel.hs
@@ -112,6 +112,9 @@ module Convex.ThreatModel (
   monitorThreatModel,
   monitorLocalThreatModel,
 
+  -- * Wallet selection
+  SigningWallet (..),
+
   -- * Cardano API helpers
   -- $cardanoHelpers
   projectAda,
@@ -153,7 +156,7 @@ import Convex.Class (MockChainState, MonadMockchain (..), coverageData, getUtxo,
 import Convex.MockChain (applyTransaction, runMockchain)
 import Convex.NodeParams (NodeParams, ledgerProtocolParameters)
 import Convex.ThreatModel.Cardano.Api
-import Convex.ThreatModel.Cardano.Api qualified as TM
+import Convex.ThreatModel.Cardano.Api qualified as TM (detectSigningWallet, rebalanceAndSign, txRequiredSigners)
 import Convex.ThreatModel.Pretty
 import Convex.ThreatModel.TxModifier
 import Convex.Wallet (Wallet)
@@ -178,6 +181,13 @@ data ThreatModelEnv = ThreatModelEnv
   , pparams :: LedgerProtocolParameters Era
   }
 
+-- | How to determine the wallet for re-balancing and re-signing modified transactions.
+data SigningWallet
+  = -- | Detect the signing wallet automatically from the transaction's witnesses.
+    AutoSign
+  | -- | Use the specified wallet for signing.
+    SignWith Wallet
+
 -- | Create `ThreatModelEnv`s by reapplying the given transactions in order, starting with the given chain state.
 threatModelEnvs :: NodeParams Era -> [Tx Era] -> MockChainState Era -> [ThreatModelEnv]
 threatModelEnvs params txs chainState0 = fst $ foldM go chainState0 txs
@@ -336,16 +346,22 @@ then performs full Phase 1 + Phase 2 validation via 'applyTransaction'.
 This catches vulnerabilities that would be masked by signature/fee failures
 in the simpler Phase 2-only validation.
 
+The wallet parameter controls signing:
+- @SignWith wallet@ - use the specified wallet for signing
+- @AutoSign@ - detect the signing wallet from the transaction's witnesses
+
 Usage:
 @
 result <- runMockchain0IOWith utxos params $ do
   -- ... run your actions to get a transaction ...
-  runThreatModelM Wallet.w1 unprotectedScriptOutput [env]
+  runThreatModelM (SignWith Wallet.w1) unprotectedScriptOutput [env]
+  -- or auto-detect:
+  runThreatModelM AutoSign unprotectedScriptOutput [env]
 @
 -}
 runThreatModelM
   :: (MonadMockchain Era m, MonadFail m, MonadIO m)
-  => Wallet
+  => SigningWallet
   -> ThreatModel a
   -> [ThreatModelEnv]
   -> m Property
@@ -359,10 +375,12 @@ output cluttering test results.
 
 The property still succeeds/fails correctly based on shouldValidate/shouldNotValidate
 checks, but Monitor/MonitorLocal annotations (counterexampleTM, etc.) are ignored.
+
+The wallet parameter controls signing (see 'runThreatModelM' for details).
 -}
 runThreatModelMQuiet
   :: (MonadMockchain Era m, MonadFail m, MonadIO m)
-  => Wallet
+  => SigningWallet
   -> ThreatModel a
   -> [ThreatModelEnv]
   -> m Property
@@ -373,14 +391,21 @@ runThreatModelM'
   :: (MonadMockchain Era m, MonadFail m, MonadIO m)
   => Bool
   -- ^ quiet: suppress counterexample annotations
-  -> Wallet
+  -> SigningWallet
   -> ThreatModel a
   -> [ThreatModelEnv]
   -> m Property
-runThreatModelM' quiet wallet = go False
+runThreatModelM' quiet signingWallet = go False
  where
   go b _model [] = pure $ b ==> property True
-  go b model (env : envs) = interpM initialMon model
+  go b model (env : envs) = do
+    -- Resolve wallet: use provided or detect from transaction
+    let resolvedWallet = case signingWallet of
+          SignWith w -> Right w
+          AutoSign -> TM.detectSigningWallet (currentTx env)
+    case resolvedWallet of
+      Left err -> pure $ counterexample err False
+      Right wallet -> interpM initialMon wallet model
    where
     initialMon = if quiet then id else counterexample (show info)
 
@@ -398,7 +423,7 @@ runThreatModelM' quiet wallet = go False
         , ""
         ]
 
-    interpM mon = \case
+    interpM mon wallet = \case
       Validate mods k -> do
         let (modifiedTx, modifiedUtxo) = applyTxModifier (currentTx env) (currentUTxOs env) mods
         -- Re-balance and re-sign the modified transaction
@@ -408,20 +433,20 @@ runThreatModelM' quiet wallet = go False
         (report, covData) <- validateTxM params rebalancedTx modifiedUtxo
         -- Accumulate coverage into the running MockChainState
         modifyMockChainState $ \s -> ((), s & coverageData %~ (<> covData))
-        interpM mon (k report)
+        interpM mon wallet (k report)
       Generate gen _shr k -> do
         -- Use QuickCheck's generate in IO
         a <- liftIO $ QC.generate gen
-        interpM mon (k a)
+        interpM mon wallet (k a)
       GetCtx k ->
-        interpM mon (k env)
+        interpM mon wallet (k env)
       Skip -> go b model envs
-      InPrecondition k -> interpM mon (k False)
+      InPrecondition k -> interpM mon wallet (k False)
       Fail err -> pure $ if quiet then property False else mon $ counterexample err False
-      Monitor m k -> if quiet then interpM mon k else m <$> interpM mon k
-      MonitorLocal m k -> if quiet then interpM mon k else interpM (mon . m) k
+      Monitor m k -> if quiet then interpM mon wallet k else m <$> interpM mon wallet k
+      MonitorLocal m k -> if quiet then interpM mon wallet k else interpM (mon . m) wallet k
       Done{} -> go True model envs
-      Named _n k -> interpM mon k
+      Named _n k -> interpM mon wallet k
 
 -- | Extract the name from a threat model, if it was defined with 'Named'.
 getThreatModelName :: ThreatModel a -> Maybe String
@@ -434,19 +459,30 @@ getThreatModelName _ = Nothing
   Rebalancing failures (e.g., "No change output found") are treated as skipped
   because they indicate the transaction modification cannot be applied to this
   particular transaction, similar to a precondition failure.
+
+  The wallet parameter controls signing:
+  - @SignWith wallet@ - use the specified wallet for signing
+  - @AutoSign@ - detect the signing wallet from the transaction's witnesses
 -}
 runThreatModelCheck
   :: (MonadMockchain Era m, MonadFail m, MonadIO m)
-  => Wallet
+  => SigningWallet
   -> ThreatModel a
   -> [ThreatModelEnv]
   -> m ThreatModelOutcome
-runThreatModelCheck wallet = go False
+runThreatModelCheck signingWallet = go False
  where
   go b _model [] = pure $ if b then TMPassed else TMSkipped
-  go b model (env : envs) = checkInterp model
+  go b model (env : envs) = do
+    -- Resolve wallet: use provided or detect from transaction
+    let resolvedWallet = case signingWallet of
+          SignWith w -> Right w
+          AutoSign -> TM.detectSigningWallet (currentTx env)
+    case resolvedWallet of
+      Left err -> pure (TMError err) -- Continue to next env would lose the error, so return it
+      Right wallet -> checkInterp wallet model
    where
-    checkInterp = \case
+    checkInterp wallet = \case
       Validate mods k -> do
         let (modifiedTx, modifiedUtxo) = applyTxModifier (currentTx env) (currentUTxOs env) mods
         params <- askNodeParams
@@ -458,19 +494,19 @@ runThreatModelCheck wallet = go False
           Right rebalancedTx -> do
             (report, covData) <- validateTxM params rebalancedTx modifiedUtxo
             modifyMockChainState $ \s -> ((), s & coverageData %~ (<> covData))
-            checkInterp (k report)
+            checkInterp wallet (k report)
       Generate gen _shr k -> do
         a <- liftIO $ QC.generate gen
-        checkInterp (k a)
+        checkInterp wallet (k a)
       GetCtx k ->
-        checkInterp (k env)
+        checkInterp wallet (k env)
       Skip -> go b model envs
-      InPrecondition k -> checkInterp (k False)
+      InPrecondition k -> checkInterp wallet (k False)
       Fail err -> pure (TMFailed err)
-      Monitor _m k -> checkInterp k -- No Property to wrap; drop monitoring
-      MonitorLocal _m k -> checkInterp k -- No Property to wrap; drop monitoring
+      Monitor _m k -> checkInterp wallet k -- No Property to wrap; drop monitoring
+      MonitorLocal _m k -> checkInterp wallet k -- No Property to wrap; drop monitoring
       Done{} -> go True model envs
-      Named _n k -> checkInterp k
+      Named _n k -> checkInterp wallet k
 
 {- | Check a precondition. If the argument threat model fails, the evaluation of the current
   transaction is skipped. If all transactions in an evaluation of `runThreatModel` are skipped
diff --git a/src/testing-interface/lib/Convex/ThreatModel/Cardano/Api.hs b/src/testing-interface/lib/Convex/ThreatModel/Cardano/Api.hs
@@ -4,7 +4,86 @@
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE TypeApplications #-}
 
-module Convex.ThreatModel.Cardano.Api where
+module Convex.ThreatModel.Cardano.Api (
+  -- * Types
+  Era,
+  LedgerEra,
+
+  -- * TxOut accessors
+  addressOfTxOut,
+  valueOfTxOut,
+  datumOfTxOut,
+  referenceScriptOfTxOut,
+
+  -- * Redeemer and script data
+  redeemerOfTxIn,
+  recomputeScriptData,
+  emptyTxBodyScriptData,
+  addScriptData,
+  updateRedeemer,
+  addMintingRedeemer,
+  recomputeScriptDataForMint,
+  addDatum,
+  toMaryAssetName,
+
+  -- * Address utilities
+  paymentCredentialToAddressAny,
+  scriptAddressAny,
+  keyAddressAny,
+  isKeyAddressAny,
+
+  -- * Datum/Redeemer conversion
+  toCtxUTxODatum,
+  txOutDatum,
+  toScriptData,
+
+  -- * Transaction utilities
+  dummyTxId,
+  makeTxOut,
+  txSigners,
+  mockWalletHashes,
+  detectSigningWallet,
+  txRequiredSigners,
+  txInputs,
+  txReferenceInputs,
+  txOutputs,
+
+  -- * Value utilities
+  leqValue,
+  projectAda,
+
+  -- * Validation
+  ValidityReport (..),
+  validateTx,
+  validateTxM,
+  buildMockState,
+
+  -- * Rebalancing
+  rebalanceAndSignM,
+  rebalanceAndSign,
+  updateExecutionUnits,
+  updateTxRedeemersWithExUnits,
+  updateScriptDataExUnits,
+  recalculateScriptIntegrityHash,
+  getScriptLanguage,
+  getTxFeeCoin,
+  setTxFeeCoin,
+  setTxOutputsList,
+  adjustChangeOutputM,
+  adjustChangeOutput,
+  replaceAt,
+
+  -- * Validity interval
+  convValidityInterval,
+
+  -- * UTxO utilities
+  restrictUTxO,
+
+  -- * Coverage
+  extractCoverageFromValidationError,
+  unescapeHaskellString,
+  extractCoverageAnnotations,
+) where
 
 import Cardano.Api
 
@@ -48,7 +127,7 @@ import Data.ByteString.Short qualified as SBS
 import Data.Either (isRight)
 import Data.Foldable (foldrM)
 import Data.Map qualified as Map
-import Data.Maybe (listToMaybe)
+import Data.Maybe (listToMaybe, mapMaybe)
 import Data.Maybe.Strict
 import Data.SOP.NonEmpty (NonEmpty (NonEmptyOne))
 import Data.Sequence.Strict qualified as Seq
@@ -260,6 +339,18 @@ txSigners (Tx _ wits) = [toHash wit | ShelleyKeyWitness _ (WitVKey wit _) <- wit
 mockWalletHashes :: [(Hash PaymentKey, Wallet)]
 mockWalletHashes = map (\w -> (Wallet.verificationKeyHash w, w)) mockWallets
 
+{- | Detect which mock wallet signed a transaction by examining its witnesses.
+Returns an error message if no known mock wallet is found among the signers.
+-}
+detectSigningWallet :: Tx Era -> Either String Wallet
+detectSigningWallet tx =
+  case txSigners tx of
+    [] -> Left "Transaction has no signers — cannot determine wallet for threat model"
+    signers ->
+      case mapMaybe (\h -> lookup h mockWalletHashes) signers of
+        (w : _) -> Right w
+        [] -> Left "Transaction signers do not match any known mock wallet"
+
 -- | Get the required signers from the transaction body (not witnesses).
 txRequiredSigners :: Tx Era -> [Hash PaymentKey]
 txRequiredSigners (Tx (ShelleyTxBody _ body _ _ _ _) _) =
diff --git a/src/testing-interface/test/AikenKingOfCardanoSpec.hs b/src/testing-interface/test/AikenKingOfCardanoSpec.hs
@@ -67,7 +67,7 @@ import Convex.TestingInterface (
   TestingInterface (..),
   propRunActionsWithOptions,
  )
-import Convex.ThreatModel (ThreatModelEnv (..), runThreatModelM)
+import Convex.ThreatModel (SigningWallet (SignWith), ThreatModelEnv (..), runThreatModelM)
 import Convex.ThreatModel.Cardano.Api (dummyTxId)
 import Convex.ThreatModel.LargeData (largeDataAttackWith)
 import Convex.ThreatModel.SelfReferenceInjection (selfReferenceInjection)
@@ -694,7 +694,7 @@ propKingUnprotectedOutput opts = monadicIO $ do
             }
 
     -- Run the threat model INSIDE MockchainT with full Phase 1 + Phase 2 validation
-    lift $ runThreatModelM Wallet.w1 unprotectedScriptOutput [env]
+    lift $ runThreatModelM (SignWith Wallet.w1) unprotectedScriptOutput [env]
 
   case result of
     (Left err, _) -> do
diff --git a/src/testing-interface/test/AikenLendingSpec.hs b/src/testing-interface/test/AikenLendingSpec.hs
@@ -76,7 +76,7 @@ import Convex.TestingInterface (
   TestingInterface (..),
   propRunActionsWithOptions,
  )
-import Convex.ThreatModel (ThreatModelEnv (..), runThreatModelMQuiet)
+import Convex.ThreatModel (SigningWallet (SignWith), ThreatModelEnv (..), runThreatModelMQuiet)
 import Convex.ThreatModel.Cardano.Api ()
 import Convex.ThreatModel.InputDuplication (inputDuplication)
 import Convex.ThreatModel.UnprotectedScriptOutput (unprotectedScriptOutput)
@@ -709,7 +709,7 @@ propLendingVulnerableToInputDuplication opts = QC.expectFailure $ monadicIO $ do
 
         -- Run inputDuplication threat model
         -- It should find the second loan UTxO and try adding it as another input
-        lift $ runThreatModelMQuiet Wallet.w3 inputDuplication [env]
+        lift $ runThreatModelMQuiet (SignWith Wallet.w3) inputDuplication [env]
       _ -> fail "Expected at least 2 loan request UTxOs"
 
   case result of
diff --git a/src/testing-interface/test/BountySpec.hs b/src/testing-interface/test/BountySpec.hs
diff --git a/src/testing-interface/test/PingPongSpec.hs b/src/testing-interface/test/PingPongSpec.hs
