{-# 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
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
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 (forall a (m :: * -> *). (Serial a, MonadPut m) => a -> m ()
serialize PubKeyI
key)]
    P2WSH (WitnessProgramSH WitnessStack
stack Script
scr) -> WitnessStack
stack forall a. Semigroup a => a -> a -> a
<> [Put -> ByteString
runPutS (forall a (m :: * -> *). (Serial a, MonadPut m) => a -> m ()
serialize Script
scr)]
    WitnessProgram
EmptyWitnessProgram -> 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
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
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
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
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
_ | forall (t :: * -> *) a. Foldable t => t a -> Int
length WitnessStack
witness forall a. Eq a => a -> a -> Bool
== Int
2 -> do
        TxSignature
sig <- Network -> ByteString -> Either String TxSignature
decodeTxSig Network
net forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
head WitnessStack
witness
        PubKeyI
pubkey <- forall a. Get a -> ByteString -> Either String a
runGetS forall a (m :: * -> *). (Serial a, MonadGet m) => m a
deserialize forall a b. (a -> b) -> a -> b
$ WitnessStack
witness forall a. [a] -> Int -> a
!! Int
1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramPKH -> WitnessProgram
P2WPKH forall a b. (a -> b) -> a -> b
$ TxSignature -> PubKeyI -> WitnessProgramPKH
WitnessProgramPKH TxSignature
sig PubKeyI
pubkey
    PayWitnessScriptHash Hash256
_ | Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null WitnessStack
witness) -> do
        Script
redeemScript <- forall a. Get a -> ByteString -> Either String a
runGetS forall a (m :: * -> *). (Serial a, MonadGet m) => m a
deserialize forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
last WitnessStack
witness
        forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramSH -> WitnessProgram
P2WSH forall a b. (a -> b) -> a -> b
$ WitnessStack -> Script -> WitnessProgramSH
WitnessProgramSH (forall a. [a] -> [a]
init WitnessStack
witness) Script
redeemScript
    ScriptOutput
_
        | forall (t :: * -> *) a. Foldable t => t a -> Bool
null WitnessStack
witness -> forall (m :: * -> *) a. Monad m => a -> m a
return WitnessProgram
EmptyWitnessProgram
        | Bool
otherwise -> 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) -> forall (m :: * -> *) a. Monad m => a -> m a
return (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
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. a -> Maybe a
Just ScriptOutput
so,) forall a b. (a -> b) -> a -> b
$ case (ScriptOutput
so, WitnessStack
st) of
            (PayPK PubKeyI
_, [ByteString
sigBS]) ->
                TxSignature -> SimpleInput
SpendPK 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Network -> ByteString -> Either String TxSignature
decodeTxSig Network
net ByteString
sigBS forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Get a -> ByteString -> Either String a
runGetS forall a (m :: * -> *). (Serial a, MonadGet m) => m a
deserialize ByteString
keyBS
            (PayMulSig [PubKeyI]
_ Int
_, ByteString
"" : WitnessStack
sigsBS) ->
                [TxSignature] -> SimpleInput
SpendMulSig forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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)
_ -> forall a b. a -> Either a b
Left String
"decodeWitnessInput: Non-standard script output"
    WitnessProgram
EmptyWitnessProgram -> 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)) -> forall {m :: * -> *}.
Monad m =>
TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig PubKeyI
pk
    (PayScriptHash{}, RegularInput (SpendPKHash TxSignature
sig PubKeyI
pk)) -> forall {m :: * -> *}.
Monad m =>
TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig PubKeyI
pk
    (PayWitnessScriptHash{}, ScriptHashInput SimpleInput
i ScriptOutput
o) -> forall {m :: * -> *}.
Monad m =>
SimpleInput -> ScriptOutput -> m WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o
    (PayScriptHash{}, ScriptHashInput SimpleInput
i ScriptOutput
o) -> forall {m :: * -> *}.
Monad m =>
SimpleInput -> ScriptOutput -> m WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o
    (ScriptOutput, ScriptInput)
_ -> forall a b. a -> Either a b
Left String
"calcWitnessProgram: Invalid segwit SigInput"
  where
    p2wpkh :: TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig = forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramPKH -> WitnessProgram
P2WPKH 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 = forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. WitnessProgramSH -> WitnessProgram
P2WSH 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 (forall a (m :: * -> *). (Serial a, MonadPut m) => a -> m ()
serialize PubKeyI
k)]
    SpendMulSig [TxSignature]
sigs -> ByteString
"" forall a. a -> [a] -> [a]
: 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