tx-signer.go

package wallet

import (
	"bytes"
	"github.com/btccom/mrsign/bip32util"
	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
	"github.com/btcsuite/fastsha256"
	"github.com/pkg/errors"
	"sync"
	"fmt"
)

// SignatureVersion represents the version of the transaction
// signature hashing algorithm.
type SignatureVersion int

const (
	// SigV0 is the legacy signature hashing algorithm
	SigV0 SignatureVersion = 0

	// SigV1 is the signature hashing algorithm specified
	// in BIP143, and used in segwit scripts.
	SigV1 SignatureVersion = 1
)

// allowedP2sh represents the list of script types
// we are prepared to accept if the script is P2SH.
var allowedP2sh = map[txscript.ScriptClass]bool{
	txscript.WitnessPubKeyHashTy: true,
	txscript.WitnessScriptHashTy: true,
	txscript.PubKeyHashTy:        true,
	txscript.PubKeyTy:            true,
	txscript.MultiSigTy:          true,
}

// canSign is the list of scriptTypes we can directly
// sign
var canSign = map[txscript.ScriptClass]bool{
	txscript.PubKeyHashTy: true,
	txscript.PubKeyTy:     true,
	txscript.MultiSigTy:   true,
}

// IsAllowedP2shType returns whether the provided
// script type is a valid P2SH redeem script.
func IsAllowedP2shType(sc txscript.ScriptClass) bool {
	_, ok := allowedP2sh[sc]
	return ok
}

// CanSignType returns whether the provided script
// type can be directly signed.
func CanSignType(sc txscript.ScriptClass) bool {
	_, ok := canSign[sc]
	return ok
}

// InputSignData is the data required to verify/sign a
// transaction.
type InputSignData struct {
	Path  *bip32util.Path
	TxOut *wire.TxOut

	RedeemScript  []byte
	WitnessScript []byte
	SigVersion    SignatureVersion
}

func (data *InputSignData) Init(txOut *wire.TxOut, rs *[]byte, ws *[]byte) error {
	isP2SH, isP2WSH, sigVersion, err := CheckScriptDetails(txOut.PkScript, rs, ws)
	if err != nil {
		return err
	}

	if isP2SH {
		data.RedeemScript = *rs
	}

	if isP2WSH {
		data.WitnessScript = *ws
	}

	data.TxOut = txOut
	data.SigVersion = sigVersion

	return nil
}

func CheckScriptDetails(spk []byte, rs *[]byte, ws *[]byte) (bool, bool, SignatureVersion, error) {
	spkType := txscript.GetScriptClass(spk)
	if spkType != txscript.ScriptHashTy && !CanSignType(spkType) {
		return false, false, SigV0, errors.New("Unsupported scriptPubKey type")
	}

	p2sh := false
	p2wsh := false
	sigVersion := SigV0
	signType := spkType

	if signType == txscript.ScriptHashTy {
		if nil == rs {
			return false, false, SigV0, errors.New("Redeem script required by scriptPubKey")
		}

		rsType := txscript.GetScriptClass(*rs)
		if !IsAllowedP2shType(rsType) {
			return false, false, SigV0, errors.New("Unsupported redeemScript type")
		}

		p2sh = true
		signType = rsType
	}

	if signType == txscript.WitnessPubKeyHashTy {
		sigVersion = SigV1
		signType = txscript.PubKeyHashTy
	} else if signType == txscript.WitnessScriptHashTy {
		if nil == ws {
			return false, false, SigV0, errors.New("Witness script required but not provided")
		}

		p2wsh = true
	}

	return p2sh, p2wsh, sigVersion, nil
}

func ParseScript(pkScript []byte, chainParams *chaincfg.Params) (txscript.ScriptClass, [][]byte, []btcutil.Address, int, error) {
	scriptType, vAddr, nReqSigs, err := txscript.ExtractPkScriptAddrs(pkScript, chainParams)
	if err != nil {
		return txscript.NonStandardTy, nil, nil, 0, err
	}

	nAddr := len(vAddr)
	var vSol [][]byte

	switch scriptType {
	case txscript.ScriptHashTy, txscript.WitnessScriptHashTy, txscript.WitnessPubKeyHashTy, txscript.PubKeyTy, txscript.PubKeyHashTy:
		if nAddr != 1 {
			if scriptType == txscript.PubKeyTy {
				return txscript.NonStandardTy, nil, nil, 0, errors.Errorf("Parsed script was %s but had an invalid public key", txscript.PubKeyTy)
			} else {
				return txscript.NonStandardTy, nil, nil, 0, errors.Errorf("Parsed script was %s but was invalid", scriptType)
			}
		}

		vSol = make([][]byte, 1)
		vSol[0] = vAddr[0].ScriptAddress()
	case txscript.MultiSigTy:
		nPubkey, _, err := txscript.CalcMultiSigStats(pkScript)
		if err != nil {
			return txscript.NonStandardTy, nil, nil, 0, err
		}
		if nAddr != nPubkey {
			return txscript.NonStandardTy, nil, nil, 0, errors.Errorf("An invalid public key was found in the multisig script")
		}
		vSol = make([][]byte, nPubkey)
		for i := 0; i < nPubkey; i++ {
			vSol[i] = vAddr[i].ScriptAddress()
		}
		break
	default:
		panic(errors.New("not sure whats going on"))
	}

	return scriptType, vSol, vAddr, nReqSigs, nil
}

type ScriptData struct {
	Type      txscript.ScriptClass
	NumSigs   int
	Script    []byte
	Solution  [][]byte
	Addresses []btcutil.Address
}

func (sd *ScriptData) CanSign() bool {
	return CanSignType(sd.Type)
}

func (sd *ScriptData) IsAllowedP2SH() bool {
	return IsAllowedP2shType(sd.Type)
}

func (sd *ScriptData) Parse(pkScript []byte, params *chaincfg.Params) error {
	scriptType, vSol, vAddr, nSigs, err := ParseScript(pkScript, params)
	if err != nil {
		return err
	}

	sd.Type = scriptType
	sd.Script = pkScript
	sd.Solution = vSol
	sd.NumSigs = nSigs
	sd.Addresses = vAddr

	return nil
}

type TxSigner struct {
	sync.RWMutex
	params    *chaincfg.Params
	vSigner   map[int]*InputSigner
	tx        *wire.MsgTx
	sigHashes *txscript.TxSigHashes
}

func (signer *TxSigner) Init(params *chaincfg.Params, tx *wire.MsgTx) error {
	signer.Lock()
	defer signer.Unlock()
	signer.params = params
	signer.tx = tx
	signer.vSigner = make(map[int]*InputSigner, len(tx.TxIn))
	return nil
}

func (signer *TxSigner) Build() (*wire.MsgTx, error) {
	signer.RLock()
	defer signer.RUnlock()
	txCopy := *signer.tx

	for i := 0; i < len(txCopy.TxIn); i++ {
		if input, exists := signer.vSigner[i]; exists {
			sig, wit, err := input.SerializeSigs()
			if err != nil {
				return nil, errors.New("Fatal error - unable to serialize signatures")
			}

			txCopy.TxIn[i].SignatureScript = sig
			txCopy.TxIn[i].Witness = wit
		}
	}

	return &txCopy, nil
}

func (signer *TxSigner) GetSigHashes() *txscript.TxSigHashes {
	if nil == signer.sigHashes {
		signer.sigHashes = txscript.NewTxSigHashes(signer.tx)
	}

	return signer.sigHashes
}

func (signer *TxSigner) Input(nInput int, signData *InputSignData, lookupKey txscript.KeyClosure) (*InputSigner, error) {
	signer.Lock()
	defer signer.Unlock()

	if inputSigner, exists := signer.vSigner[nInput]; exists {
		return inputSigner, nil
	} else {
		numInputs := len(signer.tx.TxIn)
		if nInput < 0 || nInput > numInputs {
			return nil, errors.Errorf("Requested out of range input %d, but transaction has %d", nInput, numInputs)
		}

		inputSigner := &InputSigner{}
		var sigHashes *txscript.TxSigHashes
		if signData.SigVersion == 1 {
			sigHashes = signer.GetSigHashes()
		}

		err := inputSigner.Init(signer.params, sigHashes, signer.tx, nInput, signData, lookupKey)
		if err != nil {
			return nil, err
		}

		signer.vSigner[nInput] = inputSigner

		return inputSigner, nil
	}
}

type InputSigner struct {
	sync.RWMutex
	tx            *wire.MsgTx
	nInput        int
	params        *chaincfg.Params
	ScriptPubKey  *ScriptData
	RedeemScript  *ScriptData
	WitnessScript *ScriptData
	sigHashes     *txscript.TxSigHashes
	signData      *InputSignData
	sigVersion    SignatureVersion
	signScript    *ScriptData
	keyClosure    txscript.KeyClosure
	requiredSigs  int
	sigs          map[int]*txscript.SignatureInfo
	keys          map[int]*txscript.PublicKeyInfo
}

func (input *InputSigner) Init(params *chaincfg.Params, sigHashes *txscript.TxSigHashes,
	tx *wire.MsgTx, nInput int, inputData *InputSignData, lookupKey txscript.KeyClosure) error {

	input.Lock()
	defer input.Unlock()
	if nInput < 0 || nInput > len(tx.TxIn) {
		return errors.Errorf("Input %d does not exist in transaction", nInput)
	}

	txin := tx.TxIn[nInput]

	var sigVer SignatureVersion
	var signScript *ScriptData
	var spk, rs, ws *ScriptData
	var err error

	var redeemScript []byte
	if inputData.RedeemScript != nil {
		redeemScript = inputData.RedeemScript
	}
	var witnessScript []byte
	if inputData.RedeemScript != nil {
		witnessScript = inputData.WitnessScript
	}

	sigVer, _, spk, rs, ws, signScript, err = Solve(params,
		inputData.TxOut.PkScript, txin.SignatureScript, &txin.Witness,
		&redeemScript, &witnessScript)

	if err != nil {
		return err
	}

	flags := txscript.StandardVerifyFlags
	engine, err := txscript.NewEngine(spk.Script, tx, nInput, flags, nil, nil, inputData.TxOut.Value)
	if err != nil {
		panic(err)
	}

	requiredSigs, sigs, keys, err := ExtractSignatures(signScript, engine)

	if err != nil {
		return err
	}

	input.tx = tx
	input.params = params
	input.nInput = nInput
	input.signScript = signScript
	input.ScriptPubKey = spk
	input.RedeemScript = rs
	input.WitnessScript = ws
	input.signData = inputData
	input.sigVersion = sigVer
	input.requiredSigs = requiredSigs
	input.sigHashes = sigHashes
	input.sigs = sigs
	input.keys = keys
	input.keyClosure = lookupKey
	fmt.Printf("init input %d with %d signatures\n", nInput, len(sigs))
	return nil
}

func (input *InputSigner) Sign(sigHash txscript.SigHashType) ([]*txscript.SignatureInfo, error) {
	signed := 0
	input.Lock()
	defer input.Unlock()

	sigs := make([]*txscript.SignatureInfo, 0, len(input.signScript.Addresses))
	for i, solution := range input.signScript.Addresses {
		if input.sigs[i] != nil {
			signed++
			continue
		}

		if signed < input.signScript.NumSigs {
			key, _, err := input.keyClosure.GetKey(solution)
			if err != nil {
				continue
			}

			sigBytes, err := Sign(input.tx, input.nInput, input.signScript.Script, input.signData.TxOut.Value,
				sigHash, input.sigVersion, key, input.sigHashes)
			if err != nil {
				return nil, err
			}
			signature, err := btcec.ParseDERSignature(sigBytes, btcec.S256())
			if err != nil {
				return nil, err
			}
			input.sigs[i] = &txscript.SignatureInfo{
				HashType:  sigHash,
				Signature: signature,
			}

			sigs = append(sigs, input.sigs[i])
			signed++
		}
	}

	if signed == 0 {
		return nil, errors.Errorf("Unable to sign input %d", input.nInput)
	}
	fmt.Printf("Signing input %d (got %d)\n", input.nInput, signed)

	return sigs, nil
}

func (input *InputSigner) SerializeSigs() ([]byte, [][]byte, error) {
	input.RLock()
	defer input.RUnlock()

	return SerializeSignature(input.ScriptPubKey, input.RedeemScript, input.WitnessScript, input.sigs, input.keys)
}

// Takes the sigChunks to extract from, plus the *ScriptData and sigVersion
// to extract and verify signatures
func ExtractSignatures(data *ScriptData, engine *txscript.Engine) (
	int, map[int]*txscript.SignatureInfo, map[int]*txscript.PublicKeyInfo, error) {

	// maybe required sigs is already known by this point. NB for later.
	// requiredSigs, signatures, publicKeys
	var requiredSigs int
	var sigs map[int]*txscript.SignatureInfo
	var keys map[int]*txscript.PublicKeyInfo

	idx := 0
	ops, err := engine.ExecuteSignOp(true)
	if err != nil {
		return 0, nil, nil, err
	}

	path, pos, err := ops.IsComplete(idx)
	if err != nil {
		return 0, nil, nil, err
	}

	complete := path && pos

	switch data.Type {
	case txscript.PubKeyHashTy:
		requiredSigs = 1
		if complete {
			opKeys, opSigs, err := ops.GetSignOps(idx)
			if err != nil {
				return 0, nil, nil, err
			}

			// validate script here
			sigs = make(map[int]*txscript.SignatureInfo, 1)
			sigs[0] = opSigs[0]

			keys = make(map[int]*txscript.PublicKeyInfo, 1)
			keys[0] = opKeys[0]
		}
	case txscript.PubKeyTy:
		requiredSigs = 1
		if complete {
			opKeys, opSigs, err := ops.GetSignOps(idx)
			if err != nil {
				return 0, nil, nil, err
			}

			// validate script here
			sigs = make(map[int]*txscript.SignatureInfo, 1)
			sigs[0] = opSigs[0]

			keys = make(map[int]*txscript.PublicKeyInfo, 1)
			keys[0] = opKeys[0]
		}
	case txscript.MultiSigTy:
		var err error
		var nKeys int
		nKeys, requiredSigs, err = txscript.CalcMultiSigStats(data.Script)
		if err != nil {
			return 0, nil, nil, err
		}

		var opKeys map[int]*txscript.PublicKeyInfo
		var opSigs map[int]*txscript.SignatureInfo
		if complete {
			opKeys, opSigs, err = ops.GetSignOps(idx)
			if err != nil {
				return 0, nil, nil, err
			}
		} else {
			opKeys, opSigs, err = ops.GetIncompleteOps(idx)
			if err != nil {
				return 0, nil, nil, err
			}
		}

		sigs = make(map[int]*txscript.SignatureInfo, requiredSigs)
		keys = make(map[int]*txscript.PublicKeyInfo, nKeys)
		for i := 0; i < nKeys; i++ {
			if sig, ok := opSigs[i]; ok {
				sigs[i] = sig
			}
			keys[i] = opKeys[i]
		}

	default:
		return 0, nil, nil, errors.New("Unsupported script type")
	}

	return requiredSigs, sigs, keys, nil
}

func serializeSolution(scriptType txscript.ScriptClass, sigs map[int]*txscript.SignatureInfo, pubkeys map[int]*txscript.PublicKeyInfo) ([][]byte, error) {
	var data [][]byte
	nKeys := len(pubkeys)
	nSigs := len(sigs)

	switch scriptType {
	case txscript.PubKeyTy:
		if len(sigs) == 1 {
			data = make([][]byte, 1)
			data = append(data, sigs[0].Serialize())
		}
	case txscript.PubKeyHashTy:
		if len(sigs) == 1 && nKeys == 1 {
			pubKey, _ := pubkeys[0].Serialize()

			data = make([][]byte, 2)
			data = append(data, [][]byte{sigs[0].Serialize(), pubKey}...)
			data = make([][]byte, 1+nSigs)
		}
	case txscript.MultiSigTy:
		data = append(data, []byte{})
		for i := 0; i < nKeys; i++ {
			if sigs[i] != nil {
				data = append(data, sigs[i].Serialize())
			}
		}
	default:
		return nil, errors.New("Unsupported script type used in serializeSolution")
	}

	return data, nil
}

// Errors from this function should be considered a bug if it does not
// work with scripts that were successfully Solve()'d.
func SerializeSignature(spk *ScriptData, rs *ScriptData, ws *ScriptData, sigs map[int]*txscript.SignatureInfo, pubkeys map[int]*txscript.PublicKeyInfo) ([]byte, wire.TxWitness, error) {
	p2sh := false

	var sigData [][]byte = nil
	var witnessData [][]byte = nil
	var err error

	solution := spk

	if solution.CanSign() {
		sigData, err = serializeSolution(spk.Type, sigs, pubkeys)
		if err != nil {
			return nil, nil, err
		}
	}

	if solution.Type == txscript.ScriptHashTy {
		p2sh = true
		if rs.CanSign() {
			sigData, err = serializeSolution(rs.Type, sigs, pubkeys)
			if err != nil {
				return nil, nil, err
			}
		}
		solution = rs
	}

	if solution.Type == txscript.WitnessPubKeyHashTy {
		witnessData, err = serializeSolution(txscript.PubKeyHashTy, sigs, pubkeys)
		if err != nil {
			return nil, nil, err
		}
	} else if solution.Type == txscript.WitnessScriptHashTy {
		if ws.CanSign() {
			witnessData, err = serializeSolution(txscript.PubKeyHashTy, sigs, pubkeys)
			if err != nil {
				return nil, nil, err
			}
			witnessData = append(witnessData, ws.Script)
		}
		solution = ws
	}

	if p2sh {
		sigData = append(sigData, rs.Script)
	}

	script, err := PushDataToScript(sigData)
	if err != nil {
		return nil, nil, err
	}

	return script, wire.TxWitness(witnessData), nil
}

func PushDataToScript(pushDatas [][]byte) ([]byte, error) {
	builder := txscript.NewScriptBuilder()
	for i, l := 0, len(pushDatas); i < l; i++ {
		builder.AddData(pushDatas[i])
	}
	return builder.Script()
}

func Sign(tx *wire.MsgTx, nIn int, subScript []byte, amount int64, sigHashType txscript.SigHashType,
	sigVer SignatureVersion, key *btcec.PrivateKey, sigHashes *txscript.TxSigHashes) ([]byte, error) {
	if sigVer == SigV0 {
		return txscript.RawTxInSignature(tx, nIn, subScript, sigHashType, key)
	} else {
		return txscript.RawTxInWitnessSignature(tx, sigHashes, nIn, amount, subScript, sigHashType, key)
	}
}

// Solve takes a pkScript, a sigScript (possibly empty), a witness (possibly
// empty), and a InputSignData struct, and returns the sigVersion, the chunks
// of data containing signatures (from scriptSig or witness), the *ScriptData
// for the scriptPubKey, redeemScript, and witnessScript, and also the 'sign
// script' which is directly signed in the signature hash.
func Solve(params *chaincfg.Params, pkScript []byte, sigScript []byte, witness *wire.TxWitness, rsScript *[]byte, wsScript *[]byte) (
	SignatureVersion, [][]byte, *ScriptData, *ScriptData, *ScriptData, *ScriptData, error) {

	sigVersion := SigV0

	var sigChunks [][]byte
	var scriptPubKey *ScriptData
	var redeemScript *ScriptData
	var witnessScript *ScriptData
	var solution *ScriptData = &ScriptData{}

	err := solution.Parse(pkScript, params)
	if err != nil {
		return SigV0, nil, nil, nil, nil, nil, err
	}

	// Check supported 1: must be P2SH, P2WSH, P2WPKH, or a normal script
	if solution.Type != txscript.ScriptHashTy && !solution.IsAllowedP2SH() {
		return SigV0, nil, nil, nil, nil, nil, errors.Errorf("ScriptPubKey not supported, was type %s", solution.Type)
	}

	scriptPubKey = solution
	if solution.CanSign() {
		sigChunks, err = txscript.PushedData(sigScript)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}
	}

	if solution.Type == txscript.ScriptHashTy {
		// Take all data pushed, redeemScript might be at the end.
		chunks, err := txscript.PushedData(sigScript)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}

		// Look for the redeemScript in signData or the input script.
		rs, err := findScriptAndCheck(chunks, rsScript)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}

		// Qualify the redeemScript against the scriptPubKey
		redeemScriptHash := btcutil.Hash160(rs)
		if !bytes.Equal(redeemScriptHash, solution.Solution[0]) {
			return SigV0, nil, nil, nil, nil, nil, errors.Errorf("The scriptPubKey indicates that our redeemScript is wrong.")
		}

		// Parse the redeemScript
		redeemScript = &ScriptData{}
		err = redeemScript.Parse(rs, params)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}

		// Check supported 2: P2SH scripts must pass this
		if !redeemScript.IsAllowedP2SH() {
			return SigV0, nil, nil, nil, nil, nil, errors.Errorf("RedeemScript not supported, was type %s", solution.Type)
		}

		// Set sigChunks (chunks, but drop the last item)
		sigChunks = removeLast(chunks)

		// Finally, update solution
		solution = redeemScript
	}

	if solution.Type == txscript.WitnessPubKeyHashTy {
		sigVersion = SigV1
		situAddr, err := btcutil.NewAddressPubKeyHash(solution.Solution[0], params)
		situScript, err := txscript.PayToAddrScript(situAddr)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}

		witnessKeyHash := &ScriptData{}
		err = witnessKeyHash.Parse(situScript, params)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}

		solution = witnessKeyHash
	} else if solution.Type == txscript.WitnessScriptHashTy {
		sigVersion = SigV1

		chunks := [][]byte(*witness)
		// Look for the redeemScript in signData or the input script.
		ws, err := findScriptAndCheck(chunks, wsScript)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}

		// Qualify the redeemScript against the scriptPubKey
		witnessScriptHash := fastsha256.Sum256(ws)
		if &witnessScriptHash[0] != &solution.Solution[0][0] {
			return SigV0, nil, nil, nil, nil, nil, errors.Errorf("Based on previous information, the witnessScript seems incorrect.")
		}

		// Parse the redeemScript
		witnessScript = &ScriptData{}
		err = witnessScript.Parse(ws, params)
		if err != nil {
			return SigV0, nil, nil, nil, nil, nil, err
		}

		// Set sigChunks (chunks, but drop the last item)

		sigChunks = removeLast(chunks)
		solution = redeemScript
	}

	return sigVersion, sigChunks, scriptPubKey, redeemScript, witnessScript, solution, nil
}

// This function makes it convenient to take the redeemScript/witnessScript
// from either a decompiled scriptSig/witness or if provided as the commitedScript.
// If the committed script is provided and the scriptData non-empty, they are
// compared for consistency.
func findScriptAndCheck(scriptData [][]byte, committedScript *[]byte) ([]byte, error) {
	size := len(scriptData)
	if size > 0 {
		elem := scriptData[size-1]

		if committedScript != nil {
			if !bytes.Equal(elem, *committedScript) {
				return nil, errors.Errorf("The last element of scriptData was not the expected committed script.")
			}
		}
		return elem, nil
	} else {
		if committedScript == nil {
			return nil, errors.Errorf("No committed script was provided, and scriptData was empty.")
		}
		return *committedScript, nil
	}
}

// removes last []byte from a [][]byte
func removeLast(chunks [][]byte) [][]byte {
	length := len(chunks)
	if length > 0 {
		return chunks[:length-1]
	} else {
		return make([][]byte, 0, 0)
	}
}