wire/msgblock+msgtx: user block-level script slab

Roasbeef · Dec 16, 2023 · 8c4da83 · 8c4da83
1 parent d7396dc
commit 8c4da83
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 19 deletions.
diff --git a/wire/msgblock.go b/wire/msgblock.go
@@ -86,17 +86,19 @@ func (msg *MsgBlock) BtcDecode(r io.Reader, pver uint32, enc MessageEncoding) er
 		return messageError("MsgBlock.BtcDecode", str)
 	}
 
+	scriptBuf := scriptPool.Borrow()
 	msg.Transactions = make([]*MsgTx, 0, txCount)
 	for i := uint64(0); i < txCount; i++ {
 		tx := MsgTx{}
-		err := tx.btcDecode(r, pver, enc, buf)
+		err := tx.btcDecode(r, pver, enc, buf, scriptBuf[:])
 		if err != nil {
+			scriptPool.Return(scriptBuf)
 			binarySerializer.Return(buf)
 			return err
 		}
 		msg.Transactions = append(msg.Transactions, &tx)
 	}
-
+	scriptPool.Return(scriptBuf)
 	binarySerializer.Return(buf)
 
 	return nil
@@ -164,22 +166,25 @@ func (msg *MsgBlock) DeserializeTxLoc(r *bytes.Buffer) ([]TxLoc, error) {
 		return nil, messageError("MsgBlock.DeserializeTxLoc", str)
 	}
 
+	scriptBuf := scriptPool.Borrow()
+
 	// Deserialize each transaction while keeping track of its location
 	// within the byte stream.
 	msg.Transactions = make([]*MsgTx, 0, txCount)
 	txLocs := make([]TxLoc, txCount)
 	for i := uint64(0); i < txCount; i++ {
 		txLocs[i].TxStart = fullLen - r.Len()
 		tx := MsgTx{}
-		err := tx.btcDecode(r, 0, WitnessEncoding, buf)
+		err := tx.btcDecode(r, 0, WitnessEncoding, buf, scriptBuf[:])
 		if err != nil {
+			scriptPool.Return(scriptBuf)
 			binarySerializer.Return(buf)
 			return nil, err
 		}
 		msg.Transactions = append(msg.Transactions, &tx)
 		txLocs[i].TxLen = (fullLen - r.Len()) - txLocs[i].TxStart
 	}
-
+	scriptPool.Return(scriptBuf)
 	binarySerializer.Return(buf)
 
 	return txLocs, nil

diff --git a/wire/msgtx.go b/wire/msgtx.go
@@ -456,13 +456,15 @@ func (msg *MsgTx) Copy() *MsgTx {
 // database, as opposed to decoding transactions from the wire.
 func (msg *MsgTx) BtcDecode(r io.Reader, pver uint32, enc MessageEncoding) error {
 	buf := binarySerializer.Borrow()
-	err := msg.btcDecode(r, pver, enc, buf)
+	sbuf := scriptPool.Borrow()
+	err := msg.btcDecode(r, pver, enc, buf, sbuf[:])
+	scriptPool.Return(sbuf)
 	binarySerializer.Return(buf)
 	return err
 }
 
 func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
-	buf []byte) error {
+	buf, sbuf []byte) error {
 
 	if _, err := io.ReadFull(r, buf[:4]); err != nil {
 		return err
@@ -509,9 +511,6 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 		return messageError("MsgTx.BtcDecode", str)
 	}
 
-	scriptBuf := scriptPool.Borrow()
-	sbuf := scriptBuf[:]
-
 	// Deserialize the inputs.
 	var totalScriptSize uint64
 	txIns := make([]TxIn, count)
@@ -523,7 +522,6 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 		msg.TxIn[i] = ti
 		err = readTxInBuf(r, pver, msg.Version, ti, buf, sbuf)
 		if err != nil {
-			scriptPool.Return(scriptBuf)
 			return err
 		}
 		totalScriptSize += uint64(len(ti.SignatureScript))
@@ -532,15 +530,13 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 
 	count, err = ReadVarIntBuf(r, pver, buf)
 	if err != nil {
-		scriptPool.Return(scriptBuf)
 		return err
 	}
 
 	// Prevent more output transactions than could possibly fit into a
 	// message.  It would be possible to cause memory exhaustion and panics
 	// without a sane upper bound on this count.
 	if count > uint64(maxTxOutPerMessage) {
-		scriptPool.Return(scriptBuf)
 		str := fmt.Sprintf("too many output transactions to fit into "+
 			"max message size [count %d, max %d]", count,
 			maxTxOutPerMessage)
@@ -557,7 +553,6 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 		msg.TxOut[i] = to
 		err = readTxOutBuf(r, pver, msg.Version, to, buf, sbuf)
 		if err != nil {
-			scriptPool.Return(scriptBuf)
 			return err
 		}
 		totalScriptSize += uint64(len(to.PkScript))
@@ -573,14 +568,12 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 			// varint which encodes the number of stack items.
 			witCount, err := ReadVarIntBuf(r, pver, buf)
 			if err != nil {
-				scriptPool.Return(scriptBuf)
 				return err
 			}
 
 			// Prevent a possible memory exhaustion attack by
 			// limiting the witCount value to a sane upper bound.
 			if witCount > maxWitnessItemsPerInput {
-				scriptPool.Return(scriptBuf)
 				str := fmt.Sprintf("too many witness items to fit "+
 					"into max message size [count %d, max %d]",
 					witCount, maxWitnessItemsPerInput)
@@ -597,7 +590,6 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 					"script witness item",
 				)
 				if err != nil {
-					scriptPool.Return(scriptBuf)
 					return err
 				}
 				totalScriptSize += uint64(len(txin.Witness[j]))
@@ -607,7 +599,6 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 	}
 
 	if _, err := io.ReadFull(r, buf[:4]); err != nil {
-		scriptPool.Return(scriptBuf)
 		return err
 	}
 	msg.LockTime = littleEndian.Uint32(buf[:4])
@@ -670,8 +661,6 @@ func (msg *MsgTx) btcDecode(r io.Reader, pver uint32, enc MessageEncoding,
 		offset += scriptSize
 	}
 
-	scriptPool.Return(scriptBuf)
-
 	return nil
 }