{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE NoFieldSelectors #-}

-- |
-- Module      : Haskoin.Transaction.Segwit
-- Copyright   : No rights reserved
-- License     : MIT
-- Maintainer  : jprupp@protonmail.ch
-- Stability   : experimental
-- Portability : POSIX
--
-- Types to represent segregated witness data and auxilliary functions to
-- manipulate it.  See [BIP 141](https://github.com/bitcoin/bips/blob/master/bip-0141.mediawiki)
-- and [BIP 143](https://github.com/bitcoin/bips/blob/master/bip-0143.mediawiki) for
-- details.
module Haskoin.Transaction.Segwit
  ( -- * Segwit
    WitnessProgram (..),
    WitnessProgramPKH (..),
    WitnessProgramSH (..),
    isSegwit,
    viewWitnessProgram,
    decodeWitnessInput,
    calcWitnessProgram,
    simpleInputStack,
    toWitnessStack,
  )
where

import Crypto.Secp256k1
import Data.ByteString (ByteString)
import Data.Bytes.Get (runGetS)
import Data.Bytes.Put (runPutS)
import Data.Bytes.Serial (Serial (deserialize, serialize))
import Haskoin.Crypto.Keys.Common
import Haskoin.Network.Data
import Haskoin.Script.Common
import Haskoin.Script.SigHash
import Haskoin.Script.Standard
import Haskoin.Transaction.Common
import Haskoin.Util.Marshal

-- | Test if a 'ScriptOutput' is P2WPKH or P2WSH
--
-- @since 0.11.0.0
isSegwit :: ScriptOutput -> Bool
isSegwit :: ScriptOutput -> Bool
isSegwit = \case
  PayWitnessPKHash {} -> Bool
True
  PayWitnessScriptHash {} -> Bool
True
  ScriptOutput
_ -> Bool
False

-- | High level represenation of a (v0) witness program
--
-- @since 0.11.0.0
data WitnessProgram
  = P2WPKH WitnessProgramPKH
  | P2WSH WitnessProgramSH
  | EmptyWitnessProgram
  deriving (WitnessProgram -> WitnessProgram -> Bool
(WitnessProgram -> WitnessProgram -> Bool)
-> (WitnessProgram -> WitnessProgram -> Bool) -> Eq WitnessProgram
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WitnessProgram -> WitnessProgram -> Bool
== :: WitnessProgram -> WitnessProgram -> Bool
$c/= :: WitnessProgram -> WitnessProgram -> Bool
/= :: WitnessProgram -> WitnessProgram -> Bool
Eq)

-- | Encode a witness program
--
-- @since 0.11.0.0
toWitnessStack :: Network -> Ctx -> WitnessProgram -> WitnessStack
toWitnessStack :: Network -> Ctx -> WitnessProgram -> WitnessStack
toWitnessStack Network
net Ctx
ctx = \case
  P2WPKH (WitnessProgramPKH TxSignature
sig PublicKey
key) ->
    [Network -> Ctx -> TxSignature -> ByteString
encodeTxSig Network
net Ctx
ctx TxSignature
sig, Ctx -> PublicKey -> ByteString
forall s a. Marshal s a => s -> a -> ByteString
marshal Ctx
ctx PublicKey
key]
  P2WSH (WitnessProgramSH WitnessStack
stack Script
scr) ->
    WitnessStack
stack WitnessStack -> WitnessStack -> WitnessStack
forall a. Semigroup a => a -> a -> a
<> [Put -> ByteString
runPutS (Script -> Put
forall a (m :: * -> *). (Serial a, MonadPut m) => a -> m ()
forall (m :: * -> *). MonadPut m => Script -> m ()
serialize Script
scr)]
  WitnessProgram
EmptyWitnessProgram ->
    WitnessStack
forall a. Monoid a => a
mempty

-- | High level representation of a P2WPKH witness
--
-- @since 0.11.0.0
data WitnessProgramPKH = WitnessProgramPKH
  { WitnessProgramPKH -> TxSignature
signature :: !TxSignature,
    WitnessProgramPKH -> PublicKey
key :: !PublicKey
  }
  deriving (WitnessProgramPKH -> WitnessProgramPKH -> Bool
(WitnessProgramPKH -> WitnessProgramPKH -> Bool)
-> (WitnessProgramPKH -> WitnessProgramPKH -> Bool)
-> Eq WitnessProgramPKH
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
== :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
$c/= :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
/= :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
Eq)

-- | High-level representation of a P2WSH witness
--
-- @since 0.11.0.0
data WitnessProgramSH = WitnessProgramSH
  { WitnessProgramSH -> WitnessStack
stack :: ![ByteString],
    WitnessProgramSH -> Script
script :: !Script
  }
  deriving (WitnessProgramSH -> WitnessProgramSH -> Bool
(WitnessProgramSH -> WitnessProgramSH -> Bool)
-> (WitnessProgramSH -> WitnessProgramSH -> Bool)
-> Eq WitnessProgramSH
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WitnessProgramSH -> WitnessProgramSH -> Bool
== :: WitnessProgramSH -> WitnessProgramSH -> Bool
$c/= :: WitnessProgramSH -> WitnessProgramSH -> Bool
/= :: WitnessProgramSH -> WitnessProgramSH -> Bool
Eq, Int -> WitnessProgramSH -> ShowS
[WitnessProgramSH] -> ShowS
WitnessProgramSH -> String
(Int -> WitnessProgramSH -> ShowS)
-> (WitnessProgramSH -> String)
-> ([WitnessProgramSH] -> ShowS)
-> Show WitnessProgramSH
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> WitnessProgramSH -> ShowS
showsPrec :: Int -> WitnessProgramSH -> ShowS
$cshow :: WitnessProgramSH -> String
show :: WitnessProgramSH -> String
$cshowList :: [WitnessProgramSH] -> ShowS
showList :: [WitnessProgramSH] -> ShowS
Show)

-- | Calculate the witness program from the transaction data
--
-- @since 0.11.0.0
viewWitnessProgram ::
  Network ->
  Ctx ->
  ScriptOutput ->
  WitnessStack ->
  Either String WitnessProgram
viewWitnessProgram :: Network
-> Ctx
-> ScriptOutput
-> WitnessStack
-> Either String WitnessProgram
viewWitnessProgram Network
net Ctx
ctx ScriptOutput
so WitnessStack
witness = case ScriptOutput
so of
  PayWitnessPKHash Hash160
_ | WitnessStack -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length WitnessStack
witness Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 -> do
    TxSignature
sig <- Network -> Ctx -> ByteString -> Either String TxSignature
decodeTxSig Network
net Ctx
ctx (WitnessStack -> ByteString
forall a. HasCallStack => [a] -> a
head WitnessStack
witness)
    PublicKey
pubkey <- Ctx -> ByteString -> Either String PublicKey
forall s a. Marshal s a => s -> ByteString -> Either String a
unmarshal Ctx
ctx (ByteString -> Either String PublicKey)
-> ByteString -> Either String PublicKey
forall a b. (a -> b) -> a -> b
$ WitnessStack
witness WitnessStack -> Int -> ByteString
forall a. HasCallStack => [a] -> Int -> a
!! Int
1
    WitnessProgram -> Either String WitnessProgram
forall a. a -> Either String a
forall (m :: * -> *) a. Monad m => a -> m a
return (WitnessProgram -> Either String WitnessProgram)
-> (WitnessProgramPKH -> WitnessProgram)
-> WitnessProgramPKH
-> Either String WitnessProgram
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramPKH -> WitnessProgram
P2WPKH (WitnessProgramPKH -> Either String WitnessProgram)
-> WitnessProgramPKH -> Either String WitnessProgram
forall a b. (a -> b) -> a -> b
$ TxSignature -> PublicKey -> WitnessProgramPKH
WitnessProgramPKH TxSignature
sig PublicKey
pubkey
  PayWitnessScriptHash Hash256
_ | Bool -> Bool
not (WitnessStack -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null WitnessStack
witness) -> do
    Script
redeemScript <- Get Script -> ByteString -> Either String Script
forall a. Get a -> ByteString -> Either String a
runGetS Get Script
forall a (m :: * -> *). (Serial a, MonadGet m) => m a
forall (m :: * -> *). MonadGet m => m Script
deserialize (ByteString -> Either String Script)
-> ByteString -> Either String Script
forall a b. (a -> b) -> a -> b
$ WitnessStack -> ByteString
forall a. HasCallStack => [a] -> a
last WitnessStack
witness
    WitnessProgram -> Either String WitnessProgram
forall a. a -> Either String a
forall (m :: * -> *) a. Monad m => a -> m a
return (WitnessProgram -> Either String WitnessProgram)
-> (WitnessProgramSH -> WitnessProgram)
-> WitnessProgramSH
-> Either String WitnessProgram
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramSH -> WitnessProgram
P2WSH (WitnessProgramSH -> Either String WitnessProgram)
-> WitnessProgramSH -> Either String WitnessProgram
forall a b. (a -> b) -> a -> b
$ WitnessStack -> Script -> WitnessProgramSH
WitnessProgramSH (WitnessStack -> WitnessStack
forall a. HasCallStack => [a] -> [a]
init WitnessStack
witness) Script
redeemScript
  ScriptOutput
_
    | WitnessStack -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null WitnessStack
witness -> WitnessProgram -> Either String WitnessProgram
forall a. a -> Either String a
forall (m :: * -> *) a. Monad m => a -> m a
return WitnessProgram
EmptyWitnessProgram
    | Bool
otherwise -> String -> Either String WitnessProgram
forall a b. a -> Either a b
Left String
"viewWitnessProgram: Invalid witness program"

-- | Analyze the witness, trying to match it with standard input structures
--
-- @since 0.11.0.0
decodeWitnessInput ::
  Network ->
  Ctx ->
  WitnessProgram ->
  Either String (Maybe ScriptOutput, SimpleInput)
decodeWitnessInput :: Network
-> Ctx
-> WitnessProgram
-> Either String (Maybe ScriptOutput, SimpleInput)
decodeWitnessInput Network
net Ctx
ctx = \case
  P2WPKH (WitnessProgramPKH TxSignature
sig PublicKey
key) -> (Maybe ScriptOutput, SimpleInput)
-> Either String (Maybe ScriptOutput, SimpleInput)
forall a. a -> Either String a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ScriptOutput
forall a. Maybe a
Nothing, TxSignature -> PublicKey -> SimpleInput
SpendPKHash TxSignature
sig PublicKey
key)
  P2WSH (WitnessProgramSH WitnessStack
st Script
scr) -> do
    ScriptOutput
so <- Ctx -> Script -> Either String ScriptOutput
decodeOutput Ctx
ctx Script
scr
    (SimpleInput -> (Maybe ScriptOutput, SimpleInput))
-> Either String SimpleInput
-> Either String (Maybe ScriptOutput, SimpleInput)
forall a b. (a -> b) -> Either String a -> Either String b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ScriptOutput -> Maybe ScriptOutput
forall a. a -> Maybe a
Just ScriptOutput
so,) (Either String SimpleInput
 -> Either String (Maybe ScriptOutput, SimpleInput))
-> Either String SimpleInput
-> Either String (Maybe ScriptOutput, SimpleInput)
forall a b. (a -> b) -> a -> b
$ case (ScriptOutput
so, WitnessStack
st) of
      (PayPK PublicKey
_, [ByteString
sigBS]) ->
        TxSignature -> SimpleInput
SpendPK (TxSignature -> SimpleInput)
-> Either String TxSignature -> Either String SimpleInput
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Network -> Ctx -> ByteString -> Either String TxSignature
decodeTxSig Network
net Ctx
ctx ByteString
sigBS
      (PayPKHash Hash160
_, [ByteString
sigBS, ByteString
keyBS]) ->
        TxSignature -> PublicKey -> SimpleInput
SpendPKHash
          (TxSignature -> PublicKey -> SimpleInput)
-> Either String TxSignature
-> Either String (PublicKey -> SimpleInput)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Network -> Ctx -> ByteString -> Either String TxSignature
decodeTxSig Network
net Ctx
ctx ByteString
sigBS
          Either String (PublicKey -> SimpleInput)
-> Either String PublicKey -> Either String SimpleInput
forall a b.
Either String (a -> b) -> Either String a -> Either String b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ctx -> ByteString -> Either String PublicKey
forall s a. Marshal s a => s -> ByteString -> Either String a
unmarshal Ctx
ctx ByteString
keyBS
      (PayMulSig [PublicKey]
_ Int
_, ByteString
"" : WitnessStack
sigsBS) ->
        [TxSignature] -> SimpleInput
SpendMulSig
          ([TxSignature] -> SimpleInput)
-> Either String [TxSignature] -> Either String SimpleInput
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ByteString -> Either String TxSignature)
-> WitnessStack -> Either String [TxSignature]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Network -> Ctx -> ByteString -> Either String TxSignature
decodeTxSig Network
net Ctx
ctx) WitnessStack
sigsBS
      (ScriptOutput, WitnessStack)
_ -> String -> Either String SimpleInput
forall a b. a -> Either a b
Left String
"decodeWitnessInput: Non-standard script output"
  WitnessProgram
EmptyWitnessProgram -> String -> Either String (Maybe ScriptOutput, SimpleInput)
forall a b. a -> Either a b
Left String
"decodeWitnessInput: Empty witness program"

-- | Create the witness program for a standard input
--
-- @since 0.11.0.0
calcWitnessProgram ::
  Network ->
  Ctx ->
  ScriptOutput ->
  ScriptInput ->
  Either String WitnessProgram
calcWitnessProgram :: Network
-> Ctx
-> ScriptOutput
-> ScriptInput
-> Either String WitnessProgram
calcWitnessProgram Network
net Ctx
ctx ScriptOutput
so ScriptInput
si = case (ScriptOutput
so, ScriptInput
si) of
  (PayWitnessPKHash {}, RegularInput (SpendPKHash TxSignature
sig PublicKey
pk)) ->
    WitnessProgram -> Either String WitnessProgram
forall a b. b -> Either a b
Right (WitnessProgram -> Either String WitnessProgram)
-> WitnessProgram -> Either String WitnessProgram
forall a b. (a -> b) -> a -> b
$ TxSignature -> PublicKey -> WitnessProgram
p2wpkh TxSignature
sig PublicKey
pk
  (PayScriptHash {}, RegularInput (SpendPKHash TxSignature
sig PublicKey
pk)) ->
    WitnessProgram -> Either String WitnessProgram
forall a b. b -> Either a b
Right (WitnessProgram -> Either String WitnessProgram)
-> WitnessProgram -> Either String WitnessProgram
forall a b. (a -> b) -> a -> b
$ TxSignature -> PublicKey -> WitnessProgram
p2wpkh TxSignature
sig PublicKey
pk
  (PayWitnessScriptHash {}, ScriptHashInput SimpleInput
i ScriptOutput
o) ->
    WitnessProgram -> Either String WitnessProgram
forall a b. b -> Either a b
Right (WitnessProgram -> Either String WitnessProgram)
-> WitnessProgram -> Either String WitnessProgram
forall a b. (a -> b) -> a -> b
$ SimpleInput -> ScriptOutput -> WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o
  (PayScriptHash {}, ScriptHashInput SimpleInput
i ScriptOutput
o) ->
    WitnessProgram -> Either String WitnessProgram
forall a b. b -> Either a b
Right (WitnessProgram -> Either String WitnessProgram)
-> WitnessProgram -> Either String WitnessProgram
forall a b. (a -> b) -> a -> b
$ SimpleInput -> ScriptOutput -> WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o
  (ScriptOutput, ScriptInput)
_ -> String -> Either String WitnessProgram
forall a b. a -> Either a b
Left String
"calcWitnessProgram: Invalid segwit SigInput"
  where
    p2wpkh :: TxSignature -> PublicKey -> WitnessProgram
p2wpkh TxSignature
sig =
      WitnessProgramPKH -> WitnessProgram
P2WPKH (WitnessProgramPKH -> WitnessProgram)
-> (PublicKey -> WitnessProgramPKH) -> PublicKey -> WitnessProgram
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TxSignature -> PublicKey -> WitnessProgramPKH
WitnessProgramPKH TxSignature
sig
    p2wsh :: SimpleInput -> ScriptOutput -> WitnessProgram
p2wsh SimpleInput
i =
      WitnessProgramSH -> WitnessProgram
P2WSH (WitnessProgramSH -> WitnessProgram)
-> (ScriptOutput -> WitnessProgramSH)
-> ScriptOutput
-> WitnessProgram
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessStack -> Script -> WitnessProgramSH
WitnessProgramSH (Network -> Ctx -> SimpleInput -> WitnessStack
simpleInputStack Network
net Ctx
ctx SimpleInput
i) (Script -> WitnessProgramSH)
-> (ScriptOutput -> Script) -> ScriptOutput -> WitnessProgramSH
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx -> ScriptOutput -> Script
encodeOutput Ctx
ctx

-- | Create the witness stack required to spend a standard P2WSH input
--
-- @since 0.11.0.0
simpleInputStack :: Network -> Ctx -> SimpleInput -> [ByteString]
simpleInputStack :: Network -> Ctx -> SimpleInput -> WitnessStack
simpleInputStack Network
net Ctx
ctx = \case
  SpendPK TxSignature
sig -> [TxSignature -> ByteString
f TxSignature
sig]
  SpendPKHash TxSignature
sig PublicKey
k -> [TxSignature -> ByteString
f TxSignature
sig, Ctx -> PublicKey -> ByteString
forall s a. Marshal s a => s -> a -> ByteString
marshal Ctx
ctx PublicKey
k]
  SpendMulSig [TxSignature]
sigs -> ByteString
"" ByteString -> WitnessStack -> WitnessStack
forall a. a -> [a] -> [a]
: (TxSignature -> ByteString) -> [TxSignature] -> WitnessStack
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TxSignature -> ByteString
f [TxSignature]
sigs
  where
    f :: TxSignature -> ByteString
f TxSignature
TxSignatureEmpty = ByteString
""
    f TxSignature
sig = Network -> Ctx -> TxSignature -> ByteString
encodeTxSig Network
net Ctx
ctx TxSignature
sig