From f2f8c4533bd817775adc0526684d2ea605979785 Mon Sep 17 00:00:00 2001
From: bitromortac <bitromortac@protonmail.com>
Date: Fri, 22 Oct 2021 11:59:44 +0200
Subject: [PATCH] implement option_shutdown_anysegwit

https://github.com/lightningnetwork/lightning-rfc/pull/672

We check the received shutdown script against higher segwit versions and
accept closing to that script if option_shutdown_anysegwit has been
negotiated.
---
 electrum/lnpeer.py                 | 22 +++++++++++++++++-----
 electrum/lnutil.py                 |  7 +++++++
 electrum/lnworker.py               |  3 ++-
 electrum/tests/test_transaction.py | 12 ++++++++++--
 electrum/transaction.py            | 19 +++++++++++++++++++
 5 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/electrum/lnpeer.py b/electrum/lnpeer.py
index 316d10e00..8d8309937 100644
--- a/electrum/lnpeer.py
+++ b/electrum/lnpeer.py
@@ -495,6 +495,9 @@ class Peer(Logger):
         self.lnworker.peer_closed(self)
         self.got_disconnected.set()
 
+    def is_shutdown_anysegwit(self):
+        return self.features.supports(LnFeatures.OPTION_SHUTDOWN_ANYSEGWIT_OPT)
+
     def is_static_remotekey(self):
         return self.features.supports(LnFeatures.OPTION_STATIC_REMOTEKEY_OPT)
 
@@ -1692,11 +1695,20 @@ class Peer(Logger):
         if their_upfront_scriptpubkey:
             if not (their_scriptpubkey == their_upfront_scriptpubkey):
                 raise UpfrontShutdownScriptViolation("remote didn't use upfront shutdown script it commited to in channel opening")
-        # BOLT-02 restrict the scriptpubkey to some templates:
-        if not (match_script_against_template(their_scriptpubkey, transaction.SCRIPTPUBKEY_TEMPLATE_WITNESS_V0)
-                or match_script_against_template(their_scriptpubkey, transaction.SCRIPTPUBKEY_TEMPLATE_P2SH)
-                or match_script_against_template(their_scriptpubkey, transaction.SCRIPTPUBKEY_TEMPLATE_P2PKH)):
-            raise Exception(f'scriptpubkey in received shutdown message does not conform to any template: {their_scriptpubkey.hex()}')
+        else:
+            # BOLT-02 restrict the scriptpubkey to some templates:
+            # order by decreasing dust limit
+            if match_script_against_template(their_scriptpubkey, transaction.SCRIPTPUBKEY_TEMPLATE_P2PKH):
+                pass
+            elif match_script_against_template(their_scriptpubkey, transaction.SCRIPTPUBKEY_TEMPLATE_P2SH):
+                pass
+            elif self.is_shutdown_anysegwit() and match_script_against_template(their_scriptpubkey, transaction.SCRIPTPUBKEY_TEMPLATE_ANYSEGWIT):
+                pass
+            elif match_script_against_template(their_scriptpubkey, transaction.SCRIPTPUBKEY_TEMPLATE_WITNESS_V0):
+                pass
+            else:
+                raise Exception(f'scriptpubkey in received shutdown message does not conform to any template: {their_scriptpubkey.hex()}')
+
         chan_id = chan.channel_id
         if chan_id in self.shutdown_received:
             self.shutdown_received[chan_id].set_result(payload)
diff --git a/electrum/lnutil.py b/electrum/lnutil.py
index 2e3441c9a..6971de53d 100644
--- a/electrum/lnutil.py
+++ b/electrum/lnutil.py
@@ -1026,6 +1026,12 @@ class LnFeatures(IntFlag):
     _ln_feature_contexts[OPTION_TRAMPOLINE_ROUTING_REQ] = (LNFC.INIT | LNFC.NODE_ANN | LNFC.INVOICE)
     _ln_feature_contexts[OPTION_TRAMPOLINE_ROUTING_OPT] = (LNFC.INIT | LNFC.NODE_ANN | LNFC.INVOICE)
 
+    OPTION_SHUTDOWN_ANYSEGWIT_REQ = 1 << 26
+    OPTION_SHUTDOWN_ANYSEGWIT_OPT = 1 << 27
+
+    _ln_feature_contexts[OPTION_SHUTDOWN_ANYSEGWIT_REQ] = (LNFC.INIT | LNFC.NODE_ANN)
+    _ln_feature_contexts[OPTION_SHUTDOWN_ANYSEGWIT_OPT] = (LNFC.INIT | LNFC.NODE_ANN)
+
     # temporary
     OPTION_TRAMPOLINE_ROUTING_REQ_ECLAIR = 1 << 50
     OPTION_TRAMPOLINE_ROUTING_OPT_ECLAIR = 1 << 51
@@ -1114,6 +1120,7 @@ LN_FEATURES_IMPLEMENTED = (
         | LnFeatures.PAYMENT_SECRET_OPT | LnFeatures.PAYMENT_SECRET_REQ
         | LnFeatures.BASIC_MPP_OPT | LnFeatures.BASIC_MPP_REQ
         | LnFeatures.OPTION_TRAMPOLINE_ROUTING_OPT | LnFeatures.OPTION_TRAMPOLINE_ROUTING_REQ
+        | LnFeatures.OPTION_SHUTDOWN_ANYSEGWIT_OPT | LnFeatures.OPTION_SHUTDOWN_ANYSEGWIT_REQ
 )
 
 
diff --git a/electrum/lnworker.py b/electrum/lnworker.py
index 7a60bd885..601261538 100644
--- a/electrum/lnworker.py
+++ b/electrum/lnworker.py
@@ -175,7 +175,8 @@ LNWALLET_FEATURES = BASE_FEATURES\
     | LnFeatures.OPTION_STATIC_REMOTEKEY_REQ\
     | LnFeatures.GOSSIP_QUERIES_REQ\
     | LnFeatures.BASIC_MPP_OPT\
-    | LnFeatures.OPTION_TRAMPOLINE_ROUTING_OPT
+    | LnFeatures.OPTION_TRAMPOLINE_ROUTING_OPT\
+    | LnFeatures.OPTION_SHUTDOWN_ANYSEGWIT_OPT\
 
 LNGOSSIP_FEATURES = BASE_FEATURES\
     | LnFeatures.GOSSIP_QUERIES_OPT\
diff --git a/electrum/tests/test_transaction.py b/electrum/tests/test_transaction.py
index c14ef61d1..3b33ac6e6 100644
--- a/electrum/tests/test_transaction.py
+++ b/electrum/tests/test_transaction.py
@@ -3,7 +3,8 @@ from typing import NamedTuple, Union
 from electrum import transaction, bitcoin
 from electrum.transaction import (convert_raw_tx_to_hex, tx_from_any, Transaction,
                                   PartialTransaction, TxOutpoint, PartialTxInput,
-                                  PartialTxOutput, Sighash)
+                                  PartialTxOutput, Sighash, match_script_against_template,
+                                  SCRIPTPUBKEY_TEMPLATE_ANYSEGWIT)
 from electrum.util import bh2u, bfh
 from electrum.bitcoin import (deserialize_privkey, opcodes,
                               construct_script, construct_witness)
@@ -78,6 +79,13 @@ class TestBCDataStream(ElectrumTestCase):
 
 
 class TestTransaction(ElectrumTestCase):
+    def test_match_against_script_template(self):
+        script = bfh(construct_script([opcodes.OP_5, bytes(29)]))
+        self.assertTrue(match_script_against_template(script, SCRIPTPUBKEY_TEMPLATE_ANYSEGWIT))
+        script = bfh(construct_script([opcodes.OP_NOP, bytes(30)]))
+        self.assertFalse(match_script_against_template(script, SCRIPTPUBKEY_TEMPLATE_ANYSEGWIT))
+        script = bfh(construct_script([opcodes.OP_0, bytes(50)]))
+        self.assertFalse(match_script_against_template(script, SCRIPTPUBKEY_TEMPLATE_ANYSEGWIT))
 
     def test_tx_update_signatures(self):
         tx = tx_from_any("cHNidP8BAFUBAAAAASpcmpT83pj1WBzQAWLGChOTbOt1OJ6mW/OGM7Qk60AxAAAAAAD/////AUBCDwAAAAAAGXapFCMKw3g0BzpCFG8R74QUrpKf6q/DiKwAAAAAAAAA")
@@ -927,7 +935,7 @@ class TestSighashTypes(ElectrumTestCase):
     txin = PartialTxInput(prevout=prevout)
     txin.nsequence=0xffffffff
     txin.script_type='p2sh-p2wsh'
-    txin.witness_script = bfh('56210307b8ae49ac90a048e9b53357a2354b3334e9c8bee813ecb98e99a7e07e8c3ba32103b28f0c28bfab54554ae8c658ac5c3e0ce6e79ad336331f78c428dd43eea8449b21034b8113d703413d57761b8b9781957b8c0ac1dfe69f492580ca4195f50376ba4a21033400f6afecb833092a9a21cfdf1ed1376e58c5d1f47de74683123987e967a8f42103a6d48b1131e94ba04d9737d61acdaa1322008af9602b3b14862c07a1789aac162102d8b661b0b3302ee2f162b09e07a55ad5dfbe673a9f01d9f0c19617681024306b56ae') 
+    txin.witness_script = bfh('56210307b8ae49ac90a048e9b53357a2354b3334e9c8bee813ecb98e99a7e07e8c3ba32103b28f0c28bfab54554ae8c658ac5c3e0ce6e79ad336331f78c428dd43eea8449b21034b8113d703413d57761b8b9781957b8c0ac1dfe69f492580ca4195f50376ba4a21033400f6afecb833092a9a21cfdf1ed1376e58c5d1f47de74683123987e967a8f42103a6d48b1131e94ba04d9737d61acdaa1322008af9602b3b14862c07a1789aac162102d8b661b0b3302ee2f162b09e07a55ad5dfbe673a9f01d9f0c19617681024306b56ae')
     txin.redeem_script = bfh('0020a16b5755f7f6f96dbd65f5f0d6ab9418b89af4b1f14a1bb8a09062c35f0dcb54')
     txin._trusted_value_sats = 987654321
     #Output of Transaction
diff --git a/electrum/transaction.py b/electrum/transaction.py
index 288211957..bd6580240 100644
--- a/electrum/transaction.py
+++ b/electrum/transaction.py
@@ -431,7 +431,23 @@ class OPPushDataGeneric:
                or (isinstance(item, type) and issubclass(item, cls))
 
 
+class OPGeneric:
+    def __init__(self, matcher: Callable=None):
+        if matcher is not None:
+            self.matcher = matcher
+
+    def match(self, op) -> bool:
+        return self.matcher(op)
+
+    @classmethod
+    def is_instance(cls, item):
+        # accept objects that are instances of this class
+        # or other classes that are subclasses
+        return isinstance(item, cls) \
+               or (isinstance(item, type) and issubclass(item, cls))
+
 OPPushDataPubkey = OPPushDataGeneric(lambda x: x in (33, 65))
+OP_ANYSEGWIT_VERSION = OPGeneric(lambda x: x in list(range(opcodes.OP_1, opcodes.OP_16 + 1)))
 
 SCRIPTPUBKEY_TEMPLATE_P2PKH = [opcodes.OP_DUP, opcodes.OP_HASH160,
                                OPPushDataGeneric(lambda x: x == 20),
@@ -440,6 +456,7 @@ SCRIPTPUBKEY_TEMPLATE_P2SH = [opcodes.OP_HASH160, OPPushDataGeneric(lambda x: x
 SCRIPTPUBKEY_TEMPLATE_WITNESS_V0 = [opcodes.OP_0, OPPushDataGeneric(lambda x: x in (20, 32))]
 SCRIPTPUBKEY_TEMPLATE_P2WPKH = [opcodes.OP_0, OPPushDataGeneric(lambda x: x == 20)]
 SCRIPTPUBKEY_TEMPLATE_P2WSH = [opcodes.OP_0, OPPushDataGeneric(lambda x: x == 32)]
+SCRIPTPUBKEY_TEMPLATE_ANYSEGWIT = [OP_ANYSEGWIT_VERSION, OPPushDataGeneric(lambda x: x in list(range(2, 40 + 1)))]
 
 
 def match_script_against_template(script, template) -> bool:
@@ -459,6 +476,8 @@ def match_script_against_template(script, template) -> bool:
         script_item = script[i]
         if OPPushDataGeneric.is_instance(template_item) and template_item.check_data_len(script_item[0]):
             continue
+        if OPGeneric.is_instance(template_item) and template_item.match(script_item[0]):
+            continue
         if template_item != script_item[0]:
             return False
     return True