Fix HTLC signing process

nblockchain · Jun 22, 2021 · 0389f6a · 0389f6a
1 parent 80089f0
commit 0389f6a
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 25 deletions.
diff --git a/src/DotNetLightning.Core/Channel/Channel.fs b/src/DotNetLightning.Core/Channel/Channel.fs
@@ -6,6 +6,7 @@ open DotNetLightning.Utils.Aether
 open DotNetLightning.Chain
 open DotNetLightning.Crypto
 open DotNetLightning.Transactions
+open DotNetLightning.Transactions.Transactions
 open DotNetLightning.Serialization
 open DotNetLightning.Serialization.Msgs
 open NBitcoin
@@ -1245,7 +1246,7 @@ and Channel = {
                 let htlcSigs =
                     sortedHTLCTXs
                     |> List.map(
-                            (fun htlc -> channelPrivKeys.SignHtlcTx htlc.Value remoteNextPerCommitmentPoint)
+                            (fun htlc -> signHtlcTx htlc channelPrivKeys remoteNextPerCommitmentPoint)
                             >> fst
                             >> (fun txSig -> txSig.Signature)
                             )
@@ -1326,29 +1327,30 @@ and Channel = {
                 let sortedHTLCTXs = Commitments.Helpers.sortBothHTLCs htlcTimeoutTxs htlcSuccessTxs
                 do! Commitments.checkSignatureCountMismatch sortedHTLCTXs msg
 
-                let _localHTLCSigs, sortedHTLCTXs =
-                    let localHtlcSigsAndHTLCTxs =
-                        sortedHTLCTXs |> List.map(fun htlc ->
-                            channelPrivKeys.SignHtlcTx htlc.Value localPerCommitmentPoint
-                        )
-                    localHtlcSigsAndHTLCTxs |> List.map(fst), localHtlcSigsAndHTLCTxs |> List.map(snd) |> Seq.cast<IHTLCTx> |> List.ofSeq
+                let htlcTxsAndSignatures =
+                    sortedHTLCTXs 
+                    |> List.zip (msg.HTLCSignatures)
+                    |> List.map(fun (remoteSig, htlc) ->
+                         htlc, signHtlcTx htlc channelPrivKeys localPerCommitmentPoint |> fst, remoteSig
+                    )
 
                 let remoteHTLCPubKey = localPerCommitmentPoint.DeriveHtlcPubKey remoteChannelKeys.HtlcBasepoint
 
-                let checkHTLCSig (htlc: IHTLCTx, remoteECDSASig: LNECDSASignature): Result<_, _> =
-                    let remoteS = TransactionSignature(remoteECDSASig.Value, SigHash.All)
+                let checkHTLCSig (htlc: IHTLCTx, localSignature: TransactionSignature, remoteECDSASig: LNECDSASignature): Result<_, _> =
+                    let remoteSignature = TransactionSignature(remoteECDSASig.Value, SigHash.All)
                     match htlc with
                     | :? HTLCTimeoutTx ->
-                        (Transactions.checkTxFinalized (htlc.Value) (0) (seq [(remoteHTLCPubKey.RawPubKey(), remoteS)]))
+                        (Transactions.checkTxFinalized (htlc.Value) (0) (seq [(remoteHTLCPubKey.RawPubKey(), remoteSignature)]))
                         |> Result.map(box)
                     // we cannot check that htlc-success tx are spendable because we need the payment preimage; thus we only check the remote sig
                     | :? HTLCSuccessTx ->
-                        (Transactions.checkSigAndAdd (htlc) (remoteS) (remoteHTLCPubKey.RawPubKey()))
+                        (htlc :?> HTLCTimeoutTx).Finalize(localSignature, remoteSignature)
+                        |> Result.map(FinalizedTx)
                         |> Result.map(box)
                     | _ -> failwith "Unreachable!"
 
                 let! txList =
-                    List.zip sortedHTLCTXs msg.HTLCSignatures
+                    htlcTxsAndSignatures
                     |> List.map(checkHTLCSig)
                     |> List.sequenceResultA
                     |> expectTransactionErrors

diff --git a/src/DotNetLightning.Core/Crypto/KeyExtensions.fs b/src/DotNetLightning.Core/Crypto/KeyExtensions.fs
@@ -9,6 +9,7 @@ open DotNetLightning.Core.Utils.Extensions
 
 open ResultUtils
 open ResultUtils.Portability
+open NBitcoin.BuilderExtensions
 
 [<AutoOpen>]
 module NBitcoinArithmethicExtensions =
@@ -280,19 +281,6 @@ module KeyExtensions =
             | Some signature -> (signature, psbt)
             | None -> failwithf "Failed to get signature for %A with funding pub key (%A). This should never happen" psbt fundingPubKey
 
-        member this.SignHtlcTx (psbt: PSBT)
-                               (perCommitmentPoint: PerCommitmentPoint)
-                                   : TransactionSignature * PSBT =
-            let htlcPrivKey = perCommitmentPoint.DeriveHtlcPrivKey this.HtlcBasepointSecret
-            let htlcPubKey = htlcPrivKey.HtlcPubKey()
-            psbt.SignWithKeys(htlcPrivKey.RawKey()) |> ignore
-            match psbt.GetMatchingSig(htlcPubKey.RawPubKey()) with
-            | Some signature -> (signature, psbt)
-            | None ->
-                failwithf
-                    "failed to get htlc signature for %A. with htlc pubkey (%A) and perCommitmentPoint (%A)"
-                    psbt htlcPubKey perCommitmentPoint
-
     /// This is the node-wide master key which is also used for
     /// transport-level encryption. The channel's keys are derived from
     /// this via BIP32 key derivation where `channelIndex` is the child

diff --git a/src/DotNetLightning.Core/Transactions/Transactions.fs b/src/DotNetLightning.Core/Transactions/Transactions.fs
@@ -519,6 +519,13 @@ module Transactions =
             | Error e -> failwithf "%A" e
 
     let sign(tx, key) = signCore(tx, key, true)
+    let signHtlcTx (htlc: IHTLCTx) (channelPrivKeys: ChannelPrivKeys) (perCommitmentPoint: PerCommitmentPoint) =
+        let htlcPrivKey =
+            perCommitmentPoint.DeriveHtlcPrivKey
+                channelPrivKeys.HtlcBasepointSecret
+
+        sign(htlc, htlcPrivKey.RawKey())
+
     let makeHTLCTimeoutTx (commitTx: Transaction)
                           (localDustLimit: Money)
                           (localRevocationPubKey: RevocationPubKey)