Waste less memory if sighash optimizations are on

Legacy transaction deep copy code mandated by the Bitcoin protocol
caused large amounts of data to be copied needlessly. If the
optimization for SigHashAll is set in chaincfg/params.go, these
extra copies are avoided by directly writing the pkScript and
decorations to a buffer and then hashing to get a witness hash,
while using the cached hash for the prefix.

Fixes #126.

txscript/internal_test.go

@@ -12,7 +12,12 @@ interface.  The functions are only exported while the tests are being run.
 
 package txscript
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/decred/dcrd/chaincfg/chainhash"
+	"github.com/decred/dcrd/wire"
+)
 
 // TstMaxScriptSize makes the internal maxScriptSize constant available to the
 // test package.
@@ -84,6 +89,12 @@ func (vm *Engine) TstSetPC(script, off int) {
 	vm.scriptOff = off
 }
 
+// TstCalcSignatureHash is an exported version for testing.
+func TstCalcSignatureHash(script []parsedOpcode, hashType SigHashType,
+	tx *wire.MsgTx, idx int, cachedPrefix *chainhash.Hash) ([]byte, error) {
+	return calcSignatureHash(script, hashType, tx, idx, cachedPrefix)
+}
+
 // Internal tests for opcode parsing with bad data templates.
 func TestParseOpcode(t *testing.T) {
 	// Deep copy the array and make one of the opcodes invalid by setting it

txscript/script.go

@@ -281,20 +281,55 @@ func removeOpcodeByData(pkscript []parsedOpcode, data []byte) []parsedOpcode {
 		}
 	}
 	return retScript
-
 }
 
-// CalcSignatureHash is an exported version for testing.
-func CalcSignatureHash(script []parsedOpcode, hashType SigHashType,
-	tx *wire.MsgTx, idx int, cachedPrefix *chainhash.Hash) ([]byte, error) {
-	return calcSignatureHash(script, hashType, tx, idx, cachedPrefix)
+// optimizedWitnessSigningHash is a fast calculation of the witness for
+// signing to be used when signing with the SigHashAll optimization flag on.
+// It writes the serialized pkScript as it would be in the form of a serialized
+// msgTx with a witness signing version to a buffer, then hashes that.
+func optimizedWitnessSigningHash(script []parsedOpcode, hashType SigHashType,
+	tx *wire.MsgTx, idx int) (chainhash.Hash, error) {
+	var buf bytes.Buffer
+	var versionB [4]byte
+	if hashType&sigHashMask != SigHashAllValue {
+		binary.LittleEndian.PutUint32(versionB[:],
+			uint32(wire.WitnessSigningMsgTxVersion()))
+		buf.Write(versionB[:])
+	} else {
+		binary.LittleEndian.PutUint32(versionB[:],
+			uint32(wire.WitnessValueSigningMsgTxVersion()))
+		buf.Write(versionB[:])
+	}
+	txInLen := len(tx.TxIn)
+	wire.WriteVarInt(&buf, uint64(txInLen))
+	for i := 0; i < idx; i++ {
+		buf.WriteByte(0x00)
+	}
+	sigScript, _ := unparseScript(script)
+	err := wire.WriteVarBytes(&buf, sigScript)
+	if err != nil {
+		return chainhash.Hash{}, err
+	}
+	for i := idx + 1; i < txInLen; i++ {
+		buf.WriteByte(0x00)
+	}
+
+	return chainhash.HashFuncH(buf.Bytes()), nil
 }
 
 // calcSignatureHash will, given a script and hash type for the current script
 // engine instance, calculate the signature hash to be used for signing and
 // verification.
 func calcSignatureHash(script []parsedOpcode, hashType SigHashType,
 	tx *wire.MsgTx, idx int, cachedPrefix *chainhash.Hash) ([]byte, error) {
+	// Only use the optimization for SigHashAll with
+	// SigHashAnyoneCanPay disabled, which makes up
+	// the vast majority of outputs on the blockchain.
+	usesOptimization := cachedPrefix != nil &&
+		(hashType&sigHashMask == SigHashAll) &&
+		(hashType&SigHashAnyOneCanPay == 0) &&
+		chaincfg.SigHashOptimization
+
 	// The SigHashSingle signature type signs only the corresponding input
 	// and output (the output with the same index number as the input).
 	//
@@ -322,62 +357,86 @@ func calcSignatureHash(script []parsedOpcode, hashType SigHashType,
 	}
 
 	// Remove all instances of OP_CODESEPARATOR from the script.
+	// Decred: OP_CODESEPARATOR is disabled, remove?
 	script = removeOpcode(script, OP_CODESEPARATOR)
 
 	// Make a deep copy of the transaction, zeroing out the script for all
-	// inputs that are not currently being processed.
-	txCopy := tx.Copy()
-	for i := range txCopy.TxIn {
-		if i == idx {
-			// UnparseScript cannot fail here because removeOpcode
-			// above only returns a valid script.
-			sigScript, _ := unparseScript(script)
-			txCopy.TxIn[idx].SignatureScript = sigScript
-		} else {
-			txCopy.TxIn[i].SignatureScript = nil
+	// inputs that are not currently being processed. If the SigHashAll
+	// optimization is enabled, calculate the witness hash by quickly
+	// writing the version and pkScript to a buffer and then hashing it.
+	var txCopy *wire.MsgTx
+	var witnessHash chainhash.Hash
+	if usesOptimization {
+		var err error
+		witnessHash, err = optimizedWitnessSigningHash(script, hashType, tx, idx)
+		if err != nil {
+			return nil, err
 		}
-	}
-
-	switch hashType & sigHashMask {
-	case SigHashNone:
-		txCopy.TxOut = txCopy.TxOut[0:0] // Empty slice.
+	} else {
+		// No optimization available, or optimization disabled.
+		// Do the complete copying of the transaction.
+		txCopy = tx.Copy()
 		for i := range txCopy.TxIn {
-			if i != idx {
-				txCopy.TxIn[i].Sequence = 0
+			if i == idx {
+				// UnparseScript cannot fail here because removeOpcode
+				// above only returns a valid script.
+				sigScript, _ := unparseScript(script)
+				txCopy.TxIn[idx].SignatureScript = sigScript
+			} else {
+				txCopy.TxIn[i].SignatureScript = nil
 			}
 		}
 
-	case SigHashSingle:
-		// Resize output array to up to and including requested index.
-		txCopy.TxOut = txCopy.TxOut[:idx+1]
+		switch hashType & sigHashMask {
+		case SigHashNone:
+			txCopy.TxOut = txCopy.TxOut[0:0] // Empty slice.
+			for i := range txCopy.TxIn {
+				if i != idx {
+					txCopy.TxIn[i].Sequence = 0
+				}
+			}
 
-		// All but current output get zeroed out.
-		for i := 0; i < idx; i++ {
-			txCopy.TxOut[i].Value = -1
-			txCopy.TxOut[i].PkScript = nil
-		}
+		case SigHashSingle:
+			// Resize output array to up to and including requested index.
+			txCopy.TxOut = txCopy.TxOut[:idx+1]
 
-		// Sequence on all other inputs is 0, too.
-		for i := range txCopy.TxIn {
-			if i != idx {
-				txCopy.TxIn[i].Sequence = 0
+			// All but current output get zeroed out.
+			for i := 0; i < idx; i++ {
+				txCopy.TxOut[i].Value = -1
+				txCopy.TxOut[i].PkScript = nil
+			}
+
+			// Sequence on all other inputs is 0, too.
+			for i := range txCopy.TxIn {
+				if i != idx {
+					txCopy.TxIn[i].Sequence = 0
+				}
 			}
+
+		default:
+			// Consensus treats undefined hashtypes like normal SigHashAll
+			// for purposes of hash generation.
+			fallthrough
+		case SigHashOld:
+			fallthrough
+		case SigHashAllValue:
+			fallthrough
+		case SigHashAll:
+			// Nothing special here.
+		}
+		if hashType&SigHashAnyOneCanPay != 0 {
+			txCopy.TxIn = txCopy.TxIn[idx : idx+1]
+			idx = 0
 		}
 
-	default:
-		// Consensus treats undefined hashtypes like normal SigHashAll
-		// for purposes of hash generation.
-		fallthrough
-	case SigHashOld:
-		fallthrough
-	case SigHashAllValue:
-		fallthrough
-	case SigHashAll:
-		// Nothing special here.
-	}
-	if hashType&SigHashAnyOneCanPay != 0 {
-		txCopy.TxIn = txCopy.TxIn[idx : idx+1]
-		idx = 0
+		// If the ValueIn is to be included in what we're signing, sign
+		// the witness hash that includes it. Otherwise, just sign the
+		// prefix and signature scripts.
+		if hashType&sigHashMask != SigHashAllValue {
+			witnessHash = txCopy.TxShaWitnessSigning()
+		} else {
+			witnessHash = txCopy.TxShaWitnessValueSigning()
+		}
 	}
 
 	// The final hash (message to sign) is the hash of:
@@ -387,31 +446,20 @@ func calcSignatureHash(script []parsedOpcode, hashType SigHashType,
 	var wbuf bytes.Buffer
 	binary.Write(&wbuf, binary.LittleEndian, uint32(hashType))
 
-	// Optimization for SIGHASH_ALL. In this case, the prefix hash is
-	// the same as the transaction hash because only the inputs have
-	// been modified, so don't bother to do the wasteful O(N^2) extra
-	// hash here.
+	// Optimization for SIGHASH_ALL. This almost fixes the O(N^2)
+	// behaviour seen in Bitcoin, except N many 0x00 bytes are
+	// written during the serialization of the witness. TODO:
+	// Add a softfork flag that changes writing N many 0x00 bytes
+	// to simply writing the index.
 	// The caching only works if the "anyone can pay flag" is also
 	// disabled.
 	var prefixHash chainhash.Hash
-	if cachedPrefix != nil &&
-		(hashType&sigHashMask == SigHashAll) &&
-		(hashType&SigHashAnyOneCanPay == 0) &&
-		chaincfg.SigHashOptimization {
+	if usesOptimization {
 		prefixHash = *cachedPrefix
 	} else {
 		prefixHash = txCopy.TxSha()
 	}
 
-	// If the ValueIn is to be included in what we're signing, sign
-	// the witness hash that includes it. Otherwise, just sign the
-	// prefix and signature scripts.
-	var witnessHash chainhash.Hash
-	if hashType&sigHashMask != SigHashAllValue {
-		witnessHash = txCopy.TxShaWitnessSigning()
-	} else {
-		witnessHash = txCopy.TxShaWitnessValueSigning()
-	}
 	wbuf.Write(prefixHash.Bytes())
 	wbuf.Write(witnessHash.Bytes())
 	return chainhash.HashFuncB(wbuf.Bytes()), nil

txscript/script_test.go

@@ -532,13 +532,14 @@ func TestCalcSignatureHash(t *testing.T) {
 	pops, _ := txscript.TstParseScript([]byte{0x01, 0x01, 0x02, 0x03})
 
 	// Test prefix caching.
-	msg1, err := txscript.CalcSignatureHash(pops, txscript.SigHashAll, tx, 0, nil)
+	msg1, err := txscript.TstCalcSignatureHash(pops, txscript.SigHashAll, tx, 0,
+		nil)
 	if err != nil {
 		t.Fatalf("unexpected error %v", err.Error())
 	}
 
 	prefixHash := tx.TxSha()
-	msg2, err := txscript.CalcSignatureHash(pops, txscript.SigHashAll, tx, 0,
+	msg2, err := txscript.TstCalcSignatureHash(pops, txscript.SigHashAll, tx, 0,
 		&prefixHash)
 	if err != nil {
 		t.Fatalf("unexpected error %v", err.Error())
@@ -563,7 +564,7 @@ func TestCalcSignatureHash(t *testing.T) {
 
 	// Move the index and make sure that we get a whole new hash, despite
 	// using the same TxOuts.
-	msg3, err := txscript.CalcSignatureHash(pops, txscript.SigHashAll, tx, 1,
+	msg3, err := txscript.TstCalcSignatureHash(pops, txscript.SigHashAll, tx, 1,
 		&prefixHash)
 	if err != nil {
 		t.Fatalf("unexpected error %v", err.Error())

wire/common.go

@@ -6,6 +6,7 @@
 package wire
 
 import (
+	"bytes"
 	"crypto/rand"
 	"encoding/binary"
 	"fmt"
@@ -434,6 +435,13 @@ func writeVarInt(w io.Writer, pver uint32, val uint64) error {
 	return err
 }
 
+// WriteVarInt is the exported form of writeVarBytes. It uses the latest
+// version of the wire protocol as stored in the package.
+func WriteVarInt(buf *bytes.Buffer, val uint64) error {
+	w := io.Writer(buf)
+	return writeVarInt(w, ProtocolVersion, val)
+}
+
 // VarIntSerializeSize returns the number of bytes it would take to serialize
 // val as a variable length integer.
 func VarIntSerializeSize(val uint64) int {
@@ -533,7 +541,7 @@ func readVarBytes(r io.Reader, pver uint32, maxAllowed uint32,
 	return b, nil
 }
 
-// writeVarInt serializes a variable length byte array to w as a varInt
+// writeVarBytes serializes a variable length byte array to w as a varInt
 // containing the number of bytes, followed by the bytes themselves.
 func writeVarBytes(w io.Writer, pver uint32, bytes []byte) error {
 	slen := uint64(len(bytes))
@@ -549,6 +557,13 @@ func writeVarBytes(w io.Writer, pver uint32, bytes []byte) error {
 	return nil
 }
 
+// WriteVarBytes is the exported form of writeVarBytes. It uses the latest
+// version of the wire protocol as stored in the package.
+func WriteVarBytes(buf *bytes.Buffer, bytes []byte) error {
+	w := io.Writer(buf)
+	return writeVarBytes(w, ProtocolVersion, bytes)
+}
+
 // randomUint64 returns a cryptographically random uint64 value.  This
 // unexported version takes a reader primarily to ensure the error paths
 // can be properly tested by passing a fake reader in the tests.