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

{- |
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 Data.ByteString (ByteString)
import Data.Bytes.Get
import Data.Bytes.Put
import Data.Bytes.Serial
import Haskoin.Data
import Haskoin.Keys.Common
import Haskoin.Script
import Haskoin.Transaction.Common

{- | 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
/= :: WitnessProgram -> WitnessProgram -> Bool
$c/= :: WitnessProgram -> WitnessProgram -> Bool
== :: WitnessProgram -> WitnessProgram -> Bool
$c== :: WitnessProgram -> WitnessProgram -> Bool
Eq, Int -> WitnessProgram -> ShowS
[WitnessProgram] -> ShowS
WitnessProgram -> String
(Int -> WitnessProgram -> ShowS)
-> (WitnessProgram -> String)
-> ([WitnessProgram] -> ShowS)
-> Show WitnessProgram
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [WitnessProgram] -> ShowS
$cshowList :: [WitnessProgram] -> ShowS
show :: WitnessProgram -> String
$cshow :: WitnessProgram -> String
showsPrec :: Int -> WitnessProgram -> ShowS
$cshowsPrec :: Int -> WitnessProgram -> ShowS
Show)

{- | Encode a witness program

 @since 0.11.0.0
-}
toWitnessStack :: WitnessProgram -> WitnessStack
toWitnessStack :: WitnessProgram -> WitnessStack
toWitnessStack = \case
    P2WPKH (WitnessProgramPKH TxSignature
sig PubKeyI
key) -> [TxSignature -> ByteString
encodeTxSig TxSignature
sig, Put -> ByteString
runPutS (PubKeyI -> Put
forall a (m :: * -> *). (Serial a, MonadPut m) => a -> m ()
serialize PubKeyI
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 ()
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
witnessSignature :: !TxSignature
    , WitnessProgramPKH -> PubKeyI
witnessPubKey :: !PubKeyI
    }
    deriving (WitnessProgramPKH -> WitnessProgramPKH -> Bool
(WitnessProgramPKH -> WitnessProgramPKH -> Bool)
-> (WitnessProgramPKH -> WitnessProgramPKH -> Bool)
-> Eq WitnessProgramPKH
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
$c/= :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
== :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
$c== :: WitnessProgramPKH -> WitnessProgramPKH -> Bool
Eq, Int -> WitnessProgramPKH -> ShowS
[WitnessProgramPKH] -> ShowS
WitnessProgramPKH -> String
(Int -> WitnessProgramPKH -> ShowS)
-> (WitnessProgramPKH -> String)
-> ([WitnessProgramPKH] -> ShowS)
-> Show WitnessProgramPKH
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [WitnessProgramPKH] -> ShowS
$cshowList :: [WitnessProgramPKH] -> ShowS
show :: WitnessProgramPKH -> String
$cshow :: WitnessProgramPKH -> String
showsPrec :: Int -> WitnessProgramPKH -> ShowS
$cshowsPrec :: Int -> WitnessProgramPKH -> ShowS
Show)

{- | High-level representation of a P2WSH witness

 @since 0.11.0.0
-}
data WitnessProgramSH = WitnessProgramSH
    { WitnessProgramSH -> WitnessStack
witnessScriptHashStack :: ![ByteString]
    , WitnessProgramSH -> Script
witnessScriptHashScript :: !Script
    }
    deriving (WitnessProgramSH -> WitnessProgramSH -> Bool
(WitnessProgramSH -> WitnessProgramSH -> Bool)
-> (WitnessProgramSH -> WitnessProgramSH -> Bool)
-> Eq WitnessProgramSH
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: WitnessProgramSH -> WitnessProgramSH -> Bool
$c/= :: WitnessProgramSH -> WitnessProgramSH -> Bool
== :: WitnessProgramSH -> WitnessProgramSH -> Bool
$c== :: 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
showList :: [WitnessProgramSH] -> ShowS
$cshowList :: [WitnessProgramSH] -> ShowS
show :: WitnessProgramSH -> String
$cshow :: WitnessProgramSH -> String
showsPrec :: Int -> WitnessProgramSH -> ShowS
$cshowsPrec :: Int -> WitnessProgramSH -> ShowS
Show)

{- | Calculate the witness program from the transaction data

 @since 0.11.0.0
-}
viewWitnessProgram ::
    Network -> ScriptOutput -> WitnessStack -> Either String WitnessProgram
viewWitnessProgram :: Network
-> ScriptOutput -> WitnessStack -> Either String WitnessProgram
viewWitnessProgram Network
net ScriptOutput
so WitnessStack
witness = case ScriptOutput
so of
    PayWitnessPKHash Hash160
_ | WitnessStack -> 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 -> ByteString -> Either String TxSignature
decodeTxSig Network
net (ByteString -> Either String TxSignature)
-> ByteString -> Either String TxSignature
forall a b. (a -> b) -> a -> b
$ WitnessStack -> ByteString
forall a. [a] -> a
head WitnessStack
witness
        PubKeyI
pubkey <- Get PubKeyI -> ByteString -> Either String PubKeyI
forall a. Get a -> ByteString -> Either String a
runGetS Get PubKeyI
forall a (m :: * -> *). (Serial a, MonadGet m) => m a
deserialize (ByteString -> Either String PubKeyI)
-> ByteString -> Either String PubKeyI
forall a b. (a -> b) -> a -> b
$ WitnessStack
witness WitnessStack -> Int -> ByteString
forall a. [a] -> Int -> a
!! Int
1
        WitnessProgram -> Either String WitnessProgram
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 -> PubKeyI -> WitnessProgramPKH
WitnessProgramPKH TxSignature
sig PubKeyI
pubkey
    PayWitnessScriptHash Hash256
_ | Bool -> Bool
not (WitnessStack -> 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
deserialize (ByteString -> Either String Script)
-> ByteString -> Either String Script
forall a b. (a -> b) -> a -> b
$ WitnessStack -> ByteString
forall a. [a] -> a
last WitnessStack
witness
        WitnessProgram -> Either String WitnessProgram
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. [a] -> [a]
init WitnessStack
witness) Script
redeemScript
    ScriptOutput
_
        | WitnessStack -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null WitnessStack
witness -> WitnessProgram -> Either String WitnessProgram
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 ->
    WitnessProgram ->
    Either String (Maybe ScriptOutput, SimpleInput)
decodeWitnessInput :: Network
-> WitnessProgram
-> Either String (Maybe ScriptOutput, SimpleInput)
decodeWitnessInput Network
net = \case
    P2WPKH (WitnessProgramPKH TxSignature
sig PubKeyI
key) -> (Maybe ScriptOutput, SimpleInput)
-> Either String (Maybe ScriptOutput, SimpleInput)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ScriptOutput
forall a. Maybe a
Nothing, TxSignature -> PubKeyI -> SimpleInput
SpendPKHash TxSignature
sig PubKeyI
key)
    P2WSH (WitnessProgramSH WitnessStack
st Script
scr) -> do
        ScriptOutput
so <- Script -> Either String ScriptOutput
decodeOutput Script
scr
        (SimpleInput -> (Maybe ScriptOutput, SimpleInput))
-> Either String SimpleInput
-> Either String (Maybe ScriptOutput, SimpleInput)
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 PubKeyI
_, [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 -> ByteString -> Either String TxSignature
decodeTxSig Network
net ByteString
sigBS
            (PayPKHash Hash160
_, [ByteString
sigBS, ByteString
keyBS]) ->
                TxSignature -> PubKeyI -> SimpleInput
SpendPKHash (TxSignature -> PubKeyI -> SimpleInput)
-> Either String TxSignature
-> Either String (PubKeyI -> SimpleInput)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Network -> ByteString -> Either String TxSignature
decodeTxSig Network
net ByteString
sigBS Either String (PubKeyI -> SimpleInput)
-> Either String PubKeyI -> Either String SimpleInput
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get PubKeyI -> ByteString -> Either String PubKeyI
forall a. Get a -> ByteString -> Either String a
runGetS Get PubKeyI
forall a (m :: * -> *). (Serial a, MonadGet m) => m a
deserialize ByteString
keyBS
            (PayMulSig [PubKeyI]
_ 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)
traverse (Network -> ByteString -> Either String TxSignature
decodeTxSig Network
net) 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 :: ScriptOutput -> ScriptInput -> Either String WitnessProgram
calcWitnessProgram :: ScriptOutput -> ScriptInput -> Either String WitnessProgram
calcWitnessProgram ScriptOutput
so ScriptInput
si = case (ScriptOutput
so, ScriptInput
si) of
    (PayWitnessPKHash{}, RegularInput (SpendPKHash TxSignature
sig PubKeyI
pk)) -> TxSignature -> PubKeyI -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig PubKeyI
pk
    (PayScriptHash{}, RegularInput (SpendPKHash TxSignature
sig PubKeyI
pk)) -> TxSignature -> PubKeyI -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig PubKeyI
pk
    (PayWitnessScriptHash{}, ScriptHashInput SimpleInput
i ScriptOutput
o) -> SimpleInput -> ScriptOutput -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
SimpleInput -> ScriptOutput -> m WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o
    (PayScriptHash{}, ScriptHashInput SimpleInput
i ScriptOutput
o) -> SimpleInput -> ScriptOutput -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
SimpleInput -> ScriptOutput -> m 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 -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig = WitnessProgram -> m WitnessProgram
forall (m :: * -> *) a. Monad m => a -> m a
return (WitnessProgram -> m WitnessProgram)
-> (PubKeyI -> WitnessProgram) -> PubKeyI -> m WitnessProgram
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramPKH -> WitnessProgram
P2WPKH (WitnessProgramPKH -> WitnessProgram)
-> (PubKeyI -> WitnessProgramPKH) -> PubKeyI -> WitnessProgram
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TxSignature -> PubKeyI -> WitnessProgramPKH
WitnessProgramPKH TxSignature
sig
    p2wsh :: SimpleInput -> ScriptOutput -> m WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o = WitnessProgram -> m WitnessProgram
forall (m :: * -> *) a. Monad m => a -> m a
return (WitnessProgram -> m WitnessProgram)
-> (WitnessProgramSH -> WitnessProgram)
-> WitnessProgramSH
-> m WitnessProgram
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramSH -> WitnessProgram
P2WSH (WitnessProgramSH -> m WitnessProgram)
-> WitnessProgramSH -> m WitnessProgram
forall a b. (a -> b) -> a -> b
$ WitnessStack -> Script -> WitnessProgramSH
WitnessProgramSH (SimpleInput -> WitnessStack
simpleInputStack SimpleInput
i) (ScriptOutput -> Script
encodeOutput ScriptOutput
o)

{- | Create the witness stack required to spend a standard P2WSH input

 @since 0.11.0.0
-}
simpleInputStack :: SimpleInput -> [ByteString]
simpleInputStack :: SimpleInput -> WitnessStack
simpleInputStack = \case
    SpendPK TxSignature
sig -> [TxSignature -> ByteString
f TxSignature
sig]
    SpendPKHash TxSignature
sig PubKeyI
k -> [TxSignature -> ByteString
f TxSignature
sig, Put -> ByteString
runPutS (PubKeyI -> Put
forall a (m :: * -> *). (Serial a, MonadPut m) => a -> m ()
serialize PubKeyI
k)]
    SpendMulSig [TxSignature]
sigs -> ByteString
"" ByteString -> WitnessStack -> WitnessStack
forall a. a -> [a] -> [a]
: (TxSignature -> ByteString) -> [TxSignature] -> WitnessStack
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 = TxSignature -> ByteString
encodeTxSig TxSignature
sig