{-# LANGUAGE OverloadedStrings #-}

{- | Threat model for detecting Invalid Datum Index vulnerabilities.

An Invalid Datum Index Attack mutates the constructor index of a @Constr@ datum
on a script output to a value outside the expected range. If a validator's
@FromData@ parser uses a wildcard or permissive @otherwise@ branch when
deserialising the constructor index, an attacker can supply an invalid index
that still matches a catch-all and is interpreted as a legitimate state.

== Consequences ==

1. __State confusion__: If the out-of-range index accidentally matches a
   catch-all or default branch in @unsafeFromBuiltinData@, the validator may
   interpret the datum as a valid (but semantically wrong) state, allowing
   unintended transitions or fund extraction.

2. __Permanent fund locking__: If the invalid index triggers a script error at
   spend time, any UTxO locked with that corrupted datum becomes permanently
   unspendable.

== Root Cause ==

Validators that use @unsafeFromBuiltinData@ or manually pattern-match on
constructor indices without explicitly rejecting unexpected values are at risk.
For example:

@
case index of
  0 -> Pinged
  1 -> Ponged
  _ -> Stopped   -- catch-all silently accepts ANY other index!
@

This means @Constr 99 []@ would decode as @Stopped@, bypassing any guard that
should have rejected it.

== Mitigation ==

A secure validator should explicitly enumerate all valid indices and call
@P.traceError@ (or equivalent) for any unexpected constructor index:

@
case index of
  0 -> Pinged
  1 -> Ponged
  2 -> Stopped
  _ -> P.traceError "PingPongState: invalid index"
@

This threat model tests whether a script output with an inline datum still
validates when the @Constr@ index is replaced with an out-of-range value.
If it does, the validator has an Invalid Datum Index vulnerability.
-}
module Convex.ThreatModel.InvalidDatumIndex (
  invalidDatumIndexAttack,
  invalidDatumIndexAttackWith,
  replaceConstrIndex,
) where

import Convex.ThreatModel

{- | Check for Invalid Datum Index vulnerabilities using constructor index 5.

This is the default configuration, which replaces the constructor index of any
inline @Constr@ datum on a script output with @5@. Valid PingPong indices are
0 (Pinged), 1 (Ponged), and 2 (Stopped), so @5@ is safely out-of-range for all
known state machines in this codebase. If the transaction still validates, the
validator has a permissive datum index check.
-}
invalidDatumIndexAttack :: ThreatModel ()
invalidDatumIndexAttack :: ThreatModel ()
invalidDatumIndexAttack = Integer -> ThreatModel ()
invalidDatumIndexAttackWith Integer
5

{- | Check for Invalid Datum Index vulnerabilities with a configurable
constructor index.

For a transaction with script outputs containing inline datums:

* Replace the @Constr@ index of the datum with @invalidIdx@
* Leave the constructor fields unchanged so any field parsing still succeeds
* If the transaction still validates, the validator does not reject
  out-of-range constructor indices — it may have a permissive catch-all.

Choose @invalidIdx@ to be outside the range of all valid constructors for the
script under test (e.g., @5@ for a 3-constructor type, @99@ for extra clarity).
-}
invalidDatumIndexAttackWith :: Integer -> ThreatModel ()
invalidDatumIndexAttackWith :: Integer -> ThreatModel ()
invalidDatumIndexAttackWith Integer
invalidIdx = [Char] -> ThreatModel () -> ThreatModel ()
forall a. [Char] -> ThreatModel a -> ThreatModel a
Named ([Char]
"Invalid Datum Index Attack (index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
invalidIdx [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")") (ThreatModel () -> ThreatModel ())
-> ThreatModel () -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$ do
  -- Precondition: transaction must spend a script input (otherwise no validator runs)
  Input
_ <- (Input -> Bool) -> ThreatModel Input
anyInputSuchThat (Bool -> Bool
not (Bool -> Bool) -> (Input -> Bool) -> Input -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AddressAny -> Bool
isKeyAddressAny (AddressAny -> Bool) -> (Input -> AddressAny) -> Input -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> AddressAny
forall t. IsInputOrOutput t => t -> AddressAny
addressOf)

  -- Get all outputs from the transaction
  [Output]
outputs <- ThreatModel [Output]
getTxOutputs

  -- Filter to script outputs with inline datums that carry a Constr datum
  let scriptOutputsWithConstr :: [Output]
scriptOutputsWithConstr = (Output -> Bool) -> [Output] -> [Output]
forall a. (a -> Bool) -> [a] -> [a]
filter Output -> Bool
isScriptOutputWithConstrInlineDatum [Output]
outputs

  -- Precondition: there must be at least one eligible target output
  ThreatModel () -> ThreatModel ()
forall a. ThreatModel a -> ThreatModel a
threatPrecondition (ThreatModel () -> ThreatModel ())
-> ThreatModel () -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$ Bool -> ThreatModel ()
ensure (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Output] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Output]
scriptOutputsWithConstr)

  -- Pick a target output
  Output
target <- [Output] -> ThreatModel Output
forall a. Show a => [a] -> ThreatModel a
pickAny [Output]
scriptOutputsWithConstr

  -- Extract the inline datum (we know it exists due to the filter above)
  ScriptData
originalDatum <- case Output -> Maybe ScriptData
getInlineDatum Output
target of
    Maybe ScriptData
Nothing -> [Char] -> ThreatModel ScriptData
forall a. [Char] -> ThreatModel a
failPrecondition [Char]
"Script output missing inline datum"
    Just ScriptData
d -> ScriptData -> ThreatModel ScriptData
forall a. a -> ThreatModel a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScriptData
d

  let mutatedDatum :: ScriptData
mutatedDatum = Integer -> ScriptData -> ScriptData
replaceConstrIndex Integer
invalidIdx ScriptData
originalDatum

  [Char] -> ThreatModel ()
counterexampleTM ([Char] -> ThreatModel ()) -> [Char] -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$
    [[Char]] -> [Char]
paragraph
      [ [Char]
"The transaction contains a script output at index"
      , TxIx -> [Char]
forall a. Show a => a -> [Char]
show (Output -> TxIx
outputIx Output
target)
      , [Char]
"with an inline Constr datum."
      ]

  [Char] -> ThreatModel ()
counterexampleTM ([Char] -> ThreatModel ()) -> [Char] -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$
    [[Char]] -> [Char]
paragraph
      [ [Char]
"Testing if the datum's constructor index can be replaced with"
      , Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
invalidIdx
      , [Char]
"while the fields are left unchanged, and the transaction still validates."
      ]

  [Char] -> ThreatModel ()
counterexampleTM ([Char] -> ThreatModel ()) -> [Char] -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$
    [[Char]] -> [Char]
paragraph
      [ [Char]
"If this validates, the script's FromData parser accepts out-of-range"
      , [Char]
"constructor indices (e.g., via a catch-all branch). An attacker could"
      , [Char]
"exploit this to:"
      , [Char]
"1) Confuse the validator about which state the datum represents"
      , [Char]
"2) Bypass state-transition guards that depend on the constructor index"
      , [Char]
"3) Lock funds permanently with an unspendable corrupted datum"
      ]

  -- This SHOULD fail - if it validates, the contract is vulnerable.
  TxModifier -> ThreatModel ()
shouldNotValidate (TxModifier -> ThreatModel ()) -> TxModifier -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$ Output -> Datum -> TxModifier
forall t. IsInputOrOutput t => t -> Datum -> TxModifier
changeDatumOf Output
target (ScriptData -> Datum
toInlineDatum ScriptData
mutatedDatum)

-- ---------------------------------------------------------------------------
-- Helpers
-- ---------------------------------------------------------------------------

{- | Replace the constructor index of a @ScriptDataConstructor@ with @newIdx@,
preserving all fields.

For other @ScriptData@ variants (@Map@, @List@, @Number@, @Bytes@) the value
is returned unchanged and the precondition filter above will have already
excluded non-Constr datums.
-}
replaceConstrIndex :: Integer -> ScriptData -> ScriptData
replaceConstrIndex :: Integer -> ScriptData -> ScriptData
replaceConstrIndex Integer
newIdx ScriptData
sd = case ScriptData
sd of
  ScriptDataConstructor Integer
_idx [ScriptData]
fields -> Integer -> [ScriptData] -> ScriptData
ScriptDataConstructor Integer
newIdx [ScriptData]
fields
  ScriptData
_ -> ScriptData
sd

-- | True when the output is at a script address and has an inline Constr datum.
isScriptOutputWithConstrInlineDatum :: Output -> Bool
isScriptOutputWithConstrInlineDatum :: Output -> Bool
isScriptOutputWithConstrInlineDatum Output
output =
  Bool -> Bool
not (AddressAny -> Bool
isKeyAddressAny (Output -> AddressAny
forall t. IsInputOrOutput t => t -> AddressAny
addressOf Output
output))
    Bool -> Bool -> Bool
&& case Output -> Maybe ScriptData
getInlineDatum Output
output of
      Just (ScriptDataConstructor{}) -> Bool
True
      Maybe ScriptData
_ -> Bool
False

-- | Extract the inline datum from an output, if present.
getInlineDatum :: Output -> Maybe ScriptData
getInlineDatum :: Output -> Maybe ScriptData
getInlineDatum Output
output =
  case TxOut CtxTx Era -> Datum
forall ctx. TxOut ctx Era -> TxOutDatum ctx Era
datumOfTxOut (Output -> TxOut CtxTx Era
outputTxOut Output
output) of
    TxOutDatumInline BabbageEraOnwards Era
_ HashableScriptData
hashableData -> ScriptData -> Maybe ScriptData
forall a. a -> Maybe a
Just (HashableScriptData -> ScriptData
getScriptData HashableScriptData
hashableData)
    Datum
_ -> Maybe ScriptData
forall a. Maybe a
Nothing

-- | Wrap a @ScriptData@ as an inline datum for use with @changeDatumOf@.
toInlineDatum :: ScriptData -> Datum
toInlineDatum :: ScriptData -> Datum
toInlineDatum ScriptData
sd =
  BabbageEraOnwards Era -> HashableScriptData -> Datum
forall era ctx.
BabbageEraOnwards era -> HashableScriptData -> TxOutDatum ctx era
TxOutDatumInline BabbageEraOnwards Era
BabbageEraOnwardsConway (ScriptData -> HashableScriptData
unsafeHashableScriptData ScriptData
sd)