Transaction.hs

{-|
  Arbitrary types for Network.Haskoin.Transaction.
-}
module Network.Haskoin.Test.Transaction where
import           Control.Monad
import qualified Data.ByteString             as BS
import           Data.Either                 (fromRight)
import           Data.List                   (nub, nubBy, permutations)
import           Data.Word                   (Word64)
import           Network.Haskoin.Address
import           Network.Haskoin.Constants
import           Network.Haskoin.Keys.Common
import           Network.Haskoin.Script
import           Network.Haskoin.Test.Crypto
import           Network.Haskoin.Test.Keys
import           Network.Haskoin.Test.Script
import           Network.Haskoin.Test.Util
import           Network.Haskoin.Transaction
import           Test.QuickCheck

newtype TestCoin = TestCoin { getTestCoin :: Word64 }
    deriving (Eq, Show)

instance Coin TestCoin where
    coinValue = getTestCoin

-- | Arbitrary transaction hash (for non-existent transaction).
arbitraryTxHash :: Gen TxHash
arbitraryTxHash = TxHash <$> arbitraryHash256

-- | Arbitrary amount of Satoshi as 'Word64' (Between 1 and 21e14)
arbitrarySatoshi :: Network -> Gen TestCoin
arbitrarySatoshi net = TestCoin <$> choose (1, getMaxSatoshi net)

-- | Arbitrary 'OutPoint'.
arbitraryOutPoint :: Gen OutPoint
arbitraryOutPoint = OutPoint <$> arbitraryTxHash <*> arbitrary

-- | Arbitrary 'TxOut'.
arbitraryTxOut :: Network -> Gen TxOut
arbitraryTxOut net =
    TxOut <$> (getTestCoin <$> arbitrarySatoshi net)
          <*> (encodeOutputBS <$> arbitraryScriptOutput net)

-- | Arbitrary 'TxIn'.
arbitraryTxIn :: Network -> Gen TxIn
arbitraryTxIn net =
    TxIn <$> arbitraryOutPoint
         <*> (encodeInputBS <$> arbitraryScriptInput net)
         <*> arbitrary

-- | Arbitrary transaction. Can be regular or with witnesses.
arbitraryTx :: Network -> Gen Tx
arbitraryTx net = oneof [arbitraryLegacyTx net, arbitraryWitnessTx net]

-- | Arbitrary regular transaction.
arbitraryLegacyTx :: Network -> Gen Tx
arbitraryLegacyTx net = arbitraryWLTx net False

-- | Arbitrary witness transaction (witness data is fake).
arbitraryWitnessTx :: Network -> Gen Tx
arbitraryWitnessTx net = arbitraryWLTx net True

arbitraryWLTx :: Network -> Bool -> Gen Tx
arbitraryWLTx net wit = do
    ni <- choose (0, 5)
    no <-
        if wit
            then choose (0, 5)
            else choose
                     ( if ni == 0
                           then 2
                           else 0
                     , 5 -- avoid witness case
                      )
    inps <- vectorOf ni (arbitraryTxIn net)
    outs <- vectorOf no (arbitraryTxOut net)
    let uniqueInps = nubBy (\a b -> prevOutput a == prevOutput b) inps
    w <- if wit then vectorOf (length uniqueInps) (listOf arbitraryBS) else return []
    Tx <$> arbitrary <*> pure uniqueInps <*> pure outs <*> pure w <*> arbitrary

-- | Arbitrary transaction containing only inputs of type 'SpendPKHash',
-- 'SpendScriptHash' (multisig) and outputs of type 'PayPKHash' and 'PaySH'.
-- Only compressed public keys are used.
arbitraryAddrOnlyTx :: Network -> Gen Tx
arbitraryAddrOnlyTx net = do
    ni <- choose (0, 5)
    no <- choose (0, 5)
    inps <- vectorOf ni (arbitraryAddrOnlyTxIn net)
    outs <- vectorOf no (arbitraryAddrOnlyTxOut net)
    Tx <$> arbitrary <*> pure inps <*> pure outs <*> pure [] <*> arbitrary

-- | Like 'arbitraryAddrOnlyTx' without empty signatures in the inputs.
arbitraryAddrOnlyTxFull :: Network -> Gen Tx
arbitraryAddrOnlyTxFull net = do
    ni <- choose (0, 5)
    no <- choose (0, 5)
    inps <- vectorOf ni (arbitraryAddrOnlyTxInFull net)
    outs <- vectorOf no (arbitraryAddrOnlyTxOut net)
    Tx <$> arbitrary <*> pure inps <*> pure outs <*> pure [] <*> arbitrary

-- | Arbitrary TxIn that can only be of type 'SpendPKHash' or 'SpendScriptHash'
-- (multisig). Only compressed public keys are used.
arbitraryAddrOnlyTxIn :: Network -> Gen TxIn
arbitraryAddrOnlyTxIn net = do
    inp <- oneof [arbitraryPKHashInput net, arbitraryMulSigSHInput net]
    TxIn <$> arbitraryOutPoint <*> pure (encodeInputBS inp) <*> arbitrary

-- | like 'arbitraryAddrOnlyTxIn' with no empty signatures.
arbitraryAddrOnlyTxInFull :: Network -> Gen TxIn
arbitraryAddrOnlyTxInFull net = do
    inp <-
        oneof [arbitraryPKHashInputFullC net, arbitraryMulSigSHInputFullC net]
    TxIn <$> arbitraryOutPoint <*> pure (encodeInputBS inp) <*> arbitrary

-- | Arbitrary 'TxOut' that can only be of type 'PayPKHash' or 'PaySH'.
arbitraryAddrOnlyTxOut :: Network -> Gen TxOut
arbitraryAddrOnlyTxOut net = do
    v <- getTestCoin <$> arbitrarySatoshi net
    out <- oneof [ arbitraryPKHashOutput, arbitrarySHOutput net ]
    return $ TxOut v $ encodeOutputBS out

-- | Arbitrary 'SigInput' with the corresponding private keys used
-- to generate the 'ScriptOutput' or 'RedeemScript'.
arbitrarySigInput :: Network -> Gen (SigInput, [SecKeyI])
arbitrarySigInput net =
    oneof
        [ arbitraryPKSigInput net >>= \(si, k) -> return (si, [k])
        , arbitraryPKHashSigInput  net >>= \(si, k) -> return (si, [k])
        , arbitraryMSSigInput net
        , arbitrarySHSigInput net
        ]

-- | Arbitrary 'SigInput' with a 'ScriptOutput' of type 'PayPK'.
arbitraryPKSigInput :: Network -> Gen (SigInput, SecKeyI)
arbitraryPKSigInput net = arbitraryAnyInput net False

-- | Arbitrary 'SigInput' with a 'ScriptOutput' of type 'PayPKHash'.
arbitraryPKHashSigInput :: Network -> Gen (SigInput, SecKeyI)
arbitraryPKHashSigInput net = arbitraryAnyInput net True

arbitraryAnyInput :: Network -> Bool -> Gen (SigInput, SecKeyI)
arbitraryAnyInput net pkh = do
    (k, p) <- arbitraryKeyPair
    let out | pkh = PayPKHash $ getAddrHash160 $ pubKeyAddr net p
            | otherwise = PayPK p
    (val, op, sh) <- arbitraryInputStuff net
    return (SigInput out val op sh Nothing, k)

arbitraryInputStuff :: Network -> Gen (Word64, OutPoint, SigHash)
arbitraryInputStuff net = do
    val <- getTestCoin <$> arbitrarySatoshi net
    op <- arbitraryOutPoint
    sh <- arbitraryValidSigHash net
    return (val, op, sh)

-- | Arbitrary 'SigInput' with a 'ScriptOutput' of type 'PayMulSig'.
arbitraryMSSigInput :: Network -> Gen (SigInput, [SecKeyI])
arbitraryMSSigInput net = do
    (m, n) <- arbitraryMSParam
    ks <- vectorOf n arbitraryKeyPair
    let out = PayMulSig (map snd ks) m
    (val, op, sh) <- arbitraryInputStuff net
    perm <- choose (0, n - 1)
    let ksPerm = map fst $ take m $ permutations ks !! perm
    return (SigInput out val op sh Nothing, ksPerm)

-- | Arbitrary 'SigInput' with 'ScriptOutput' of type 'PaySH' and a
-- 'RedeemScript'.
arbitrarySHSigInput :: Network -> Gen (SigInput, [SecKeyI])
arbitrarySHSigInput net = do
    (SigInput rdm val op sh _, ks) <- oneof
        [ f <$> arbitraryPKSigInput net
        , f <$> arbitraryPKHashSigInput net
        , arbitraryMSSigInput net
        ]
    let out = PayScriptHash $ getAddrHash160 $ p2shAddr net rdm
    return (SigInput out val op sh $ Just rdm, ks)
  where
    f (si, k) = (si, [k])

-- | Arbitrary 'Tx' (empty 'TxIn'), 'SigInputs' and private keys that can be
-- passed to 'signTx' or 'detSignTx' to fully sign the 'Tx'.
arbitrarySigningData :: Network -> Gen (Tx, [SigInput], [SecKeyI])
arbitrarySigningData net = do
    v <- arbitrary
    ni <- choose (1, 5)
    no <- choose (1, 5)
    sigis <- vectorOf ni (arbitrarySigInput net)
    let uSigis = nubBy (\(a, _) (b, _) -> sigInputOP a == sigInputOP b) sigis
    inps <- forM uSigis $ \(s, _) -> TxIn (sigInputOP s) BS.empty <$> arbitrary
    outs <- vectorOf no (arbitraryTxOut net)
    l <- arbitrary
    perm <- choose (0, length inps - 1)
    let tx = Tx v (permutations inps !! perm) outs [] l
        keys = concatMap snd uSigis
    return (tx, map fst uSigis, keys)

-- | Arbitrary transaction with empty inputs.
arbitraryEmptyTx :: Network -> Gen Tx
arbitraryEmptyTx net = do
    v    <- arbitrary
    no   <- choose (1,5)
    ni   <- choose (1,5)
    outs <- vectorOf no (arbitraryTxOut net)
    ops  <- vectorOf ni arbitraryOutPoint
    t    <- arbitrary
    s    <- arbitrary
    return $ Tx v (map (\op -> TxIn op BS.empty s) (nub ops)) outs [] t

-- | Arbitrary partially-signed transactions.
arbitraryPartialTxs ::
       Network -> Gen ([Tx], [(ScriptOutput, Word64, OutPoint, Int, Int)])
arbitraryPartialTxs net = do
    tx <- arbitraryEmptyTx net
    res <-
        forM (map prevOutput $ txIn tx) $ \op -> do
            (so, val, rdmM, prvs, m, n) <- arbitraryData
            txs <- mapM (singleSig so val rdmM tx op) prvs
            return (txs, (so, val, op, m, n))
    return (concatMap fst res, map snd res)
  where
    singleSig so val rdmM tx op prv = do
        sh <- arbitraryValidSigHash net
        let sigi = SigInput so val op sh rdmM
        return . fromRight (error "Colud not decode transaction") $
            signTx net tx [sigi] [prv]
    arbitraryData = do
        (m, n) <- arbitraryMSParam
        val <- getTestCoin <$> arbitrarySatoshi net
        nPrv <- choose (m, n)
        keys <- vectorOf n arbitraryKeyPair
        perm <- choose (0, length keys - 1)
        let pubKeys = map snd keys
            prvKeys = take nPrv $ permutations (map fst keys) !! perm
        let so = PayMulSig pubKeys m
        elements
            [ (so, val, Nothing, prvKeys, m, n)
            , ( PayScriptHash $ getAddrHash160 $ p2shAddr net so
              , val
              , Just so
              , prvKeys
              , m
              , n)
            ]