{-# LANGUAGE OverloadedStrings #-}

{- | Threat model for detecting scripts that don't validate output addresses.

Many Plutus scripts validate datum state transitions but forget to check
that outputs actually go to the correct address. This allows an attacker
to redirect funds while satisfying the script's datum requirements.

Example vulnerable pattern:
@
  validator :: Datum -> Redeemer -> ScriptContext -> Bool
  validator oldState action ctx =
    let newState = getOutputDatum ctx
    in validTransition oldState action newState  -- Only checks datum, not address!
@

A secure script should also verify:
@
  && outputGoesToSameScript ctx
  && valueIsPreserved ctx
@
-}
module Convex.ThreatModel.UnprotectedScriptOutput (
  unprotectedScriptOutput,
) where

import Convex.ThreatModel

{- | Check for unprotected script output vulnerabilities.

For a transaction that spends a script UTxO and produces an output back to
the same script address:

* Try redirecting that output to the transaction signer (attacker)
* If the transaction still validates, the script doesn't properly protect
  its outputs - it only validates datum, not the output address.

This catches a common vulnerability pattern where scripts implement state
machine logic but forget to enforce that outputs stay at the script address.
-}
unprotectedScriptOutput :: ThreatModel ()
unprotectedScriptOutput :: ThreatModel ()
unprotectedScriptOutput = String -> ThreatModel () -> ThreatModel ()
forall a. String -> ThreatModel a -> ThreatModel a
Named String
"Unprotected Script Output" (ThreatModel () -> ThreatModel ())
-> ThreatModel () -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$ do
  -- The attacker is one of the transaction signers
  Hash PaymentKey
signer <- ThreatModel (Hash PaymentKey)
anySigner

  -- Find a script input (non-key address = script address)
  Input
scriptInput <- (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)
  let scriptAddr :: AddressAny
scriptAddr = Input -> AddressAny
forall t. IsInputOrOutput t => t -> AddressAny
addressOf Input
scriptInput

  -- Find an output going back to the same script address
  [Output]
outputs <- ThreatModel [Output]
getTxOutputs
  let scriptOutputs :: [Output]
scriptOutputs = (Output -> Bool) -> [Output] -> [Output]
forall a. (a -> Bool) -> [a] -> [a]
filter ((AddressAny -> AddressAny -> Bool
forall a. Eq a => a -> a -> Bool
== AddressAny
scriptAddr) (AddressAny -> Bool) -> (Output -> AddressAny) -> Output -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Output -> AddressAny
forall t. IsInputOrOutput t => t -> AddressAny
addressOf) [Output]
outputs

  -- Precondition: there must be a continuation output to the script
  -- (otherwise there's nothing to redirect)
  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]
scriptOutputs)

  Output
scriptOutput <- [Output] -> ThreatModel Output
forall a. Show a => [a] -> ThreatModel a
pickAny [Output]
scriptOutputs

  String -> ThreatModel ()
counterexampleTM (String -> ThreatModel ()) -> String -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$
    [String] -> String
paragraph
      [ String
"The transaction spends a script UTxO at"
      , Doc -> String
forall a. Show a => a -> String
show (Doc -> String) -> Doc -> String
forall a b. (a -> b) -> a -> b
$ AddressAny -> Doc
prettyAddress AddressAny
scriptAddr
      , String
"and produces a continuation output back to the same address."
      ]

  String -> ThreatModel ()
counterexampleTM (String -> ThreatModel ()) -> String -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$
    [String] -> String
paragraph
      [ String
"Testing if this output can be redirected to the signer at"
      , Doc -> String
forall a. Show a => a -> String
show (Doc -> String) -> Doc -> String
forall a b. (a -> b) -> a -> b
$ AddressAny -> Doc
prettyAddress (Hash PaymentKey -> AddressAny
keyAddressAny Hash PaymentKey
signer)
      , String
"while preserving the datum (to satisfy any datum-based validation)."
      ]

  String -> ThreatModel ()
counterexampleTM (String -> ThreatModel ()) -> String -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$
    [String] -> String
paragraph
      [ String
"If this validates, the script only checks datum state transitions"
      , String
"but doesn't verify that outputs go to the correct address."
      , String
"An attacker could steal funds by satisfying datum requirements"
      , String
"while redirecting the output to their own wallet."
      ]

  -- changeAddressOf preserves the datum, only changing the address
  -- If this validates, the script is vulnerable
  TxModifier -> ThreatModel ()
shouldNotValidate (TxModifier -> ThreatModel ()) -> TxModifier -> ThreatModel ()
forall a b. (a -> b) -> a -> b
$ Output -> AddressAny -> TxModifier
forall t. IsInputOrOutput t => t -> AddressAny -> TxModifier
changeAddressOf Output
scriptOutput (Hash PaymentKey -> AddressAny
keyAddressAny Hash PaymentKey
signer)