From 6702a2bab815491531dc07aa02ddd57591cf2045 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Fri, 19 Apr 2019 19:39:28 -0700
Subject: [PATCH] txscript: Return witness version and program in one pass

---
 txscript/standard.go | 52 +++++++++++++++++++++++++++-----------------
 1 file changed, 32 insertions(+), 20 deletions(-)

diff --git a/txscript/standard.go b/txscript/standard.go
index f0d4bfb6fd..62b8118792 100644
--- a/txscript/standard.go
+++ b/txscript/standard.go
@@ -379,25 +379,13 @@ func isWitnessScriptHashScript(script []byte) bool {
 	return extractWitnessScriptHash(script) != nil
 }
 
-// isWitnessProgramScript returns true if the passed script is a witness
-// program, and false otherwise. A witness program MUST adhere to the following
-// constraints: there must be exactly two pops (program version and the program
-// itself), the first opcode MUST be a small integer (0-16), the push data MUST
-// be canonical, and finally the size of the push data must be between 2 and 40
-// bytes.
-//
-// The length of the script must be between 4 and 42 bytes. The
-// smallest program is the witness version, followed by a data push of
-// 2 bytes.  The largest allowed witness program has a data push of
-// 40-bytes.
-//
-// NOTE: This function is only valid for version 0 scripts.  Since the function
-// does not accept a script version, the results are undefined for other script
-// versions.
-func isWitnessProgramScript(script []byte) bool {
+// extractWitnessProgramInfo returns the version and program if the passed
+// script constitutes a valid witness program. The alst return value indicates
+// whether or not the script is a valid witness program.
+func extractWitnessProgramInfo(script []byte) (int, []byte, bool) {
 	// Skip parsing if we know the program is invalid based on size.
 	if len(script) < 4 || len(script) > 42 {
-		return false
+		return 0, nil, false
 	}
 
 	const scriptVersion = 0
@@ -407,8 +395,9 @@ func isWitnessProgramScript(script []byte) bool {
 	if !tokenizer.Next() ||
 		!isSmallInt(tokenizer.Opcode()) {
 
-		return false
+		return 0, nil, false
 	}
+	version := asSmallInt(tokenizer.Opcode())
 
 	// The second opcode must be a canonical data push, the length of the
 	// data push is bounded to 40 by the initial check on overall script
@@ -416,12 +405,35 @@ func isWitnessProgramScript(script []byte) bool {
 	if !tokenizer.Next() ||
 		!isCanonicalPush(tokenizer.Opcode(), tokenizer.Data()) {
 
-		return false
+		return 0, nil, false
 	}
+	program := tokenizer.Data()
 
 	// The witness program is valid if there are no more opcodes, and we
 	// terminated without a parsing error.
-	return tokenizer.Done() && tokenizer.Err() == nil
+	valid := tokenizer.Done() && tokenizer.Err() == nil
+
+	return version, program, valid
+}
+
+// isWitnessProgramScript returns true if the passed script is a witness
+// program, and false otherwise. A witness program MUST adhere to the following
+// constraints: there must be exactly two pops (program version and the program
+// itself), the first opcode MUST be a small integer (0-16), the push data MUST
+// be canonical, and finally the size of the push data must be between 2 and 40
+// bytes.
+//
+// The length of the script must be between 4 and 42 bytes. The
+// smallest program is the witness version, followed by a data push of
+// 2 bytes.  The largest allowed witness program has a data push of
+// 40-bytes.
+//
+// NOTE: This function is only valid for version 0 scripts.  Since the function
+// does not accept a script version, the results are undefined for other script
+// versions.
+func isWitnessProgramScript(script []byte) bool {
+	_, _, valid := extractWitnessProgramInfo(script)
+	return valid
 }
 
 // isNullDataScript returns whether or not the passed script is a standard