From 45bd5b47516dad0a1f1ba6d37fba9cf1218f2a8f Mon Sep 17 00:00:00 2001
From: junderw <junderwood@bitcoinbank.co.jp>
Date: Thu, 4 Jul 2019 17:35:39 +0900
Subject: [PATCH] Check for signatures, add setSequence

---
 src/psbt.js     | 60 ++++++++++++++++++++++++++++++++++++++++++++++
 ts_src/psbt.ts  | 63 +++++++++++++++++++++++++++++++++++++++++++++++++
 types/psbt.d.ts |  1 +
 3 files changed, 124 insertions(+)

diff --git a/src/psbt.js b/src/psbt.js
index a85a23a..eaa8d62 100644
--- a/src/psbt.js
+++ b/src/psbt.js
@@ -80,16 +80,31 @@ class Psbt extends bip174_1.Psbt {
     dpew(this, 'opts', false, true);
   }
   setVersion(version) {
+    check32Bit(version);
+    checkInputsForPartialSig(this.inputs, 'setVersion');
     this.__TX.version = version;
     this.__TX_BUF_CACHE = undefined;
     return this;
   }
   setLocktime(locktime) {
+    check32Bit(locktime);
+    checkInputsForPartialSig(this.inputs, 'setLocktime');
     this.__TX.locktime = locktime;
     this.__TX_BUF_CACHE = undefined;
     return this;
   }
+  setSequence(inputIndex, sequence) {
+    check32Bit(sequence);
+    checkInputsForPartialSig(this.inputs, 'setSequence');
+    if (this.__TX.ins.length <= inputIndex) {
+      throw new Error('Input index too high');
+    }
+    this.__TX.ins[inputIndex].sequence = sequence;
+    this.__TX_BUF_CACHE = undefined;
+    return this;
+  }
   addInput(inputData) {
+    checkInputsForPartialSig(this.inputs, 'addInput');
     const self = this;
     const inputAdder = (_inputData, txBuf) => {
       if (
@@ -119,6 +134,7 @@ class Psbt extends bip174_1.Psbt {
     return super.addInput(inputData, inputAdder);
   }
   addOutput(outputData) {
+    checkInputsForPartialSig(this.inputs, 'addOutput');
     const { address } = outputData;
     if (typeof address === 'string') {
       const { network } = this.opts;
@@ -578,3 +594,47 @@ function checkTxEmpty(tx) {
     throw new Error('Format Error: Transaction ScriptSigs are not empty');
   }
 }
+function checkInputsForPartialSig(inputs, action) {
+  inputs.forEach(input => {
+    let throws = false;
+    if ((input.partialSig || []).length > 0) {
+      if (input.sighashType !== undefined) {
+        const whitelist = [];
+        const isAnyoneCanPay =
+          input.sighashType & transaction_1.Transaction.SIGHASH_ANYONECANPAY;
+        if (isAnyoneCanPay) whitelist.push('addInput');
+        if (!isAnyoneCanPay && action === 'addInput') {
+          throws = true;
+        }
+        const hashType = input.sighashType & 0x1f;
+        switch (hashType) {
+          case transaction_1.Transaction.SIGHASH_ALL:
+            break;
+          case transaction_1.Transaction.SIGHASH_SINGLE:
+          case transaction_1.Transaction.SIGHASH_NONE:
+            whitelist.push('addOutput');
+            whitelist.push('setSequence');
+            break;
+        }
+        if (whitelist.indexOf(action) === -1) {
+          throws = true;
+        }
+      } else {
+        throws = true;
+      }
+    }
+    if (throws) {
+      throw new Error('Can not modify transaction, signatures exist.');
+    }
+  });
+}
+function check32Bit(num) {
+  if (
+    typeof num !== 'number' ||
+    num !== Math.floor(num) ||
+    num > 0xffffffff ||
+    num < 0
+  ) {
+    throw new Error('Invalid 32 bit integer');
+  }
+}
diff --git a/ts_src/psbt.ts b/ts_src/psbt.ts
index 14f0c08..0b019a3 100644
--- a/ts_src/psbt.ts
+++ b/ts_src/psbt.ts
@@ -108,18 +108,34 @@ export class Psbt extends PsbtBase {
   }
 
   setVersion(version: number): this {
+    check32Bit(version);
+    checkInputsForPartialSig(this.inputs, 'setVersion');
     this.__TX.version = version;
     this.__TX_BUF_CACHE = undefined;
     return this;
   }
 
   setLocktime(locktime: number): this {
+    check32Bit(locktime);
+    checkInputsForPartialSig(this.inputs, 'setLocktime');
     this.__TX.locktime = locktime;
     this.__TX_BUF_CACHE = undefined;
     return this;
   }
 
+  setSequence(inputIndex: number, sequence: number): this {
+    check32Bit(sequence);
+    checkInputsForPartialSig(this.inputs, 'setSequence');
+    if (this.__TX.ins.length <= inputIndex) {
+      throw new Error('Input index too high');
+    }
+    this.__TX.ins[inputIndex].sequence = sequence;
+    this.__TX_BUF_CACHE = undefined;
+    return this;
+  }
+
   addInput(inputData: TransactionInput): this {
+    checkInputsForPartialSig(this.inputs, 'addInput');
     const self = this;
     const inputAdder = (
       _inputData: TransactionInput,
@@ -152,6 +168,7 @@ export class Psbt extends PsbtBase {
   }
 
   addOutput(outputData: TransactionOutput): this {
+    checkInputsForPartialSig(this.inputs, 'addOutput');
     const { address } = outputData as any;
     if (typeof address === 'string') {
       const { network } = this.opts;
@@ -734,3 +751,49 @@ function checkTxEmpty(tx: Transaction): void {
     throw new Error('Format Error: Transaction ScriptSigs are not empty');
   }
 }
+
+function checkInputsForPartialSig(inputs: PsbtInput[], action: string): void {
+  inputs.forEach(input => {
+    let throws = false;
+    if ((input.partialSig || []).length > 0) {
+      if (input.sighashType !== undefined) {
+        const whitelist: string[] = [];
+        const isAnyoneCanPay =
+          input.sighashType & Transaction.SIGHASH_ANYONECANPAY;
+        if (isAnyoneCanPay) whitelist.push('addInput');
+        if (!isAnyoneCanPay && action === 'addInput') {
+          throws = true;
+        }
+        const hashType = input.sighashType & 0x1f;
+        switch (hashType) {
+          case Transaction.SIGHASH_ALL:
+            break;
+          case Transaction.SIGHASH_SINGLE:
+          case Transaction.SIGHASH_NONE:
+            whitelist.push('addOutput');
+            whitelist.push('setSequence');
+            break;
+        }
+        if (whitelist.indexOf(action) === -1) {
+          throws = true;
+        }
+      } else {
+        throws = true;
+      }
+    }
+    if (throws) {
+      throw new Error('Can not modify transaction, signatures exist.');
+    }
+  });
+}
+
+function check32Bit(num: number): void {
+  if (
+    typeof num !== 'number' ||
+    num !== Math.floor(num) ||
+    num > 0xffffffff ||
+    num < 0
+  ) {
+    throw new Error('Invalid 32 bit integer');
+  }
+}
diff --git a/types/psbt.d.ts b/types/psbt.d.ts
index a203854..1aa6f28 100644
--- a/types/psbt.d.ts
+++ b/types/psbt.d.ts
@@ -13,6 +13,7 @@ export declare class Psbt extends PsbtBase {
     constructor(opts?: PsbtOptsOptional);
     setVersion(version: number): this;
     setLocktime(locktime: number): this;
+    setSequence(inputIndex: number, sequence: number): this;
     addInput(inputData: TransactionInput): this;
     addOutput(outputData: TransactionOutput): this;
     extractTransaction(): Transaction;