Derive weight from witness for p2wsh inputs

If witness is available for p2wsh inputs we should be using the size of
the witness to get expected input weight. Doing this incidentally
enables p2wsh inputs for the payjoin sender.

Related issue: payjoin#659

Author: arminsabouri

payjoin/src/core/psbt/mod.rs

@@ -206,7 +206,27 @@ impl InternalInputPair<'_> {
                 }
             }
             P2wpkh => Ok(InputWeightPrediction::P2WPKH_MAX),
-            P2wsh => Err(InputWeightError::NotSupported),
+            P2wsh =>
+                if !self.txin.witness.is_empty() {
+                    Ok(InputWeightPrediction::new(
+                        0,
+                        self.txin.witness.iter().map(|el| el.len()).collect::<Vec<_>>(),
+                    ))
+                } else {
+                    let iwp = self
+                        .psbtin
+                        .final_script_witness
+                        .as_ref()
+                        .filter(|w| !w.is_empty())
+                        .map(|w| {
+                            InputWeightPrediction::new(
+                                0,
+                                w.iter().map(|el| el.len()).collect::<Vec<_>>(),
+                            )
+                        })
+                        .ok_or(InputWeightError::NotSupported)?;
+                    Ok(iwp)
+                },
             P2tr => Ok(InputWeightPrediction::P2TR_KEY_DEFAULT_SIGHASH),
             _ => Err(AddressTypeError::UnknownAddressType.into()),
         }?;
@@ -245,6 +265,8 @@ pub(crate) enum InternalPsbtInputError {
     AddressType(AddressTypeError),
     InvalidScriptPubKey(AddressType),
     WeightError(InputWeightError),
+    /// Weight was provided but can be calculated from available information
+    ProvidedUnnecessaryWeight,
 }
 
 impl fmt::Display for InternalPsbtInputError {
@@ -256,6 +278,7 @@ impl fmt::Display for InternalPsbtInputError {
             Self::AddressType(_) => write!(f, "invalid address type"),
             Self::InvalidScriptPubKey(e) => write!(f, "provided script was not a valid type of {e}"),
             Self::WeightError(e) => write!(f, "{e}"),
+            Self::ProvidedUnnecessaryWeight => write!(f, "weight was provided but can be calculated from available information"),
         }
     }
 }
@@ -269,6 +292,7 @@ impl std::error::Error for InternalPsbtInputError {
             Self::AddressType(error) => Some(error),
             Self::InvalidScriptPubKey(_) => None,
             Self::WeightError(error) => Some(error),
+            Self::ProvidedUnnecessaryWeight => None,
         }
     }
 }

payjoin/src/core/receive/mod.rs

@@ -61,15 +61,13 @@ impl InputPair {
         let raw = InternalInputPair { txin: &txin, psbtin: &psbtin };
         raw.validate_utxo()?;
 
-        let expected_weight = match raw.expected_input_weight() {
-            Ok(weight) => weight,
-            Err(InputWeightError::NotSupported) =>
-                if let Some(weight) = expected_weight {
-                    weight
-                } else {
-                    return Err(InternalPsbtInputError::from(InputWeightError::NotSupported).into());
-                },
-            Err(e) => return Err(InternalPsbtInputError::from(e).into()),
+        let expected_weight = match (raw.expected_input_weight(), expected_weight) {
+            (Ok(_), Some(_)) => {
+                return Err(InternalPsbtInputError::ProvidedUnnecessaryWeight.into());
+            }
+            (Ok(weight), None) => weight,
+            (Err(InputWeightError::NotSupported), Some(expected_weight)) => expected_weight,
+            (Err(e), _) => return Err(InternalPsbtInputError::from(e).into()),
         };
 
         let input_pair = Self { expected_weight, txin, psbtin };
@@ -246,7 +244,9 @@ mod tests {
     use bitcoin::key::{PublicKey, WPubkeyHash};
     use bitcoin::secp256k1::Secp256k1;
     use bitcoin::transaction::InputWeightPrediction;
-    use bitcoin::{Amount, PubkeyHash, ScriptBuf, ScriptHash, Txid, WScriptHash, XOnlyPublicKey};
+    use bitcoin::{
+        witness, Amount, PubkeyHash, ScriptBuf, ScriptHash, Txid, WScriptHash, XOnlyPublicKey,
+    };
     use payjoin_test_utils::{DUMMY20, DUMMY32};
 
     use super::*;
@@ -457,7 +457,7 @@ mod tests {
 
         let p2wsh_pair = InputPair::new(
             TxIn { previous_output: outpoint, sequence, ..Default::default() },
-            psbt::Input { witness_utxo: Some(p2wsh_txout), ..Default::default() },
+            psbt::Input { witness_utxo: Some(p2wsh_txout.clone()), ..Default::default() },
             None,
         );
         // P2wsh is not supported when expected weight is not provided
@@ -475,7 +475,39 @@ mod tests {
                 .err()
                 .unwrap(),
             PsbtInputError::from(InvalidScriptPubKey(AddressType::P2wsh))
-        )
+        );
+
+        let mut dummy_witness = witness::Witness::new();
+        dummy_witness.push(DUMMY32);
+        let txin = TxIn {
+            previous_output: outpoint,
+            witness: dummy_witness.clone(),
+            ..Default::default()
+        };
+        let input_weight = Weight::from_non_witness_data_size(txin.base_size() as u64)
+            + Weight::from_witness_data_size(dummy_witness.size() as u64);
+
+        // Add the witness straight to the txin
+        let psbtin = psbt::Input { witness_utxo: Some(p2wsh_txout.clone()), ..Default::default() };
+        let p2wsh_pair = InputPair::new(txin, psbtin, None).expect("witness is provided for p2wsh");
+        assert_eq!(p2wsh_pair.expected_weight, input_weight);
+        // Same check but add the witness to the psbtin
+        let txin = TxIn { previous_output: outpoint, ..Default::default() };
+        let psbtin = psbt::Input {
+            witness_utxo: Some(p2wsh_txout),
+            final_script_witness: Some(dummy_witness),
+            ..Default::default()
+        };
+        let p2wsh_pair = InputPair::new(txin.clone(), psbtin.clone(), None)
+            .expect("witness is provided for p2wsh");
+        assert_eq!(p2wsh_pair.expected_weight, input_weight);
+
+        // Should error out if expected weight is provided and witness is provided
+        let p2wsh_pair = InputPair::new(txin, psbtin, Some(expected_weight));
+        assert_eq!(
+            p2wsh_pair.err().unwrap(),
+            PsbtInputError::from(InternalPsbtInputError::ProvidedUnnecessaryWeight)
+        );
     }
 
     #[test]

payjoin/src/core/receive/v1/mod.rs

@@ -1231,13 +1231,13 @@ pub(crate) mod test {
         let input_pair_1 = InputPair::new(
             TxIn { previous_output: ot1, sequence: Sequence::MAX, ..Default::default() },
             Input { witness_utxo: Some(txout.clone()), ..Default::default() },
-            Some(Weight::from_wu(777)),
+            None,
         )
         .unwrap();
         let input_pair_2 = InputPair::new(
             TxIn { previous_output: ot2, sequence: Sequence::MAX, ..Default::default() },
             Input { witness_utxo: Some(txout), ..Default::default() },
-            Some(Weight::from_wu(777)),
+            None,
         )
         .unwrap();