From 7f840cf03b60727d9204068e3e87d531e73872bd Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Mon, 1 Apr 2024 17:59:08 +0200
Subject: [PATCH] signpsbt+lnd: fix signing for P2WKH

---
 cmd/chantools/signpsbt.go | 19 ++++++++++++++-----
 lnd/signer.go             |  8 ++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/cmd/chantools/signpsbt.go b/cmd/chantools/signpsbt.go
index 877b1d6..927881b 100644
--- a/cmd/chantools/signpsbt.go
+++ b/cmd/chantools/signpsbt.go
@@ -10,6 +10,7 @@ import (
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/btcutil/hdkeychain"
 	"github.com/btcsuite/btcd/btcutil/psbt"
+	"github.com/btcsuite/btcd/txscript"
 	"github.com/lightninglabs/chantools/lnd"
 	"github.com/spf13/cobra"
 )
@@ -153,17 +154,25 @@ func signPsbt(rootKey *hdkeychain.ExtendedKey,
 		return fmt.Errorf("could not derive local key: %w", err)
 	}
 
-	if len(packet.Inputs[inputIndex].WitnessScript) == 0 {
-		return fmt.Errorf("invalid PSBT, input %d is missing witness "+
-			"script", inputIndex)
-	}
-	witnessScript := packet.Inputs[inputIndex].WitnessScript
 	if packet.Inputs[inputIndex].WitnessUtxo == nil {
 		return fmt.Errorf("invalid PSBT, input %d is missing witness "+
 			"UTXO", inputIndex)
 	}
 	utxo := packet.Inputs[inputIndex].WitnessUtxo
 
+	// The signing is a bit different for P2WPKH, we need to specify the
+	// pk script as the witness script.
+	var witnessScript []byte
+	if txscript.IsPayToWitnessPubKeyHash(utxo.PkScript) {
+		witnessScript = utxo.PkScript
+	} else {
+		if len(packet.Inputs[inputIndex].WitnessScript) == 0 {
+			return fmt.Errorf("invalid PSBT, input %d is missing "+
+				"witness script", inputIndex)
+		}
+		witnessScript = packet.Inputs[inputIndex].WitnessScript
+	}
+
 	localPrivateKey, err := localKey.ECPrivKey()
 	if err != nil {
 		return fmt.Errorf("error getting private key: %w", err)
diff --git a/lnd/signer.go b/lnd/signer.go
index da51e1c..328eb67 100644
--- a/lnd/signer.go
+++ b/lnd/signer.go
@@ -210,6 +210,14 @@ func (s *Signer) AddPartialSignatureForPrivateKey(packet *psbt.Packet,
 	if err != nil {
 		return fmt.Errorf("error creating PSBT updater: %w", err)
 	}
+
+	// If the witness script is the pk script for a P2WPKH output, then we
+	// need to blank it out for the PSBT code, otherwise it interprets it as
+	// a P2WSH.
+	if txscript.IsPayToWitnessPubKeyHash(utxo.PkScript) {
+		witnessScript = nil
+	}
+
 	status, err := updater.Sign(
 		inputIndex, ourSig, privateKey.PubKey().SerializeCompressed(),
 		nil, witnessScript,