refactor partial signature verification

src/main/java/com/craigraw/drongo/TransactionTask.java

@@ -44,9 +44,13 @@ public class TransactionTask implements Runnable {
 
                 TransactionOutput referencedOutput = referencedTransaction.getOutputs().get((int)referencedVout);
                 if(referencedOutput.getScript().containsToAddress()) {
-                    Address[] inputAddresses = referencedOutput.getScript().getToAddresses();
+                    try {
-                    input.getOutpoint().setAddresses(inputAddresses);
+                        Address[] inputAddresses = referencedOutput.getScript().getToAddresses();
-                    inputJoiner.add((inputAddresses.length == 1 ? inputAddresses[0] : Arrays.asList(inputAddresses)) + ":" + vin);
+                        input.getOutpoint().setAddresses(inputAddresses);
+                        inputJoiner.add((inputAddresses.length == 1 ? inputAddresses[0] : Arrays.asList(inputAddresses)) + ":" + vin);
+                    } catch(NonStandardScriptException e) {
+                        //Cannot happen
+                    }
                 } else {
                     log.warn("Could not determine nature of referenced input tx: " + referencedTxID + ":" + referencedVout);
                 }
@@ -62,9 +66,13 @@ public class TransactionTask implements Runnable {
         for(TransactionOutput output : transaction.getOutputs()) {
             try {
                 if(output.getScript().containsToAddress()) {
-                    Address[] outputAddresses = output.getScript().getToAddresses();
+                    try {
-                    output.setAddresses(outputAddresses);
+                        Address[] outputAddresses = output.getScript().getToAddresses();
-                    outputJoiner.add((outputAddresses.length == 1 ? outputAddresses[0] : Arrays.asList(outputAddresses)) + ":" + vout + " (" + output.getValue() + ")");
+                        output.setAddresses(outputAddresses);
+                        outputJoiner.add((outputAddresses.length == 1 ? outputAddresses[0] : Arrays.asList(outputAddresses)) + ":" + vout + " (" + output.getValue() + ")");
+                    } catch(NonStandardScriptException e) {
+                        //Cannot happen
+                    }
                 }
             } catch(ProtocolException e) {
                 log.debug("Invalid script for output " + vout + " detected (" + e.getMessage() + "). Skipping...");

src/main/java/com/craigraw/drongo/protocol/NonStandardScriptException.java

@@ -0,0 +1,19 @@
+package com.craigraw.drongo.protocol;
+
+public class NonStandardScriptException extends Exception {
+    public NonStandardScriptException() {
+        super();
+    }
+
+    public NonStandardScriptException(String message) {
+        super(message);
+    }
+
+    public NonStandardScriptException(Throwable cause) {
+        super(cause);
+    }
+
+    public NonStandardScriptException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

src/main/java/com/craigraw/drongo/protocol/Script.java

@@ -132,7 +132,7 @@ public class Script {
     /**
      * Gets the destination address from this script, if it's in the required form.
      */
-    public Address[] getToAddresses() {
+    public Address[] getToAddresses() throws NonStandardScriptException {
         if (ScriptPattern.isP2PK(this))
             return new Address[] { new P2PKAddress( ScriptPattern.extractPKFromP2PK(this)) };
         else if (ScriptPattern.isP2PKH(this))
@@ -146,7 +146,19 @@ public class Script {
         else if (ScriptPattern.isSentToMultisig(this))
             return ScriptPattern.extractMultisigAddresses(this);
         else
-            throw new ProtocolException("Cannot cast this script to an address");
+            throw new NonStandardScriptException("Cannot find addresses in non standard script: " + toString());
+    }
+
+    public int getNumRequiredSignatures() throws NonStandardScriptException {
+        if(ScriptPattern.isP2PK(this) || ScriptPattern.isP2PKH(this) || ScriptPattern.isP2WPKH(this)) {
+            return 1;
+        }
+
+        if(ScriptPattern.isSentToMultisig(this)) {
+            return ScriptPattern.extractMultisigThreshold(this);
+        }
+
+        throw new NonStandardScriptException("Cannot find number of required signatures for non standard script: " + toString());
     }
 
     public static int decodeFromOpN(int opcode) {

src/main/java/com/craigraw/drongo/protocol/ScriptPattern.java

@@ -134,6 +134,10 @@ public class ScriptPattern {
         return true;
     }
 
+    public static int extractMultisigThreshold(Script script) {
+        return decodeFromOpN(script.chunks.get(0).opcode);
+    }
+
     public static Address[] extractMultisigAddresses(Script script) {
         List<Address> addresses = new ArrayList<>();
 

src/main/java/com/craigraw/drongo/protocol/Transaction.java

@@ -501,7 +501,7 @@ public class Transaction extends TransactionPart {
         }
     }
 
-    public static final void main(String[] args) {
+    public static final void main(String[] args) throws NonStandardScriptException {
         String hex = "0100000002fe3dc9208094f3ffd12645477b3dc56f60ec4fa8e6f5d67c565d1c6b9216b36e0000000000ffffffff0815cf020f013ed6cf91d29f4202e8a58726b1ac6c79da47c23d1bee0a6925f80000000000ffffffff0100f2052a010000001976a914a30741f8145e5acadf23f751864167f32e0963f788ac00000000";
         byte[] transactionBytes = Utils.hexToBytes(hex);
         Transaction transaction = new Transaction(transactionBytes);

src/main/java/com/craigraw/drongo/psbt/PSBT.java

@@ -3,9 +3,6 @@ package com.craigraw.drongo.psbt;
 import com.craigraw.drongo.ExtendedPublicKey;
 import com.craigraw.drongo.KeyDerivation;
 import com.craigraw.drongo.Utils;
-import com.craigraw.drongo.address.Address;
-import com.craigraw.drongo.address.P2PKHAddress;
-import com.craigraw.drongo.crypto.ECKey;
 import com.craigraw.drongo.protocol.*;
 import org.bouncycastle.util.encoders.Base64;
 import org.bouncycastle.util.encoders.Hex;
@@ -214,7 +211,7 @@ public class PSBT {
                     for(TransactionOutput output: transaction.getOutputs()) {
                         try {
                             log.debug(" Transaction output value: " + output.getValue() + " to addresses " + Arrays.asList(output.getScript().getToAddresses()) + " with script hex " + Hex.toHexString(output.getScript().getProgram()) + " to script " + output.getScript());
-                        } catch(ProtocolException e) {
+                        } catch(NonStandardScriptException e) {
                             log.debug(" Transaction output value: " + output.getValue() + " with script hex " + Hex.toHexString(output.getScript().getProgram()) + " to script " + output.getScript());
                         }
                     }
@@ -253,7 +250,7 @@ public class PSBT {
             int inputIndex = this.psbtInputs.size();
             PSBTInput input = new PSBTInput(inputEntries, transaction, inputIndex);
 
-            boolean verified = verifySignatures(input, inputIndex);
+            boolean verified = input.verifySignatures();
             if(verified) {
                 log.debug("Verified signatures on input #" + inputIndex);
             }
@@ -262,52 +259,6 @@ public class PSBT {
         }
     }
 
-    private boolean verifySignatures(PSBTInput input, int index) {
-        if(input.getSigHash() != null && (input.getNonWitnessUtxo() != null || input.getWitnessUtxo() != null)) {
-            int vout = (int)transaction.getInputs().get(index).getOutpoint().getIndex();
-            Script inputScript = input.getNonWitnessUtxo() != null ? input.getNonWitnessUtxo().getOutputs().get(vout).getScript() : input.getWitnessUtxo().getScript();
-
-            Script connectedScript = inputScript;
-            if(ScriptPattern.isP2SH(connectedScript)) {
-                if(input.getRedeemScript() == null) {
-                    return false;
-                } else {
-                    connectedScript = input.getRedeemScript();
-                }
-            }
-
-            if(ScriptPattern.isP2WPKH(connectedScript)) {
-                Address address = new P2PKHAddress(connectedScript.getPubKeyHash());
-                connectedScript = address.getOutputScript();
-            } else if(ScriptPattern.isP2WSH(connectedScript)) {
-                if(input.getWitnessScript() == null) {
-                    return false;
-                } else {
-                    connectedScript = input.getWitnessScript();
-                }
-            }
-
-            Sha256Hash hash = null;
-            if(input.getNonWitnessUtxo() != null) {
-                hash = transaction.hashForSignature(index, connectedScript, input.getSigHash(), false);
-            } else {
-                long prevValue = input.getWitnessUtxo().getValue();
-                hash = transaction.hashForWitnessSignature(index, connectedScript, prevValue, input.getSigHash(), false);
-            }
-
-            for(ECKey sigPublicKey : input.getPartialSignatures().keySet()) {
-                TransactionSignature signature = input.getPartialSignature(sigPublicKey);
-                if(!sigPublicKey.verify(hash, signature)) {
-                    throw new IllegalStateException("Partial signature does not verify against provided public key");
-                }
-            }
-
-            return true;
-        }
-
-        return false;
-    }
-
     private void parseOutputEntries(List<List<PSBTEntry>> outputEntryLists) {
         for(List<PSBTEntry> outputEntries : outputEntryLists) {
             PSBTEntry duplicate = findDuplicateKey(outputEntries);

src/main/java/com/craigraw/drongo/psbt/PSBTInput.java

@@ -2,6 +2,8 @@ package com.craigraw.drongo.psbt;
 
 import com.craigraw.drongo.KeyDerivation;
 import com.craigraw.drongo.Utils;
+import com.craigraw.drongo.address.Address;
+import com.craigraw.drongo.address.P2PKHAddress;
 import com.craigraw.drongo.crypto.ECKey;
 import com.craigraw.drongo.crypto.LazyECPoint;
 import com.craigraw.drongo.protocol.*;
@@ -42,6 +44,9 @@ public class PSBTInput {
     private String porCommitment;
     private Map<String, String> proprietary = new LinkedHashMap<>();
 
+    private Transaction transaction;
+    private int index;
+
     private static final Logger log = LoggerFactory.getLogger(PSBTInput.class);
 
     PSBTInput(List<PSBTEntry> inputEntries, Transaction transaction, int index) {
@@ -66,7 +71,11 @@ public class PSBTInput {
                         log.debug(" Transaction input references txid: " + input.getOutpoint().getHash() + " vout " + input.getOutpoint().getIndex() + " with script " + input.getScriptSig());
                     }
                     for(TransactionOutput output: nonWitnessTx.getOutputs()) {
-                        log.debug(" Transaction output value: " + output.getValue() + " to addresses " + Arrays.asList(output.getScript().getToAddresses()) + " with script hex " + Hex.toHexString(output.getScript().getProgram()) + " to script " + output.getScript());
+                        try {
+                            log.debug(" Transaction output value: " + output.getValue() + " to addresses " + Arrays.asList(output.getScript().getToAddresses()) + " with script hex " + Hex.toHexString(output.getScript().getProgram()) + " to script " + output.getScript());
+                        } catch(NonStandardScriptException e) {
+                            log.error("Unknown script type", e);
+                        }
                     }
                     break;
                 case PSBT_IN_WITNESS_UTXO:
@@ -79,7 +88,11 @@ public class PSBTInput {
                         throw new IllegalStateException("Witness UTXO provided for non-witness or unknown input");
                     }
                     this.witnessUtxo = witnessTxOutput;
-                    log.debug("Found input witness utxo amount " + witnessTxOutput.getValue() + " script hex " + Hex.toHexString(witnessTxOutput.getScript().getProgram()) + " script " + witnessTxOutput.getScript() + " addresses " + Arrays.asList(witnessTxOutput.getScript().getToAddresses()));
+                    try {
+                        log.debug("Found input witness utxo amount " + witnessTxOutput.getValue() + " script hex " + Hex.toHexString(witnessTxOutput.getScript().getProgram()) + " script " + witnessTxOutput.getScript() + " addresses " + Arrays.asList(witnessTxOutput.getScript().getToAddresses()));
+                    } catch(NonStandardScriptException e) {
+                        log.error("Unknown script type", e);
+                    }
                     break;
                 case PSBT_IN_PARTIAL_SIG:
                     entry.checkOneBytePlusPubKey();
@@ -168,6 +181,9 @@ public class PSBTInput {
                     log.warn("PSBT input not recognized key type: " + entry.getKeyType());
             }
         }
+
+        this.transaction = transaction;
+        this.index = index;
     }
 
     public Transaction getNonWitnessUtxo() {
@@ -221,4 +237,75 @@ public class PSBTInput {
     public Map<String, String> getProprietary() {
         return proprietary;
     }
+
+    public boolean isSigned() throws NonStandardScriptException {
+        //All partial sigs are already verified
+        int reqSigs = getConnectedScript().getNumRequiredSignatures();
+        int sigs = getPartialSignatures().size();
+        return sigs == reqSigs;
+    }
+
+    boolean verifySignatures() {
+        Transaction.SigHash localSigHash = getSigHash();
+        if(localSigHash == null) {
+            //Assume SigHash.ALL
+            localSigHash = Transaction.SigHash.ALL;
+        }
+
+        if(getNonWitnessUtxo() != null || getWitnessUtxo() != null) {
+            Script connectedScript = getConnectedScript();
+            if(connectedScript != null) {
+                Sha256Hash hash = getHashForSignature(connectedScript, localSigHash);
+
+                for(ECKey sigPublicKey : getPartialSignatures().keySet()) {
+                    TransactionSignature signature = getPartialSignature(sigPublicKey);
+                    if(!sigPublicKey.verify(hash, signature)) {
+                        throw new IllegalStateException("Partial signature does not verify against provided public key");
+                    }
+                }
+
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    public Script getConnectedScript() {
+        int vout = (int)transaction.getInputs().get(index).getOutpoint().getIndex();
+        Script connectedScript = getNonWitnessUtxo() != null ? getNonWitnessUtxo().getOutputs().get(vout).getScript() : getWitnessUtxo().getScript();
+
+        if(ScriptPattern.isP2SH(connectedScript)) {
+            if(getRedeemScript() == null) {
+                return null;
+            } else {
+                connectedScript = getRedeemScript();
+            }
+        }
+
+        if(ScriptPattern.isP2WPKH(connectedScript)) {
+            Address address = new P2PKHAddress(connectedScript.getPubKeyHash());
+            connectedScript = address.getOutputScript();
+        } else if(ScriptPattern.isP2WSH(connectedScript)) {
+            if(getWitnessScript() == null) {
+                return null;
+            } else {
+                connectedScript = getWitnessScript();
+            }
+        }
+
+        return connectedScript;
+    }
+
+    private Sha256Hash getHashForSignature(Script connectedScript, Transaction.SigHash localSigHash) {
+        Sha256Hash hash;
+        if (getNonWitnessUtxo() != null) {
+            hash = transaction.hashForSignature(index, connectedScript, localSigHash, false);
+        } else {
+            long prevValue = getWitnessUtxo().getValue();
+            hash = transaction.hashForWitnessSignature(index, connectedScript, prevValue, localSigHash, false);
+        }
+
+        return hash;
+    }
 }