lnwatcher: sweep to_remote and to_local outputs if they close

lib/lnhtlc.py

@@ -302,12 +302,12 @@ class HTLCStateMachine(PrintError):
     @property
     def points(self):
         last_small_num = self.local_state.ctn
-        next_small_num = last_small_num + 2
         this_small_num = last_small_num + 1
-        last_secret = get_per_commitment_secret_from_seed(self.local_state.per_commitment_secret_seed, 2**48-last_small_num-1)
-        this_secret = get_per_commitment_secret_from_seed(self.local_state.per_commitment_secret_seed, 2**48-this_small_num-1)
+        next_small_num = last_small_num + 2
+        last_secret = get_per_commitment_secret_from_seed(self.local_state.per_commitment_secret_seed, RevocationStore.START_INDEX - last_small_num)
+        this_secret = get_per_commitment_secret_from_seed(self.local_state.per_commitment_secret_seed, RevocationStore.START_INDEX - this_small_num)
         this_point = secret_to_pubkey(int.from_bytes(this_secret, 'big'))
-        next_secret = get_per_commitment_secret_from_seed(self.local_state.per_commitment_secret_seed, 2**48-next_small_num-1)
+        next_secret = get_per_commitment_secret_from_seed(self.local_state.per_commitment_secret_seed, RevocationStore.START_INDEX - next_small_num)
         next_point = secret_to_pubkey(int.from_bytes(next_secret, 'big'))
         return last_secret, this_point, next_point
 

lib/lnutil.py

@@ -22,11 +22,19 @@ LocalState = namedtuple("LocalState", ["ctn", "per_commitment_secret_seed", "amo
 ChannelConstraints = namedtuple("ChannelConstraints", ["capacity", "is_initiator", "funding_txn_minimum_depth"])
 #OpenChannel = namedtuple("OpenChannel", ["channel_id", "short_channel_id", "funding_outpoint", "local_config", "remote_config", "remote_state", "local_state", "constraints", "node_id"])
 
+
+class UnableToDeriveSecret(Exception): pass
+
+
 class RevocationStore:
     """ taken from lnd """
+
+    START_INDEX = 2 ** 48 - 1
+
     def __init__(self):
-        self.buckets = [None] * 48
-        self.index = 2**48 - 1
+        self.buckets = [None] * 49
+        self.index = self.START_INDEX
+
     def add_next_entry(self, hsh):
         new_element = ShachainElement(index=self.index, secret=hsh)
         bucket = count_trailing_zeros(self.index)
@@ -38,8 +46,21 @@ class RevocationStore:
                 raise Exception("hash is not derivable: {} {} {}".format(bh2u(e.secret), bh2u(this_bucket.secret), this_bucket.index))
         self.buckets[bucket] = new_element
         self.index -= 1
+
+    def retrieve_secret(self, index: int) -> bytes:
+        for bucket in self.buckets:
+            if bucket is None:
+                raise UnableToDeriveSecret()
+            try:
+                element = shachain_derive(bucket, index)
+            except UnableToDeriveSecret:
+                continue
+            return element.secret
+        raise UnableToDeriveSecret()
+
     def serialize(self):
         return {"index": self.index, "buckets": [[bh2u(k.secret), k.index] if k is not None else None for k in self.buckets]}
+
     @staticmethod
     def from_json_obj(decoded_json_obj):
         store = RevocationStore()
@@ -47,8 +68,10 @@ class RevocationStore:
         store.buckets = [k if k is None else decode(k) for k in decoded_json_obj["buckets"]]
         store.index = decoded_json_obj["index"]
         return store
+
     def __eq__(self, o):
         return type(o) is RevocationStore and self.serialize() == o.serialize()
+
     def __hash__(self):
         return hash(json.dumps(self.serialize(), sort_keys=True))
 
@@ -59,8 +82,17 @@ def count_trailing_zeros(index):
     except ValueError:
         return 48
 
-def shachain_derive(element, toIndex):
-    return ShachainElement(get_per_commitment_secret_from_seed(element.secret, toIndex, count_trailing_zeros(element.index)), toIndex)
+def shachain_derive(element, to_index):
+    def get_prefix(index, pos):
+        mask = (1 << 64) - 1 - ((1 << pos) - 1)
+        return index & mask
+    from_index = element.index
+    zeros = count_trailing_zeros(from_index)
+    if from_index != get_prefix(to_index, zeros):
+        raise UnableToDeriveSecret("prefixes are different; index not derivable")
+    return ShachainElement(
+        get_per_commitment_secret_from_seed(element.secret, to_index, zeros),
+        to_index)
 
 ShachainElement = namedtuple("ShachainElement", ["secret", "index"])
 ShachainElement.__str__ = lambda self: "ShachainElement(" + bh2u(self.secret) + "," + str(self.index) + ")"
@@ -76,26 +108,35 @@ def get_per_commitment_secret_from_seed(seed: bytes, i: int, bits: int = 48) ->
     bajts = bytes(per_commitment_secret)
     return bajts
 
-def secret_to_pubkey(secret):
+def secret_to_pubkey(secret: int) -> bytes:
     assert type(secret) is int
-    return (ecc.generator() * secret).get_public_key_bytes()
+    return ecc.ECPrivkey.from_secret_scalar(secret).get_public_key_bytes(compressed=True)
 
-def derive_pubkey(basepoint, per_commitment_point):
+def derive_pubkey(basepoint: bytes, per_commitment_point: bytes) -> bytes:
     p = ecc.ECPubkey(basepoint) + ecc.generator() * ecc.string_to_number(sha256(per_commitment_point + basepoint))
     return p.get_public_key_bytes()
 
-def derive_privkey(secret, per_commitment_point):
+def derive_privkey(secret: int, per_commitment_point: bytes) -> int:
     assert type(secret) is int
     basepoint = secret_to_pubkey(secret)
     basepoint = secret + ecc.string_to_number(sha256(per_commitment_point + basepoint))
     basepoint %= CURVE_ORDER
     return basepoint
 
-def derive_blinded_pubkey(basepoint, per_commitment_point):
+def derive_blinded_pubkey(basepoint: bytes, per_commitment_point: bytes) -> bytes:
     k1 = ecc.ECPubkey(basepoint) * ecc.string_to_number(sha256(basepoint + per_commitment_point))
     k2 = ecc.ECPubkey(per_commitment_point) * ecc.string_to_number(sha256(per_commitment_point + basepoint))
     return (k1 + k2).get_public_key_bytes()
 
+def derive_blinded_privkey(basepoint_secret: bytes, per_commitment_secret: bytes) -> bytes:
+    basepoint = ecc.ECPrivkey(basepoint_secret).get_public_key_bytes(compressed=True)
+    per_commitment_point = ecc.ECPrivkey(per_commitment_secret).get_public_key_bytes(compressed=True)
+    k1 = ecc.string_to_number(basepoint_secret) * ecc.string_to_number(sha256(basepoint + per_commitment_point))
+    k2 = ecc.string_to_number(per_commitment_secret) * ecc.string_to_number(sha256(per_commitment_point + basepoint))
+    sum = (k1 + k2) % ecc.CURVE_ORDER
+    return ecc.number_to_string(sum, CURVE_ORDER)
+
+
 def make_htlc_tx_output(amount_msat, local_feerate, revocationpubkey, local_delayedpubkey, success, to_self_delay):
     assert type(amount_msat) is int
     assert type(local_feerate) is int
@@ -250,10 +291,8 @@ def make_commitment(ctn, local_funding_pubkey, remote_funding_pubkey,
         'sequence': sequence
     }]
     # commitment tx outputs
-    local_script = bytes([opcodes.OP_IF]) + bfh(push_script(bh2u(revocation_pubkey))) + bytes([opcodes.OP_ELSE]) + bitcoin.add_number_to_script(to_self_delay) \
-                   + bytes([opcodes.OP_CSV, opcodes.OP_DROP]) + bfh(push_script(bh2u(delayed_pubkey))) + bytes([opcodes.OP_ENDIF, opcodes.OP_CHECKSIG])
-    local_address = bitcoin.redeem_script_to_address('p2wsh', bh2u(local_script))
-    remote_address = bitcoin.pubkey_to_address('p2wpkh', bh2u(remote_payment_pubkey))
+    local_address = make_commitment_output_to_local_address(revocation_pubkey, to_self_delay, delayed_pubkey)
+    remote_address = make_commitment_output_to_remote_address(remote_payment_pubkey)
     # TODO trim htlc outputs here while also considering 2nd stage htlc transactions
     fee = local_feerate * overall_weight(len(htlcs))
     assert type(fee) is int
@@ -284,6 +323,20 @@ def make_commitment(ctn, local_funding_pubkey, remote_funding_pubkey,
 
     return tx
 
+def make_commitment_output_to_local_witness_script(
+        revocation_pubkey: bytes, to_self_delay: int, delayed_pubkey: bytes) -> bytes:
+    local_script = bytes([opcodes.OP_IF]) + bfh(push_script(bh2u(revocation_pubkey))) + bytes([opcodes.OP_ELSE]) + bitcoin.add_number_to_script(to_self_delay) \
+                   + bytes([opcodes.OP_CSV, opcodes.OP_DROP]) + bfh(push_script(bh2u(delayed_pubkey))) + bytes([opcodes.OP_ENDIF, opcodes.OP_CHECKSIG])
+    return local_script
+
+def make_commitment_output_to_local_address(
+        revocation_pubkey: bytes, to_self_delay: int, delayed_pubkey: bytes) -> str:
+    local_script = make_commitment_output_to_local_witness_script(revocation_pubkey, to_self_delay, delayed_pubkey)
+    return bitcoin.redeem_script_to_address('p2wsh', bh2u(local_script))
+
+def make_commitment_output_to_remote_address(remote_payment_pubkey: bytes) -> str:
+    return bitcoin.pubkey_to_address('p2wpkh', bh2u(remote_payment_pubkey))
+
 def sign_and_get_sig_string(tx, local_config, remote_config):
     pubkeys = sorted([bh2u(local_config.multisig_key.pubkey), bh2u(remote_config.multisig_key.pubkey)])
     tx.sign({bh2u(local_config.multisig_key.pubkey): (local_config.multisig_key.privkey, True)})
@@ -306,6 +359,13 @@ def get_obscured_ctn(ctn, local, remote):
     mask = int.from_bytes(sha256(local + remote)[-6:], 'big')
     return ctn ^ mask
 
+def extract_ctn_from_tx(tx, txin_index, local_payment_basepoint, remote_payment_basepoint):
+    tx.deserialize()
+    locktime = tx.locktime
+    sequence = tx.inputs()[txin_index]['sequence']
+    obs = ((sequence & 0xffffff) << 24) + (locktime & 0xffffff)
+    return get_obscured_ctn(obs, local_payment_basepoint, remote_payment_basepoint)
+
 def overall_weight(num_htlc):
     return 500 + 172 * num_htlc + 224
 

lib/lnwatcher.py

@@ -1,12 +1,20 @@
-from .util import PrintError
-from .lnutil import funding_output_script
-from .bitcoin import redeem_script_to_address
+from .util import PrintError, bh2u, bfh, NoDynamicFeeEstimates
+from .lnutil import (funding_output_script, extract_ctn_from_tx, derive_privkey,
+                     get_per_commitment_secret_from_seed, derive_pubkey,
+                     make_commitment_output_to_remote_address,
+                     RevocationStore, UnableToDeriveSecret)
+from . import lnutil
+from .bitcoin import redeem_script_to_address, TYPE_ADDRESS
+from . import transaction
+from .transaction import Transaction
+from . import ecc
 
 class LNWatcher(PrintError):
 
     def __init__(self, network):
         self.network = network
         self.watched_channels = {}
+        self.address_status = {}  # addr -> status
 
     def parse_response(self, response):
         if response.get('error'):
@@ -15,8 +23,7 @@ class LNWatcher(PrintError):
         return response['params'], response['result']
 
     def watch_channel(self, chan, callback):
-        script = funding_output_script(chan.local_config, chan.remote_config)
-        funding_address = redeem_script_to_address('p2wsh', script)
+        funding_address = funding_address_for_channel(chan)
         self.watched_channels[funding_address] = chan, callback
         self.network.subscribe_to_addresses([funding_address], self.on_address_status)
 
@@ -25,6 +32,8 @@ class LNWatcher(PrintError):
         if not params:
             return
         addr = params[0]
+        if self.address_status.get(addr) != result:
+            self.address_status[addr] = result
             self.network.request_address_utxos(addr, self.on_utxos)
 
     def on_utxos(self, response):
@@ -34,3 +43,192 @@ class LNWatcher(PrintError):
         addr = params[0]
         chan, callback = self.watched_channels[addr]
         callback(chan, result)
+
+
+def funding_address_for_channel(chan):
+    script = funding_output_script(chan.local_config, chan.remote_config)
+    return redeem_script_to_address('p2wsh', script)
+
+
+class LNChanCloseHandler(PrintError):
+
+    def __init__(self, network, wallet, chan):
+        self.network = network
+        self.wallet = wallet
+        self.chan = chan
+        self.funding_address = funding_address_for_channel(chan)
+        network.request_address_history(self.funding_address, self.on_history)
+
+    # TODO: de-duplicate?
+    def parse_response(self, response):
+        if response.get('error'):
+            self.print_error("response error:", response)
+            return None, None
+        return response['params'], response['result']
+
+    def on_history(self, response):
+        params, result = self.parse_response(response)
+        if not params:
+            return
+        addr = params[0]
+        if self.funding_address != addr:
+            self.print_error("unexpected funding address: {} != {}"
+                             .format(self.funding_address, addr))
+            return
+        txids = set(map(lambda item: item['tx_hash'], result))
+        self.network.get_transactions(txids, self.on_tx_response)
+
+    def on_tx_response(self, response):
+        params, result = self.parse_response(response)
+        if not params:
+            return
+        tx_hash = params[0]
+        tx = Transaction(result)
+        try:
+            tx.deserialize()
+        except Exception:
+            self.print_msg("cannot deserialize transaction", tx_hash)
+            return
+        if tx_hash != tx.txid():
+            self.print_error("received tx does not match expected txid ({} != {})"
+                             .format(tx_hash, tx.txid()))
+            return
+        funding_outpoint = self.chan.funding_outpoint
+        for i, txin in enumerate(tx.inputs()):
+            if txin['prevout_hash'] == funding_outpoint.txid \
+                    and txin['prevout_n'] == funding_outpoint.output_index:
+                self.print_error("funding outpoint {} is spent by {}"
+                                 .format(funding_outpoint, tx_hash))
+                self.inspect_spending_tx(tx, i)
+                break
+
+    # TODO batch sweeps
+    def inspect_spending_tx(self, ctx, txin_idx: int):
+        chan = self.chan
+        ctn = extract_ctn_from_tx(ctx, txin_idx,
+                                  chan.local_config.payment_basepoint.pubkey,
+                                  chan.remote_config.payment_basepoint.pubkey)
+        latest_local_ctn = chan.local_state.ctn
+        latest_remote_ctn = chan.remote_state.ctn
+        self.print_error("ctx {} has ctn {}. latest local ctn is {}, latest remote ctn is {}"
+                         .format(ctx.txid(), ctn, latest_local_ctn, latest_remote_ctn))
+        # see if it is a normal unilateral close by them
+        if ctn == latest_remote_ctn:
+            their_cur_pcp = chan.remote_state.current_per_commitment_point
+            self.find_and_sweep_their_ctx_to_remote(ctx, their_cur_pcp)
+        # see if we have a revoked secret for this ctn
+        try:
+            per_commitment_secret = chan.remote_state.revocation_store.retrieve_secret(
+                RevocationStore.START_INDEX - ctn)
+        except UnableToDeriveSecret:
+            self.print_error("revocation store does not have secret for ctx {}".format(ctx.txid()))
+        else:
+            # FIXME what if we closed unilaterally?
+            #self.print_error("ctx {} is breach!! by them and we have the revocation secret. "
+            #                 "yay, free money".format(ctx.txid()))
+            their_pcp = ecc.ECPrivkey(per_commitment_secret).get_public_key_bytes(compressed=True)
+            self.find_and_sweep_their_ctx_to_remote(ctx, their_pcp)
+            self.find_and_sweep_their_ctx_to_local(ctx, per_commitment_secret)
+            # TODO sweep other outputs
+
+    def find_and_sweep_their_ctx_to_remote(self, ctx, their_pcp: bytes):
+        payment_bp_privkey = ecc.ECPrivkey(self.chan.local_config.payment_basepoint.privkey)
+        our_payment_privkey = derive_privkey(payment_bp_privkey.secret_scalar, their_pcp)
+        our_payment_privkey = ecc.ECPrivkey.from_secret_scalar(our_payment_privkey)
+        our_payment_pubkey = our_payment_privkey.get_public_key_bytes(compressed=True)
+        to_remote_address = make_commitment_output_to_remote_address(our_payment_pubkey)
+        for output_idx, (type, addr, val) in enumerate(ctx.outputs()):
+            if type == TYPE_ADDRESS and addr == to_remote_address:
+                self.print_error("found to_remote output paying to us: ctx {}:{}".
+                                 format(ctx.txid(), output_idx))
+                #self.print_error("ctx {} is normal unilateral close by them".format(ctx.txid()))
+                break
+        else:
+            return
+        self.sweep_their_ctx_to_remote(ctx, output_idx, our_payment_privkey)
+
+    def sweep_their_ctx_to_remote(self, ctx, output_idx: int, our_payment_privkey: ecc.ECPrivkey):
+        our_payment_pubkey = our_payment_privkey.get_public_key_hex(compressed=True)
+        val = ctx.outputs()[output_idx][2]
+        sweep_inputs = [{
+            'type': 'p2wpkh',
+            'x_pubkeys': [our_payment_pubkey],
+            'num_sig': 1,
+            'prevout_n': output_idx,
+            'prevout_hash': ctx.txid(),
+            'value': val,
+            'coinbase': False,
+            'signatures': [None],
+        }]
+        tx_size_bytes = 110  # approx size of p2wpkh->p2wpkh
+        try:
+            fee = self.network.config.estimate_fee(tx_size_bytes)
+        except NoDynamicFeeEstimates:
+            fee_per_kb = self.network.config.fee_per_kb(dyn=False)
+            fee = self.network.config.estimate_fee_for_feerate(fee_per_kb, tx_size_bytes)
+        sweep_outputs = [(TYPE_ADDRESS, self.wallet.get_receiving_address(), val-fee)]
+        locktime = self.network.get_local_height()
+        sweep_tx = Transaction.from_io(sweep_inputs, sweep_outputs, locktime=locktime)
+        sweep_tx.set_rbf(True)
+        sweep_tx.sign({our_payment_pubkey: (our_payment_privkey.get_secret_bytes(), True)})
+        if not sweep_tx.is_complete():
+            raise Exception('channel close sweep tx is not complete')
+        self.network.broadcast_transaction(sweep_tx,
+                                           lambda res: self.print_tx_broadcast_result('sweep_their_ctx_to_remote', res))
+
+    def find_and_sweep_their_ctx_to_local(self, ctx, per_commitment_secret: bytes):
+        per_commitment_point = ecc.ECPrivkey(per_commitment_secret).get_public_key_bytes(compressed=True)
+        revocation_privkey = lnutil.derive_blinded_privkey(self.chan.local_config.revocation_basepoint.privkey,
+                                                           per_commitment_secret)
+        revocation_pubkey = ecc.ECPrivkey(revocation_privkey).get_public_key_bytes(compressed=True)
+        to_self_delay = self.chan.local_config.to_self_delay
+        delayed_pubkey = derive_pubkey(self.chan.remote_config.delayed_basepoint.pubkey,
+                                       per_commitment_point)
+        witness_script = bh2u(lnutil.make_commitment_output_to_local_witness_script(
+            revocation_pubkey, to_self_delay, delayed_pubkey))
+        to_local_address = redeem_script_to_address('p2wsh', witness_script)
+        for output_idx, (type, addr, val) in enumerate(ctx.outputs()):
+            if type == TYPE_ADDRESS and addr == to_local_address:
+                self.print_error("found to_local output paying to them: ctx {}:{}".
+                                 format(ctx.txid(), output_idx))
+                break
+        else:
+            self.print_error('could not find to_local output in their ctx {}'.format(ctx.txid()))
+            return
+        self.sweep_their_ctx_to_local(ctx, output_idx, witness_script, revocation_privkey)
+
+    def sweep_their_ctx_to_local(self, ctx, output_idx: int, witness_script: str, revocation_privkey: bytes):
+        val = ctx.outputs()[output_idx][2]
+        sweep_inputs = [{
+            'scriptSig': '',
+            'type': 'p2wsh',
+            'signatures': [],
+            'num_sig': 0,
+            'prevout_n': output_idx,
+            'prevout_hash': ctx.txid(),
+            'value': val,
+            'coinbase': False,
+            'preimage_script': witness_script,
+        }]
+        tx_size_bytes = 200  # TODO calc size
+        try:
+            fee = self.network.config.estimate_fee(tx_size_bytes)
+        except NoDynamicFeeEstimates:
+            fee_per_kb = self.network.config.fee_per_kb(dyn=False)
+            fee = self.network.config.estimate_fee_for_feerate(fee_per_kb, tx_size_bytes)
+        sweep_outputs = [(TYPE_ADDRESS, self.wallet.get_receiving_address(), val - fee)]
+        locktime = self.network.get_local_height()
+        sweep_tx = Transaction.from_io(sweep_inputs, sweep_outputs, locktime=locktime)
+        sweep_tx.set_rbf(True)
+        revocation_sig = sweep_tx.sign_txin(0, revocation_privkey)
+        witness = transaction.construct_witness([revocation_sig, 1, witness_script])
+        sweep_tx.inputs()[0]['witness'] = witness
+        self.network.broadcast_transaction(sweep_tx,
+                                           lambda res: self.print_tx_broadcast_result('sweep_their_ctx_to_local', res))
+
+    def print_tx_broadcast_result(self, name, res):
+        error = res.get('error')
+        if error:
+            self.print_error('{} broadcast failed: {}'.format(name, error))
+        else:
+            self.print_error('{} broadcast succeeded'.format(name))

lib/lnworker.py

@@ -16,6 +16,7 @@ from .ecc import der_sig_from_sig_string
 from .transaction import Transaction
 from .lnhtlc import HTLCStateMachine
 from .lnutil import Outpoint, calc_short_channel_id
+from .lnwatcher import LNChanCloseHandler
 
 # hardcoded nodes
 node_list = [
@@ -96,6 +97,8 @@ class LNWorker(PrintError):
         outpoints = [Outpoint(x["tx_hash"], x["tx_pos"]) for x in utxos]
         if chan.funding_outpoint not in outpoints:
             self.channel_state[chan.channel_id] = "CLOSED"
+            # FIXME is this properly GC-ed? (or too soon?)
+            LNChanCloseHandler(self.network, self.wallet, chan)
         elif self.channel_state[chan.channel_id] == 'DISCONNECTED':
             peer = self.peers[chan.node_id]
             coro = peer.reestablish_channel(chan)
@@ -125,6 +128,7 @@ class LNWorker(PrintError):
         peer = self.peers[node_id]
         openingchannel = await peer.channel_establishment_flow(self.wallet, self.config, password, amount_sat, push_sat * 1000, temp_channel_id=os.urandom(32))
         self.save_channel(openingchannel)
+        self.network.lnwatcher.watch_channel(openingchannel, self.on_channel_utxos)
         self.on_channels_updated()
 
     def on_channels_updated(self):

lib/tests/test_lnutil.py

@@ -1,7 +1,10 @@
 import unittest
 import json
 from lib import bitcoin
-from lib.lnutil import RevocationStore, get_per_commitment_secret_from_seed, make_offered_htlc, make_received_htlc, make_commitment, make_htlc_tx_witness, make_htlc_tx_output, make_htlc_tx_inputs, secret_to_pubkey, derive_blinded_pubkey, derive_privkey, derive_pubkey, make_htlc_tx
+from lib.lnutil import (RevocationStore, get_per_commitment_secret_from_seed, make_offered_htlc,
+                        make_received_htlc, make_commitment, make_htlc_tx_witness, make_htlc_tx_output,
+                        make_htlc_tx_inputs, secret_to_pubkey, derive_blinded_pubkey, derive_privkey,
+                        derive_pubkey, make_htlc_tx, extract_ctn_from_tx, UnableToDeriveSecret)
 from lib.util import bh2u, bfh
 from lib.transaction import Transaction
 
@@ -425,6 +428,12 @@ class TestLNUtil(unittest.TestCase):
                     if not insert["successful"]:
                         raise Exception("Failed ({}): error wasn't received".format(test["name"]))
 
+            for insert in test["inserts"]:
+                secret = bytes.fromhex(insert["secret"])
+                index = insert["index"]
+                if insert["successful"]:
+                    self.assertEqual(secret, receiver.retrieve_secret(index))
+
             print("Passed ({})".format(test["name"]))
 
     def test_shachain_produce_consume(self):
@@ -549,7 +558,7 @@ class TestLNUtil(unittest.TestCase):
 
         local_sig = our_htlc_tx.sign_txin(0, local_privkey[:-1])
 
-        our_htlc_tx_witness = make_htlc_tx_witness(  # FIXME only correct for success=True
+        our_htlc_tx_witness = make_htlc_tx_witness(
             remotehtlcsig=bfh(remote_htlc_sig) + b"\x01",  # 0x01 is SIGHASH_ALL
             localhtlcsig=bfh(local_sig),
             payment_preimage=htlc_payment_preimage if success else b'',  # will put 00 on witness if timeout
@@ -595,6 +604,11 @@ class TestLNUtil(unittest.TestCase):
 
         self.assertEqual(str(our_commit_tx), output_commit_tx)
 
+    def test_extract_commitment_number_from_tx(self):
+        raw_tx = "02000000000101bef67e4e2fb9ddeeb3461973cd4c62abb35050b1add772995b820b584a488489000000000038b02b8007e80300000000000022002052bfef0479d7b293c27e0f1eb294bea154c63a3294ef092c19af51409bce0e2ad007000000000000220020403d394747cae42e98ff01734ad5c08f82ba123d3d9a620abda88989651e2ab5d007000000000000220020748eba944fedc8827f6b06bc44678f93c0f9e6078b35c6331ed31e75f8ce0c2db80b000000000000220020c20b5d1f8584fd90443e7b7b720136174fa4b9333c261d04dbbd012635c0f419a00f0000000000002200208c48d15160397c9731df9bc3b236656efb6665fbfe92b4a6878e88a499f741c4c0c62d0000000000160014ccf1af2f2aabee14bb40fa3851ab2301de843110e0a06a00000000002200204adb4e2f00643db396dd120d4e7dc17625f5f2c11a40d857accc862d6b7dd80e04004730440220275b0c325a5e9355650dc30c0eccfbc7efb23987c24b556b9dfdd40effca18d202206caceb2c067836c51f296740c7ae807ffcbfbf1dd3a0d56b6de9a5b247985f060147304402204fd4928835db1ccdfc40f5c78ce9bd65249b16348df81f0c44328dcdefc97d630220194d3869c38bc732dd87d13d2958015e2fc16829e74cd4377f84d215c0b7060601475221023da092f6980e58d2c037173180e9a465476026ee50f96695963e8efe436f54eb21030e9f7b623d2ccc7c9bd44d66d5ce21ce504c0acf6385a132cec6d3c39fa711c152ae3e195220"
+        tx = Transaction(raw_tx)
+        self.assertEqual(commitment_number, extract_ctn_from_tx(tx, 0, local_payment_basepoint, remote_payment_basepoint))
+
     def test_per_commitment_secret_from_seed(self):
         self.assertEqual(0x02a40c85b6f28da08dfdbe0926c53fab2de6d28c10301f8f7c4073d5e42e3148.to_bytes(byteorder="big", length=32),
                          get_per_commitment_secret_from_seed(0x0000000000000000000000000000000000000000000000000000000000000000.to_bytes(byteorder="big", length=32), 281474976710655))