{-# 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 qualified Data.Serialize             as S
import           Haskoin.Constants
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
    _                      -> 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 sig :: TxSignature
sig key :: PubKeyI
key) -> [TxSignature -> ByteString
encodeTxSig TxSignature
sig, PubKeyI -> ByteString
forall a. Serialize a => a -> ByteString
S.encode PubKeyI
key]
    P2WSH (WitnessProgramSH stack :: WitnessStack
stack scr :: Script
scr) -> WitnessStack
stack WitnessStack -> WitnessStack -> WitnessStack
forall a. Semigroup a => a -> a -> a
<> [Script -> ByteString
forall a. Serialize a => a -> ByteString
S.encode Script
scr]
    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 net :: Network
net so :: ScriptOutput
so witness :: WitnessStack
witness = case ScriptOutput
so of
    PayWitnessPKHash _ | WitnessStack -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length WitnessStack
witness Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 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 <- ByteString -> Either String PubKeyI
forall a. Serialize a => ByteString -> Either String a
S.decode (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
!! 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 _ | Bool -> Bool
not (WitnessStack -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null WitnessStack
witness) -> do
        Script
redeemScript <- ByteString -> Either String Script
forall a. Serialize a => ByteString -> Either String a
S.decode (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
    _ | 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 "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 net :: Network
net = \case
    P2WPKH (WitnessProgramPKH sig :: TxSignature
sig key :: 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 st :: WitnessStack
st scr :: 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 _, [sigBS :: 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 _, [sigBS :: ByteString
sigBS, keyBS :: 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
<*> ByteString -> Either String PubKeyI
forall a. Serialize a => ByteString -> Either String a
S.decode ByteString
keyBS
            (PayMulSig _ _, "" : sigsBS :: 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
            _ -> String -> Either String SimpleInput
forall a b. a -> Either a b
Left "decodeWitnessInput: Non-standard script output"
    EmptyWitnessProgram -> String -> Either String (Maybe ScriptOutput, SimpleInput)
forall a b. a -> Either a b
Left "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 so :: ScriptOutput
so si :: ScriptInput
si = case (ScriptOutput
so, ScriptInput
si) of
    (PayWitnessPKHash{}, RegularInput (SpendPKHash sig :: TxSignature
sig pk :: PubKeyI
pk)) -> TxSignature -> PubKeyI -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig PubKeyI
pk
    (PayScriptHash{}, RegularInput (SpendPKHash sig :: TxSignature
sig pk :: PubKeyI
pk))    -> TxSignature -> PubKeyI -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh TxSignature
sig PubKeyI
pk
    (PayWitnessScriptHash{}, ScriptHashInput i :: SimpleInput
i o :: ScriptOutput
o)           -> SimpleInput -> ScriptOutput -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
SimpleInput -> ScriptOutput -> m WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o
    (PayScriptHash{}, ScriptHashInput i :: SimpleInput
i o :: ScriptOutput
o)                  -> SimpleInput -> ScriptOutput -> Either String WitnessProgram
forall (m :: * -> *).
Monad m =>
SimpleInput -> ScriptOutput -> m WitnessProgram
p2wsh SimpleInput
i ScriptOutput
o
    _ -> String -> Either String WitnessProgram
forall a b. a -> Either a b
Left "calcWitnessProgram: Invalid segwit SigInput"
  where
    p2wpkh :: TxSignature -> PubKeyI -> m WitnessProgram
p2wpkh sig :: 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 i :: SimpleInput
i o :: 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 sig :: TxSignature
sig       -> [TxSignature -> ByteString
f TxSignature
sig]
    SpendPKHash sig :: TxSignature
sig k :: PubKeyI
k -> [TxSignature -> ByteString
f TxSignature
sig, PubKeyI -> ByteString
forall a. Serialize a => a -> ByteString
S.encode PubKeyI
k]
    SpendMulSig sigs :: [TxSignature]
sigs  -> "" 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 TxSignatureEmpty = ""
    f sig :: TxSignature
sig              = TxSignature -> ByteString
encodeTxSig TxSignature
sig