{-# 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,
  invalidDatumIndexAttackWithGen,
  replaceConstrIndex,
) where

import Convex.ThreatModel
import Test.QuickCheck (Gen, choose)

{- | Default invalid-datum-index attack. The replacement constructor index is
drawn per transaction from a curated range starting at 3 (above the typical
0, 1, 2 indices of most Plutus sum types), so QuickCheck explores the
parameter space without accidentally hitting a valid index.
-}
invalidDatumIndexAttack :: ThreatModel ()
invalidDatumIndexAttack :: ThreatModel ()
invalidDatumIndexAttack = Gen Integer -> ThreatModel ()
invalidDatumIndexAttackWithGen ((Integer, Integer) -> Gen Integer
forall a. Random a => (a, a) -> Gen a
choose (Integer
3, Integer
100))

{- | Invalid-datum-index attack with a fixed constructor index. Keep using
this for deterministic regression tests and golden seeds.
-}
invalidDatumIndexAttackWith :: Integer -> ThreatModel ()
invalidDatumIndexAttackWith :: Integer -> ThreatModel ()
invalidDatumIndexAttackWith = Gen Integer -> ThreatModel ()
invalidDatumIndexAttackWithGen (Gen Integer -> ThreatModel ())
-> (Integer -> Gen Integer) -> Integer -> ThreatModel ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Gen Integer
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

{- | Invalid-datum-index attack parameterised by a generator for the
replacement constructor index. This is the primitive the other two forms
delegate to.
-}
invalidDatumIndexAttackWithGen :: Gen Integer -> ThreatModel ()
invalidDatumIndexAttackWithGen :: Gen Integer -> ThreatModel ()
invalidDatumIndexAttackWithGen Gen Integer
invalidIdxGen =
  String -> ThreatModel () -> ThreatModel ()
forall a. String -> ThreatModel a -> ThreatModel a
Named String
"Invalid Datum Index Attack" (ThreatModel () -> ThreatModel ())
-> ThreatModel () -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$ do
    Integer
invalidIdx <- Gen Integer -> (Integer -> [Integer]) -> ThreatModel Integer
forall a. Show a => Gen a -> (a -> [a]) -> ThreatModel a
forAllTM Gen Integer
invalidIdxGen Integer -> [Integer]
forall a. a -> [a]
noShrink

    -- Negative indices belong to the negative-integer attack, not this model.
    Bool -> ThreatModel ()
ensure (Integer
invalidIdx Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0)

    Output
target <- (Output -> Bool) -> ThreatModel Output
anyGuardedOutputSuchThat Output -> Bool
hasConstrInlineDatum

    -- 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 -> String -> ThreatModel ScriptData
forall a. String -> ThreatModel a
failPrecondition String
"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

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

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

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

    String -> [String] -> ThreatModel ()
tabulateTM String
"invalid index" [Integer -> String
bucketIdx Integer
invalidIdx]

    -- 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)

{- | No shrinking: shrinking toward 0 could produce a valid constructor index,
making the attack vacuous and causing a false 'TMFailed'.
-}
noShrink :: a -> [a]
noShrink :: forall a. a -> [a]
noShrink a
_ = []

-- | Coarse bucket for the invalid-index distribution report.
bucketIdx :: Integer -> String
bucketIdx :: Integer -> String
bucketIdx Integer
n
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
10 = String
"003-010"
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
50 = String
"011-050"
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
100 = String
"051-100"
  | Bool
otherwise = String
"101+"

-- ---------------------------------------------------------------------------
-- 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 has an inline Constr datum.
hasConstrInlineDatum :: Output -> Bool
hasConstrInlineDatum :: Output -> Bool
hasConstrInlineDatum Output
output =
  case Output -> Maybe ScriptData
getInlineDatum Output
output of
    Just (ScriptDataConstructor{}) -> Bool
True
    Maybe ScriptData
_ -> Bool
False