Merge pull request #402 from GambolingPangolin/sign-psbt

Adds PSBT signing capability

CHANGELOG.md

@@ -14,8 +14,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - Fixes bug in `finalizeTransaction`
 
 ### Added
+- Signing support for PSBTs
 - Helper function for merging PSBTs
-- Another PSBT test
+- More PSBT tests
 
 ## 0.20.5
 ### Added

scripts/PsbtSignTest.hs

@@ -0,0 +1,18 @@
+module Main (main) where
+
+import qualified Data.ByteString    as BS
+import           Data.Maybe         (fromMaybe)
+import qualified Data.Serialize     as S
+import           Data.Text          (pack)
+import           Haskoin            (PartiallySignedTransaction, SecKey)
+import qualified Haskoin            as H
+import           System.Environment (getArgs)
+
+main :: IO ()
+main = do
+    keyText <- pack . head <$> getArgs
+    let key = maybe (error "Unable to decode key") H.secKeyData $ H.fromWif H.btcRegTest keyText
+    BS.interact $ S.encode . either error (onPsbt key) . S.decode
+
+onPsbt :: SecKey -> PartiallySignedTransaction -> PartiallySignedTransaction
+onPsbt key = H.signPSBT H.btcRegTest (H.secKeySigner key)

scripts/create_complex_psbt_data.sh

@@ -1,39 +1,11 @@
 #!/usr/bin/env bash
 
-# Use Bitcoin Core 22.0 to create a complex PSBT test vector
+. ./scripts/lib.sh
 
-if ! command -v bitcoind &> /dev/null; then
-    echo "Please install bitcoind"
-    exit 1
-fi
+start_bitcoind
 
-if ! command -v jq &> /dev/null; then
-    echo "Please install jq"
-    exit 1
-fi
-
-datadir=$(mktemp -d "/tmp/bitcoind-regtest-XXXXXX")
-
-btc () {
-    bitcoin-cli -regtest -datadir=$datadir "$@"
-}
-
-bitcoind -regtest -datadir=$datadir -fallbackfee=0.0001000 -daemon
-
-while ! btc getblockchaininfo &> /dev/null; do
-    echo "Waiting for bitcoind"
-    sleep 5
-done
-
-passphrase="password"
-
-create_wallet () {
-    btc -named createwallet wallet_name=$1 passphrase=$passphrase &> /dev/null
-    btc -rpcwallet=$1 walletpassphrase $passphrase 3600
-}
-
-get_address () {
-    btc -rpcwallet="miner" getnewaddress
+sign_psbt () {
+    btc -rpcwallet=$1 -named walletprocesspsbt "psbt=$2" sign=true | jq -r .psbt
 }
 
 get_priv_key () {
@@ -41,36 +13,11 @@ get_priv_key () {
     btc -rpcwallet="miner" dumpprivkey $address
 }
 
-get_public_descriptor () {
-    btc getdescriptorinfo $1 | jq -r .descriptor
-}
-
-get_descriptor_address () {
-    btc deriveaddresses $1 | jq -r ".[0]"
-}
-
 get_canonical_descriptor () {
     checksum=$(btc getdescriptorinfo "$1" | jq -r .checksum)
     echo -n "$1#$checksum"
 }
 
-import_privkey () {
-    btc -rpcwallet=$1 importprivkey $2
-}
-
-import_descriptor () {
-    request=$(jq --null-input --arg "desc" $2 '[ { desc: $desc, timestamp: "now" } ]')
-    btc -rpcwallet=$1 importmulti "$request"
-}
-
-generate_blocks () {
-    btc generatetoaddress $1 $(btc -rpcwallet="miner" getnewaddress) &> /dev/null
-}
-
-sign_psbt () {
-    btc -rpcwallet=$1 -named walletprocesspsbt "psbt=$2" sign=true | jq -r .psbt
-}
-
 create_wallet "miner"
 create_wallet "p2pkh"
 create_wallet "p2sh-pk"
@@ -369,5 +316,4 @@ add_summary_field "final_tx" $final_tx
 echo $summary_entries | jq '. | from_entries' > /tmp/psbt_vectors.json
 echo "Done."
 
-btc stop
-rm -r $datadir
+stop_bitcoind

scripts/lib.sh

@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+
+# Use Bitcoin Core 22.0 to create a complex PSBT test vector
+
+if ! command -v bitcoind &> /dev/null; then
+    echo "Please install bitcoind"
+    exit 1
+fi
+
+if ! command -v jq &> /dev/null; then
+    echo "Please install jq"
+    exit 1
+fi
+
+export datadir=$(mktemp -d "/tmp/bitcoind-regtest-XXXXXX")
+
+btc () {
+    bitcoin-cli -regtest -datadir=$datadir "$@"
+}
+
+start_bitcoind () {
+    bitcoind -regtest -datadir=$datadir -fallbackfee=0.0001000 -txindex=1 -daemon
+    while ! btc getblockchaininfo &> /dev/null; do
+        echo "Waiting for bitcoind"
+        sleep 5
+    done
+}
+
+stop_bitcoind () {
+    btc stop
+    rm -r $datadir
+}
+
+create_wallet () {
+    passphrase="password"
+    btc -named createwallet wallet_name=$1 passphrase=$passphrase &> /dev/null
+    btc -rpcwallet=$1 walletpassphrase $passphrase 3600
+}
+
+get_address () {
+    btc -rpcwallet="miner" getnewaddress
+}
+
+get_public_descriptor () {
+    btc getdescriptorinfo $1 | jq -r .descriptor
+}
+
+get_descriptor_address () {
+    btc deriveaddresses $1 | jq -r ".[0]"
+}
+
+import_descriptor () {
+    request=$(jq --null-input --arg "desc" $2 '[ { desc: $desc, timestamp: "now" } ]')
+    btc -rpcwallet=$1 importmulti "$request"
+}
+
+generate_blocks () {
+    btc generatetoaddress $1 $(btc -rpcwallet="miner" getnewaddress) &> /dev/null
+}

scripts/psbt_sign_test.sh

@@ -0,0 +1,179 @@
+#!/usr/bin/env bash
+
+. ./scripts/lib.sh
+
+cabal build
+start_bitcoind
+
+create_wallet "miner"
+create_wallet "payer"
+create_wallet "signer"
+
+generate_blocks 200
+
+
+address1=$(get_address)
+privkey1=$(btc -rpcwallet="miner" dumpprivkey $address1)
+pubkey1=$(btc -rpcwallet="miner" getaddressinfo $address1 | jq -r .pubkey)
+
+address2=$(get_address)
+privkey2=$(btc -rpcwallet="miner" dumpprivkey $address2)
+pubkey2=$(btc -rpcwallet="miner" getaddressinfo $address2 | jq -r .pubkey)
+
+address3=$(get_address)
+pubkey3=$(btc -rpcwallet="miner" getaddressinfo $address3 | jq -r .pubkey)
+
+sign () {
+    base64 -d |
+        cabal exec runghc -- ./scripts/PsbtSignTest.hs "$1" |
+        base64 -w0
+}
+
+build () {
+
+    descriptor=$1
+
+    signer_descriptor=$(get_public_descriptor $descriptor)
+    signer_addr=$(get_descriptor_address $signer_descriptor)
+
+    suitable_outputs=$(btc -rpcwallet=miner listunspent | jq ". | map(select(.spendable) | select(.amount > 1))")
+
+    txid=$(jq -r ".[0].txid" <<< $suitable_outputs)
+    vout=$(jq -r ".[0].vout" <<< $suitable_outputs)
+
+    funding_inputs=$(
+        jq --null-input \
+            --arg txid $txid \
+            --arg vout $vout \
+            '[{ txid: $txid, vout: $vout | tonumber }]'
+    )
+    funding_outputs=$(
+        jq --null-input \
+            --arg addr $signer_addr \
+            '{ ($addr): 1 }'
+           )
+    funding_options=$(jq --null-input "{ changePosition: 1 }")
+    funding_psbt=$(
+        btc -named -rpcwallet=miner walletcreatefundedpsbt \
+            "inputs=$funding_inputs" \
+            "outputs=$funding_outputs" \
+            "options=$funding_options"|
+            jq -r ".psbt"
+    )
+    signed_funding_psbt=$(
+        btc -rpcwallet=miner walletprocesspsbt $funding_psbt true | jq -r .psbt
+    )
+    funding_tx=$(btc -named finalizepsbt "psbt=$signed_funding_psbt" extract=true | jq -r .hex)
+    funding_txid=$(btc sendrawtransaction $funding_tx)
+
+    import_descriptor signer "$signer_descriptor" &> /dev/null
+    generate_blocks 10
+
+    signer_inputs=$(
+        jq --null-input \
+            --arg txid $funding_txid \
+            '[{ txid: $txid, vout: 0 }]'
+    )
+    signer_outputs=$(
+        jq --null-input \
+            --arg addr $(get_address) \
+            '{ ($addr): 0.8 }'
+    )
+    btc -named -rpcwallet=signer walletcreatefundedpsbt \
+        "inputs=$signer_inputs" \
+        "outputs=$signer_outputs" |
+        jq -r .psbt
+}
+
+echo "p2pkh"
+{
+    signed_psbt=$(build "pkh($pubkey1)" | sign $privkey1)
+    tx=$(btc -named finalizepsbt "psbt=$signed_psbt" extract=true | jq -r .hex)
+    btc sendrawtransaction $tx 1> /dev/null
+}
+
+generate_blocks 10
+
+echo "p2sh-pk"
+{
+    signed_psbt=$(build "sh(pk($pubkey1))" | sign $privkey1)
+    tx=$(btc -named finalizepsbt "psbt=$signed_psbt" extract=true | jq -r .hex)
+    btc sendrawtransaction $tx 1> /dev/null
+}
+
+generate_blocks 10
+
+echo "p2sh-ms"
+{
+    descriptor="sh(sortedmulti(2,$pubkey1,$pubkey2,$pubkey3))"
+    unsigned_psbt=$(build $descriptor)
+    signed_psbt_1=$(sign $privkey1 <<< $unsigned_psbt)
+    signed_psbt_2=$(sign $privkey2 <<< $unsigned_psbt)
+    psbt_list=$(
+        jq --null-input \
+            --arg psbt1 $signed_psbt_1 \
+            --arg psbt2 $signed_psbt_2 \
+            '[ $psbt1, $psbt2 ]'
+    )
+    merged_psbt=$(btc combinepsbt "$psbt_list")
+    tx=$(btc -named finalizepsbt "psbt=$merged_psbt" extract=true | jq -r .hex)
+    btc sendrawtransaction $tx 1> /dev/null
+}
+
+generate_blocks 10
+
+echo "p2sh-wpkh"
+{
+    signed_psbt=$(build "sh(wpkh($pubkey1))" | sign $privkey1)
+    tx=$(btc -named finalizepsbt "psbt=$signed_psbt" extract=true | jq -r .hex)
+    btc sendrawtransaction $tx 1> /dev/null
+}
+
+generate_blocks 10
+
+echo "p2sh-p2wsh-ms"
+{
+    descriptor="sh(wsh(sortedmulti(2,$pubkey1,$pubkey2,$pubkey3)))"
+    unsigned_psbt=$(build $descriptor)
+    signed_psbt_1=$(sign $privkey1 <<< $unsigned_psbt)
+    signed_psbt_2=$(sign $privkey2 <<< $unsigned_psbt)
+    psbt_list=$(
+        jq --null-input \
+            --arg psbt1 $signed_psbt_1 \
+            --arg psbt2 $signed_psbt_2 \
+            '[ $psbt1, $psbt2 ]'
+    )
+    merged_psbt=$(btc combinepsbt "$psbt_list")
+    tx=$(btc -named finalizepsbt "psbt=$merged_psbt" extract=true | jq -r .hex)
+    btc sendrawtransaction $tx 1> /dev/null
+}
+
+generate_blocks 10
+
+echo "p2wpkh"
+{
+    signed_psbt=$(build "wpkh($pubkey1)" | sign $privkey1)
+    tx=$(btc -named finalizepsbt "psbt=$signed_psbt" extract=true | jq -r .hex)
+    btc sendrawtransaction $tx 1> /dev/null
+}
+
+generate_blocks 10
+
+echo "p2wsh-ms"
+{
+    descriptor="wsh(sortedmulti(2,$pubkey1,$pubkey2,$pubkey3))"
+    unsigned_psbt=$(build $descriptor)
+    signed_psbt_1=$(sign $privkey1 <<< $unsigned_psbt)
+    signed_psbt_2=$(sign $privkey2 <<< $unsigned_psbt)
+    psbt_list=$(
+        jq --null-input \
+            --arg psbt1 $signed_psbt_1 \
+            --arg psbt2 $signed_psbt_2 \
+            '[ $psbt1, $psbt2 ]'
+    )
+    merged_psbt=$(btc combinepsbt "$psbt_list")
+    tx=$(btc -named finalizepsbt "psbt=$merged_psbt" extract=true | jq -r .hex)
+    btc sendrawtransaction $tx 1> /dev/null
+}
+
+stop_bitcoind

src/Haskoin/Transaction/Partial.hs

@@ -29,38 +29,59 @@ module Haskoin.Transaction.Partial
     , emptyPSBT
     , emptyInput
     , emptyOutput
+
+    -- ** Signing
+    , PsbtSigner
+    , getSignerKey
+    , secKeySigner
+    , xPrvSigner
+    , signPSBT
     ) where
 
-import           Control.Applicative        ((<|>))
+import           Control.Applicative         ((<|>))
 import           Control.DeepSeq
-import           Control.Monad              (guard, replicateM, void, foldM)
-import           Data.ByteString            (ByteString)
-import qualified Data.ByteString            as B
-import           Data.Bytes.Get             (runGetS)
-import           Data.Bytes.Put             (runPutS)
-import           Data.Bytes.Serial          (Serial (..))
-import           Data.HashMap.Strict        (HashMap)
-import qualified Data.HashMap.Strict        as HashMap
-import           Data.Hashable              (Hashable)
-import           Data.List                  (foldl')
-import           Data.Maybe                 (fromMaybe, isJust)
-import           Data.Serialize             (Get, Put, Serialize)
-import qualified Data.Serialize             as S
-import           GHC.Generics               (Generic)
-import           GHC.Word                   (Word32, Word8)
-import           Haskoin.Address            (Address (..), pubKeyAddr)
-import           Haskoin.Keys               (Fingerprint, KeyIndex, PubKeyI)
-import           Haskoin.Network            (VarInt (..), VarString (..),
-                                             putVarInt)
-import           Haskoin.Script             (Script (..), ScriptOp (..),
-                                             ScriptOutput (..), SigHash,
-                                             decodeOutput, decodeOutputBS,
-                                             encodeOutputBS, isPayScriptHash,
-                                             opPushData, toP2SH, toP2WSH)
-import           Haskoin.Transaction.Common (Tx (..), TxOut, WitnessStack,
-                                             outPointIndex, prevOutput,
-                                             scriptInput, scriptOutput)
-import           Haskoin.Util               (eitherToMaybe)
+import           Control.Monad               (foldM, guard, replicateM, void)
+import           Data.ByteString             (ByteString)
+import qualified Data.ByteString             as B
+import           Data.Bytes.Get              (runGetS)
+import           Data.Bytes.Put              (runPutS)
+import           Data.Bytes.Serial           (Serial (..))
+import           Data.Either                 (fromRight)
+import           Data.HashMap.Strict         (HashMap)
+import qualified Data.HashMap.Strict         as HM
+import qualified Data.HashMap.Strict         as HashMap
+import           Data.Hashable               (Hashable)
+import           Data.List                   (foldl')
+import           Data.Maybe                  (fromMaybe, isJust)
+import           Data.Serialize              (Get, Put, Serialize)
+import qualified Data.Serialize              as S
+import           GHC.Generics                (Generic)
+import           GHC.Word                    (Word32, Word8)
+import           Haskoin.Address             (Address (..), pubKeyAddr)
+import           Haskoin.Constants           (Network)
+import           Haskoin.Crypto              (SecKey, derivePubKey)
+import           Haskoin.Keys                (DerivPath, DerivPathI (Deriv),
+                                              Fingerprint, KeyIndex, PubKeyI,
+                                              SecKeyI (SecKeyI), XPrvKey,
+                                              derivePath, deriveXPubKey,
+                                              listToPath, pathToList,
+                                              pubKeyCompressed, pubKeyPoint,
+                                              xPrvKey, xPubFP)
+import           Haskoin.Network             (VarInt (..), VarString (..),
+                                              putVarInt)
+import           Haskoin.Script              (Script (..), ScriptOp (..),
+                                              ScriptOutput (..), SigHash,
+                                              decodeOutput, decodeOutputBS,
+                                              encodeOutputBS, encodeTxSig,
+                                              isPayScriptHash, opPushData,
+                                              sigHashAll, toP2SH, toP2WSH)
+import           Haskoin.Transaction.Builder (SigInput (..), makeSignature)
+import           Haskoin.Transaction.Common  (Tx (..), TxOut, WitnessStack,
+                                              outPointIndex, outValue,
+                                              prevOutput, scriptInput,
+                                              scriptOutput)
+import           Haskoin.Transaction.Segwit  (isSegwit)
+import           Haskoin.Util                (eitherToMaybe)
 
 -- | PSBT data type as specified in
 -- [BIP-174](https://github.com/bitcoin/bips/blob/master/bip-0174.mediawiki).
@@ -143,7 +164,7 @@ merge _ _ = Nothing
 -- @since 0.21.0
 mergeMany :: [PartiallySignedTransaction] -> Maybe PartiallySignedTransaction
 mergeMany (psbt : psbts) = foldM merge psbt psbts
-mergeMany _ = Nothing
+mergeMany _              = Nothing
 
 mergeInput :: Input -> Input -> Input
 mergeInput a b = Input
@@ -185,6 +206,151 @@ mergeOutput a b = Output
             outputUnknown a <> outputUnknown b
     }
 
+-- | A abstraction which covers varying key configurations.  Use the 'Semigroup'
+-- instance to create signers for sets of keys: `signerA <> signerB` can sign
+-- anything for which `signerA` or `signerB` could sign.
+--
+-- @since 0.21@
+newtype PsbtSigner = PsbtSigner
+    { unPsbtSigner ::
+        PubKeyI ->
+        Maybe (Fingerprint, DerivPath) ->
+        Maybe SecKey
+    }
+
+instance Semigroup PsbtSigner where
+    PsbtSigner signer1 <> PsbtSigner signer2 =
+        PsbtSigner $ \pubKey origin ->
+            signer1 pubKey origin <|> signer2 pubKey origin
+
+instance Monoid PsbtSigner where
+    mempty = PsbtSigner $ \_ _ -> Nothing
+
+-- | Fetch the secret key for the given 'PubKeyI' if possible.
+--
+-- @since 0.21@
+getSignerKey :: PsbtSigner -> PubKeyI -> Maybe (Fingerprint, DerivPath) -> Maybe SecKey
+getSignerKey = unPsbtSigner
+
+-- | This signer can sign for one key.
+--
+-- @since 0.21@
+secKeySigner :: SecKey -> PsbtSigner
+secKeySigner theSecKey = PsbtSigner signer
+  where
+    signer requiredKey _
+        | pubKeyPoint requiredKey == derivePubKey theSecKey = Just theSecKey
+        | otherwise = Nothing
+
+-- | This signer can sign with any child key, provided that derivation information is present.
+--
+-- @since 0.21@
+xPrvSigner ::
+    XPrvKey ->
+    -- | Origin data, if the input key is explicitly a child key
+    Maybe (Fingerprint, DerivPath) ->
+    PsbtSigner
+xPrvSigner xprv origin = PsbtSigner signer
+  where
+    signer pubKey (Just hdData)
+        | result@(Just theSecKey) <- maybe noOrigin onOrigin origin hdData
+        , pubKeyPoint pubKey == derivePubKey theSecKey = result
+    signer _ _ = Nothing
+
+    noOrigin (fp, path)
+        | thisFP == fp = Just $ deriveSecKey path
+        | otherwise = Nothing
+
+    onOrigin (originFP, originPath) (fp, path)
+        | thisFP == fp = Just $ deriveSecKey path
+        | originFP == fp =
+            deriveSecKey <$> adjustPath (pathToList originPath) (pathToList path)
+        | otherwise = Nothing
+
+    deriveSecKey path = xPrvKey $ derivePath path xprv
+
+    thisFP = xPubFP $ deriveXPubKey xprv
+
+    -- The origin path should be a prefix of the target path if we match the
+    -- origin fingerprint.  We need to remove this prefix.
+    adjustPath :: [KeyIndex] -> [KeyIndex] -> Maybe DerivPath
+    adjustPath (originIx : originTail) (thisIx : thisTail)
+        | originIx == thisIx = adjustPath originTail thisTail
+        | otherwise = Nothing
+    adjustPath [] thePath = Just $ listToPath thePath
+    adjustPath _ _ = Nothing
+
+-- | Update a PSBT with signatures when possible.  This function uses
+-- 'inputHDKeypaths' in order to calculate secret keys.
+--
+-- @since 0.21@
+signPSBT ::
+    Network ->
+    PsbtSigner ->
+    PartiallySignedTransaction ->
+    PartiallySignedTransaction
+signPSBT net signer psbt =
+    psbt
+        { inputs = addSigsForInput net signer tx <$> zip [0 :: Int ..] (inputs psbt)
+        }
+  where
+    tx = unsignedTransaction psbt
+
+addSigsForInput :: Network -> PsbtSigner -> Tx -> (Int, Input) -> Input
+addSigsForInput net signer tx (ix, input) =
+    maybe input (onPrevTxOut net signer tx ix input) $
+        (Left <$> nonWitnessUtxo input) <|> (Right <$> witnessUtxo input)
+
+onPrevTxOut ::
+    Network ->
+    PsbtSigner ->
+    Tx ->
+    Int ->
+    Input ->
+    Either Tx TxOut ->
+    Input
+onPrevTxOut net signer tx ix input prevTxData =
+    input
+        { partialSigs = newSigs <> partialSigs input
+        }
+  where
+    newSigs = HM.mapWithKey sigForInput sigKeys
+    sigForInput thePubKey theSecKey =
+        encodeTxSig . makeSignature net tx ix theSigInput $
+            SecKeyI theSecKey (pubKeyCompressed thePubKey)
+
+    theSigInput =
+        SigInput
+            { -- Must be the segwit input script for segwit spends (even nested)
+              sigInputScript = fromMaybe theInputScript segwitInput
+            , sigInputValue = outValue prevTxOut
+            , sigInputOP = thePrevOutPoint
+            , sigInputSH = fromMaybe sigHashAll $ sigHashType input
+            , -- Must be the witness script for segwit spends (even nested)
+              sigInputRedeem = theWitnessScript <|> theRedeemScript
+            }
+
+    prevTxOut = either ((!! (fromIntegral . outPointIndex) thePrevOutPoint) . txOut) id prevTxData
+    thePrevOutPoint = prevOutput $ txIn tx !! ix
+
+    segwitInput = justWhen isSegwit theInputScript <|> (justWhen isSegwit =<< theRedeemScript)
+
+    theInputScript = fromRight inputScriptErr $ (decodeOutputBS . scriptOutput) prevTxOut
+    inputScriptErr = error "addSigsForInput: Unable to decode input script"
+
+    theRedeemScript = case decodeOutput <$> inputRedeemScript input of
+        Just (Right script) -> Just script
+        Just Left{} -> error "addSigsForInput: Unable to decode redeem script"
+        _ -> Nothing
+
+    theWitnessScript = case decodeOutput <$> inputWitnessScript input of
+        Just (Right script) -> Just script
+        Just Left{} -> error "addSigsForInput: Unable to decode witness script"
+        _ -> Nothing
+
+    sigKeys = HM.mapMaybeWithKey getSignerKey $ inputHDKeypaths input
+    getSignerKey pubKey (fp, ixs) = unPsbtSigner signer pubKey $ Just (fp, listToPath ixs)
+
 -- | Take partial signatures from all of the 'Input's and finalize the signature.
 complete :: PartiallySignedTransaction
          -> PartiallySignedTransaction
@@ -714,3 +880,6 @@ word8Enum n = Left n
 
 whenJust :: Monad m => (a -> m ()) -> Maybe a -> m ()
 whenJust = maybe (return ())
+
+justWhen :: (a -> Bool) -> a -> Maybe a
+justWhen test x = if test x then Just x else Nothing

test/Haskoin/Transaction/PartialSpec.hs

@@ -2,27 +2,25 @@
 
 module Haskoin.Transaction.PartialSpec (spec) where
 
-import           Data.ByteString             (ByteString)
+import           Data.ByteString        (ByteString)
 import           Data.Bytes.Get
 import           Data.Bytes.Put
 import           Data.Bytes.Serial
-import           Data.Either                 (fromRight, isLeft, isRight)
-import           Data.HashMap.Strict         (fromList, singleton)
-import           Data.Maybe                  (fromJust, isJust)
-import           Data.Serialize              as S
-import           Data.Text                   (Text)
-import           Test.HUnit                  (Assertion, assertBool,
-                                              assertEqual)
+import           Data.Either            (fromRight, isLeft, isRight)
+import           Data.HashMap.Strict    (fromList, singleton)
+import           Data.Maybe             (fromJust, isJust)
+import           Data.Serialize         as S
+import           Data.Text              (Text)
+import           Test.HUnit             (Assertion, assertBool, assertEqual)
 import           Test.Hspec
 import           Test.QuickCheck
 
-import           Control.Monad               ((<=<))
-import           Data.Aeson                  (FromJSON, parseJSON, withObject,
-                                              (.:))
-import           Data.Bifunctor              (first)
-import           Data.ByteString.Base64      (decodeBase64)
-import qualified Data.Text                   as Text
-import           Data.Text.Encoding          (encodeUtf8)
+import           Control.Monad          ((<=<))
+import           Data.Aeson             (FromJSON, parseJSON, withObject, (.:))
+import           Data.Bifunctor         (first)
+import           Data.ByteString.Base64 (decodeBase64)
+import qualified Data.Text              as Text
+import           Data.Text.Encoding     (encodeUtf8)
 import           Haskoin.Address
 import           Haskoin.Constants
 import           Haskoin.Crypto
@@ -31,7 +29,7 @@ import           Haskoin.Script
 import           Haskoin.Transaction
 import           Haskoin.Util
 import           Haskoin.Util.Arbitrary
-import           Haskoin.UtilSpec            (readTestFile)
+import           Haskoin.UtilSpec       (readTestFile)
 
 spec :: Spec
 spec = describe "partially signed bitcoin transaction unit tests" $ do
@@ -59,7 +57,7 @@ spec = describe "partially signed bitcoin transaction unit tests" $ do
     it "encodes and decodes psbt with final witness script" $
         (fmap (encodeHex . S.encode) . decodeHexPSBT) validVec7Hex == Right validVec7Hex
     it "handles complex psbts correctly" complexPsbtTest
-
+    it "calculates keys properly" psbtSignerTest
 
 vec2Test :: Assertion
 vec2Test = do
@@ -116,7 +114,6 @@ vec4Test = do
 vec5Test :: Assertion
 vec5Test = do
     psbt <- decodeHexPSBTM "Cannot parse validVec5" validVec5Hex
-    print psbt
     assertEqual "Correctly decode PSBT" expectedPsbt psbt
     let input = head $ inputs psbt
 
@@ -233,6 +230,38 @@ complexPsbtTest = do
         | Just{} <- witnessUtxo input = input { nonWitnessUtxo = Nothing }
         | otherwise = input
 
+psbtSignerTest :: Assertion
+psbtSignerTest = do
+    assertEqual "recover explicit secret key" (Just theSecKey) (getSignerKey signer thePubKey Nothing)
+    assertEqual
+        "recover key for origin path"
+        (Just originPathSecKey)
+        (getSignerKey signer originPathPubKey (Just (rootFP, originKeyPath)))
+    assertEqual
+        "recover key for direct path"
+        (Just directPathSecKey)
+        (getSignerKey signer directPathPubKey (Just (keyFP, directPath)))
+  where
+    signer = secKeySigner theSecKey <> xPrvSigner xprv (Just origin)
+
+    Just theSecKey = secKey "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+    thePubKey = PubKeyI { pubKeyPoint = derivePubKey theSecKey, pubKeyCompressed = True }
+
+    rootXPrv = makeXPrvKey "psbtSignerTest"
+    rootFP = xPubFP $ deriveXPubKey rootXPrv
+    xprv = derivePath keyPath rootXPrv
+    keyFP = xPubFP $ deriveXPubKey xprv
+    keyPath = Deriv :| 444
+    origin = (rootFP, keyPath)
+
+    originKeyPath = Deriv :| 444 :/ 0
+    originPathSecKey = xPrvKey $ derivePath originKeyPath rootXPrv
+    originPathPubKey = PubKeyI { pubKeyPoint = derivePubKey originPathSecKey, pubKeyCompressed = True }
+
+    directPath = Deriv :/ 1
+    directPathSecKey = xPrvKey $ derivePath directPath xprv
+    directPathPubKey = PubKeyI { pubKeyPoint = derivePubKey directPathSecKey, pubKeyCompressed = True }
+
 expectedOut :: ScriptOutput
 expectedOut = fromRight (error "could not decode expected output")
             . decodeOutputBS . fromJust $ decodeHex "a9143545e6e33b832c47050f24d3eeb93c9c03948bc787"