lnsweep: refactor: maybe_reveal_preimage_for_htlc

Author: SomberNight

electrum/lnsweep.py

@@ -426,7 +426,8 @@ def txs_htlc(
                         privkey=our_localdelayed_privkey.get_secret_bytes(),
                         is_revocation=False,
                 ):
-                    txs[actual_htlc_tx.txid() + f':{output_idx}'] = SweepInfo(
+                    prevout = actual_htlc_tx.txid() + f':{output_idx}'
+                    txs[prevout] = SweepInfo(
                         name=f'second-stage-htlc:{output_idx}',
                         cltv_abs=0,
                         txin=sweep_txin,
@@ -451,23 +452,12 @@ def txs_htlc(
             # note: it is the first stage (witness of htlc_tx) that reveals the preimage,
             #       so if we are already in second stage, it is already revealed.
             #       However, here, we don't make a distinction.
-            if not chan.lnworker.is_complete_mpp(htlc.payment_hash):
-                # - do not redeem this, it might publish the preimage of an incomplete MPP
-                # - OTOH maybe this chan just got closed, and we are still receiving new htlcs
-                #   for this MPP set. So the MPP set might still transition to complete!
-                #   The MPP_TIMEOUT is only around 2 minutes, so this window is short.
-                #   The default keep_watching logic in lnwatcher is sufficient to call us again.
-                continue
-            if htlc.payment_hash in chan.lnworker.dont_settle_htlcs:
-                prevout = ctx.txid() + ':%d' % ctx_output_idx
-                txs[prevout] = KeepWatchingTXO(
-                    name=f"our_ctx_htlc_{ctx_output_idx}_for_hold_invoice",
-                    until_height=htlc.cltv_abs,
-                )
-                continue
-            preimage = chan.lnworker.get_preimage(htlc.payment_hash)
+            preimage = _maybe_reveal_preimage_for_htlc(
+                chan=chan, htlc=htlc, txs=txs,
+                prevout=ctx.txid() + ':%d' % ctx_output_idx,
+                sweep_info_name=f"our_ctx_htlc_{ctx_output_idx}",
+            )
             if not preimage:
-                # we might not have the preimage if this is a hold invoice
                 continue
         try:
             txs_htlc(
@@ -481,6 +471,32 @@ def txs_htlc(
     return txs
 
 
+def _maybe_reveal_preimage_for_htlc(
+    *,
+    chan: 'AbstractChannel',
+    htlc: 'UpdateAddHtlc',
+    txs: Dict[str, MaybeSweepInfo],  # mutated in-place!
+    prevout: str,  # commitment txid + output_idx  (so always for first stage)
+    sweep_info_name: str,
+) -> Optional[bytes]:
+    """Given a Remote-added-HTLC, return the preimage if it's okay to reveal it on-chain."""
+    if not chan.lnworker.is_complete_mpp(htlc.payment_hash):
+        # - do not redeem this, it might publish the preimage of an incomplete MPP
+        # - OTOH maybe this chan just got closed, and we are still receiving new htlcs
+        #   for this MPP set. So the MPP set might still transition to complete!
+        #   The MPP_TIMEOUT is only around 2 minutes, so this window is short.
+        #   The default keep_watching logic in lnwatcher is sufficient to call us again.
+        return None
+    if htlc.payment_hash in chan.lnworker.dont_settle_htlcs:
+        txs[prevout] = KeepWatchingTXO(
+            name=sweep_info_name + "_for_hold_invoice",
+            until_height=htlc.cltv_abs,
+        )
+        return None
+    preimage = chan.lnworker.get_preimage(htlc.payment_hash)
+    return preimage
+
+
 def extract_ctx_secrets(chan: 'Channel', ctx: Transaction):
     # note: the remote sometimes has two valid non-revoked commitment transactions,
     # either of which could be broadcast
@@ -762,23 +778,12 @@ def tx_htlc(
         preimage = None
         is_received_htlc = direction == RECEIVED
         if not is_received_htlc and not is_revocation:
-            if not chan.lnworker.is_complete_mpp(htlc.payment_hash):
-                # - do not redeem this, it might publish the preimage of an incomplete MPP
-                # - OTOH maybe this chan just got closed, and we are still receiving new htlcs
-                #   for this MPP set. So the MPP set might still transition to complete!
-                #   The MPP_TIMEOUT is only around 2 minutes, so this window is short.
-                #   The default keep_watching logic in lnwatcher is sufficient to call us again.
-                continue
-            if htlc.payment_hash in chan.lnworker.dont_settle_htlcs:
-                prevout = ctx.txid() + ':%d' % ctx_output_idx
-                txs[prevout] = KeepWatchingTXO(
-                    name=f"their_ctx_htlc_{ctx_output_idx}_for_hold_invoice",
-                    until_height=htlc.cltv_abs,
-                )
-                continue
-            preimage = chan.lnworker.get_preimage(htlc.payment_hash)
+            preimage = _maybe_reveal_preimage_for_htlc(
+                chan=chan, htlc=htlc, txs=txs,
+                prevout=ctx.txid() + ':%d' % ctx_output_idx,
+                sweep_info_name=f"their_ctx_htlc_{ctx_output_idx}",
+            )
             if not preimage:
-                # we might not have the preimage if this is a hold invoice
                 continue
         tx_htlc(
             htlc=htlc,