combiner, finaliser and extractor for psbts

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

@@ -258,7 +258,7 @@ public class Transaction extends ChildMessage {
         int numWitnesses = inputs.size();
         for (int i = 0; i < numWitnesses; i++) {
             TransactionWitness witness = new TransactionWitness(this, payload, cursor);
-            inputs.get(i).setWitness(witness);
+            inputs.get(i).witness(witness);
             cursor += witness.getLength();
         }
     }

src/main/java/com/sparrowwallet/drongo/protocol/TransactionInput.java

@@ -78,7 +78,7 @@ public class TransactionInput extends ChildMessage {
         return scriptSig;
     }
 
-    void setScriptBytes(byte[] scriptBytes) {
+    public void setScriptBytes(byte[] scriptBytes) {
         super.payload = null;
         this.scriptSig = null;
         int oldLength = length;
@@ -96,19 +96,22 @@ public class TransactionInput extends ChildMessage {
         return witness;
     }
 
-    void setWitness(TransactionWitness witness) {
+    void witness(TransactionWitness witness) {
+        this.witness = witness;
+    }
+
+    public void setWitness(TransactionWitness witness) {
+        int existingLength = getWitness() != null ? getWitness().getLength() : 0;
+        if(getParent() != null) {
+            getParent().adjustLength(witness.getLength() - existingLength);
+        }
+
         this.witness = witness;
     }
 
     public void clearWitness() {
-        TransactionWitness witness = getWitness();
-        if(witness != null) {
-            if(getParent() != null) {
-                getParent().adjustLength(-witness.getLength());
-            }
         setWitness(null);
     }
-    }
 
     public boolean hasWitness() {
         return witness != null && witness.getPushCount() != 0;

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

@@ -314,13 +314,11 @@ public class PSBT {
     public Long getFee() {
         long fee = 0L;
 
-        for (int i = 0; i < psbtInputs.size(); i++) {
+        for(PSBTInput input : psbtInputs) {
-            PSBTInput input = psbtInputs.get(i);
+            TransactionOutput utxo = input.getUtxo();
-            if(input.getNonWitnessUtxo() != null) {
+
-                int index = (int)transaction.getInputs().get(i).getOutpoint().getIndex();
+            if(utxo != null) {
-                fee += input.getNonWitnessUtxo().getOutputs().get(index).getValue();
+                fee += utxo.getValue();
-            } else if(input.getWitnessUtxo() != null) {
-                fee += input.getWitnessUtxo().getValue();
             } else {
                 log.error("Cannot determine fee - not enough information provided on inputs");
                 return null;
@@ -410,6 +408,63 @@ public class PSBT {
         return baos.toByteArray();
     }
 
+    public void combine(PSBT... psbts) {
+        for(PSBT psbt : psbts) {
+            combine(psbt);
+        }
+    }
+
+    public void combine(PSBT psbt) {
+        byte[] txBytes = transaction.bitcoinSerialize();
+        byte[] psbtTxBytes = psbt.getTransaction().bitcoinSerialize();
+
+        if(!Arrays.equals(txBytes, psbtTxBytes)) {
+            throw new IllegalArgumentException("Provided PSBT does contain a matching global transaction");
+        }
+
+        for(PSBTInput psbtInput : psbt.getPsbtInputs()) {
+            if(psbtInput.getFinalScriptSig() != null || psbtInput.getFinalScriptWitness() != null) {
+                throw new IllegalArgumentException("Cannot combine an already finalised PSBT");
+            }
+        }
+
+        if(psbt.getVersion() != null) {
+            version = psbt.getVersion();
+        }
+
+        extendedPublicKeys.putAll(psbt.extendedPublicKeys);
+        globalProprietary.putAll(psbt.globalProprietary);
+
+        for(int i = 0; i < getPsbtInputs().size(); i++) {
+            PSBTInput thisInput = getPsbtInputs().get(i);
+            PSBTInput otherInput = psbt.getPsbtInputs().get(i);
+            thisInput.combine(otherInput);
+        }
+
+        for(int i = 0; i < getPsbtOutputs().size(); i++) {
+            PSBTOutput thisOutput = getPsbtOutputs().get(i);
+            PSBTOutput otherOutput = psbt.getPsbtOutputs().get(i);
+            thisOutput.combine(otherOutput);
+        }
+    }
+
+    public Transaction extractTransaction() {
+        for(PSBTInput psbtInput : getPsbtInputs()) {
+            if(psbtInput.getFinalScriptSig() == null) {
+                return null;
+            }
+        }
+
+        for(int i = 0; i < transaction.getInputs().size(); i++) {
+            TransactionInput txInput = transaction.getInputs().get(i);
+            PSBTInput psbtInput = getPsbtInputs().get(i);
+            txInput.setScriptBytes(psbtInput.getFinalScriptSig().getProgram());
+            txInput.setWitness(psbtInput.getFinalScriptWitness());
+        }
+
+        return transaction;
+    }
+
     public List<PSBTInput> getPsbtInputs() {
         return psbtInputs;
     }

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

@@ -250,6 +250,38 @@ public class PSBTInput {
         return entries;
     }
 
+    void combine(PSBTInput psbtInput) {
+        if(psbtInput.nonWitnessUtxo != null) {
+            nonWitnessUtxo = psbtInput.nonWitnessUtxo;
+        }
+
+        if(psbtInput.witnessUtxo != null) {
+            witnessUtxo = psbtInput.witnessUtxo;
+        }
+
+        partialSignatures.putAll(psbtInput.partialSignatures);
+
+        if(psbtInput.sigHash != null) {
+            sigHash = psbtInput.sigHash;
+        }
+
+        if(psbtInput.redeemScript != null) {
+            redeemScript = psbtInput.redeemScript;
+        }
+
+        if(psbtInput.witnessScript != null) {
+            witnessScript = psbtInput.witnessScript;
+        }
+
+        derivedPublicKeys.putAll(psbtInput.derivedPublicKeys);
+
+        if(psbtInput.porCommitment != null) {
+            porCommitment = psbtInput.porCommitment;
+        }
+
+        proprietary.putAll(psbtInput.proprietary);
+    }
+
     public Transaction getNonWitnessUtxo() {
         return nonWitnessUtxo;
     }
@@ -286,10 +318,18 @@ public class PSBTInput {
         return finalScriptSig;
     }
 
+    public void setFinalScriptSig(Script finalScriptSig) {
+        this.finalScriptSig = finalScriptSig;
+    }
+
     public TransactionWitness getFinalScriptWitness() {
         return finalScriptWitness;
     }
 
+    public void setFinalScriptWitness(TransactionWitness finalScriptWitness) {
+        this.finalScriptWitness = finalScriptWitness;
+    }
+
     public String getPorCommitment() {
         return porCommitment;
     }
@@ -384,8 +424,7 @@ public class PSBTInput {
     }
 
     public Script getSigningScript() {
-        int vout = (int)transaction.getInputs().get(index).getOutpoint().getIndex();
+        Script signingScript = getUtxo().getScript();
-        Script signingScript = getWitnessUtxo() != null ? getWitnessUtxo().getScript() : getNonWitnessUtxo().getOutputs().get(vout).getScript();
 
         if(P2SH.isScriptType(signingScript)) {
             if(getRedeemScript() != null) {
@@ -412,6 +451,21 @@ public class PSBTInput {
         return signingScript;
     }
 
+    public TransactionOutput getUtxo() {
+        int vout = (int)transaction.getInputs().get(index).getOutpoint().getIndex();
+        return getWitnessUtxo() != null ? getWitnessUtxo() : (getNonWitnessUtxo() != null ?  getNonWitnessUtxo().getOutputs().get(vout) : null);
+    }
+
+    public void clearFinalised() {
+        partialSignatures.clear();
+        sigHash = null;
+        redeemScript = null;
+        witnessScript = null;
+        derivedPublicKeys.clear();
+        porCommitment = null;
+        proprietary.clear();
+    }
+
     private Sha256Hash getHashForSignature(Script connectedScript, SigHash localSigHash) {
         Sha256Hash hash;
         if(getWitnessUtxo() != null) {

src/main/java/com/sparrowwallet/drongo/psbt/PSBTOutput.java

@@ -88,6 +88,19 @@ public class PSBTOutput {
         return entries;
     }
 
+    void combine(PSBTOutput psbtOutput) {
+        if(psbtOutput.redeemScript != null) {
+            redeemScript = psbtOutput.redeemScript;
+        }
+
+        if(psbtOutput.witnessScript != null) {
+            witnessScript = psbtOutput.witnessScript;
+        }
+
+        derivedPublicKeys.putAll(psbtOutput.derivedPublicKeys);
+        proprietary.putAll(psbtOutput.proprietary);
+    }
+
     public Script getRedeemScript() {
         return redeemScript;
     }

src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java

@@ -583,14 +583,8 @@ public class Wallet {
         Map<PSBTInput, WalletNode> signingNodes = new LinkedHashMap<>();
         Map<Script, WalletNode> walletOutputScripts = getWalletOutputScripts();
 
-        for(int inputIndex = 0; inputIndex < psbt.getTransaction().getInputs().size(); inputIndex++) {
+        for(PSBTInput psbtInput : psbt.getPsbtInputs()) {
-            TransactionInput txInput = psbt.getTransaction().getInputs().get(inputIndex);
+            TransactionOutput utxo = psbtInput.getUtxo();
-            PSBTInput psbtInput = psbt.getPsbtInputs().get(inputIndex);
-
-            TransactionOutput utxo = psbtInput.getWitnessUtxo();
-            if(utxo == null) {
-                utxo = psbtInput.getNonWitnessUtxo().getOutputs().get((int)txInput.getOutpoint().getIndex());
-            }
 
             if(utxo != null) {
                 Script scriptPubKey = utxo.getScript();
@@ -641,6 +635,45 @@ public class Wallet {
         }
     }
 
+    public void finalise(PSBT psbt) {
+        int threshold = getDefaultPolicy().getNumSignaturesRequired();
+        Map<PSBTInput, WalletNode> signingNodes = getSigningNodes(psbt);
+
+        for(PSBTInput psbtInput : psbt.getPsbtInputs()) {
+            WalletNode signingNode = signingNodes.get(psbtInput);
+            TransactionOutput utxo = psbtInput.getUtxo();
+
+            if(psbtInput.getPartialSignatures().size() >= threshold && signingNode != null && utxo != null) {
+                Transaction transaction = new Transaction();
+
+                TransactionInput txInput;
+                if(getPolicyType().equals(PolicyType.SINGLE)) {
+                    ECKey pubKey = getPubKey(signingNode);
+                    TransactionSignature transactionSignature = psbtInput.getPartialSignature(pubKey);
+                    if(transactionSignature == null) {
+                        throw new IllegalArgumentException("Pubkey of partial signature does not match wallet pubkey");
+                    }
+
+                    txInput = getScriptType().addSpendingInput(transaction, utxo, pubKey, transactionSignature);
+                } else if(getPolicyType().equals(PolicyType.MULTI)) {
+                    List<ECKey> pubKeys = getPubKeys(signingNode);
+                    List<TransactionSignature> signatures = pubKeys.stream().map(psbtInput::getPartialSignature).collect(Collectors.toList());
+                    if(pubKeys.size() != signatures.size()) {
+                        throw new IllegalArgumentException("Pubkeys of partial signatures do not match wallet pubkeys");
+                    }
+
+                    txInput = getScriptType().addMultisigSpendingInput(transaction, utxo, threshold, pubKeys, signatures);
+                } else {
+                    throw new UnsupportedOperationException("Cannot finalise PSBT for policy type " + getPolicyType());
+                }
+
+                psbtInput.setFinalScriptSig(txInput.getScriptSig());
+                psbtInput.setFinalScriptWitness(txInput.getWitness());
+                psbtInput.clearFinalised();
+            }
+        }
+    }
+
     public BitcoinUnit getAutoUnit() {
         for(KeyPurpose keyPurpose : KeyPurpose.values()) {
             for(WalletNode addressNode : getNode(keyPurpose).getChildren()) {