lnutil: RevocationStore: add type hints, asserts, sanity checks

electrum/lnpeer.py

@@ -1529,6 +1529,11 @@ async def on_channel_reestablish(self, chan: Channel, msg):
         # sanity checks of received values
         assert their_next_local_ctn >= 0  # already done by lnmsg, as type is u64
         assert their_oldest_unrevoked_remote_ctn >= 0
+        if max(their_next_local_ctn, their_oldest_unrevoked_remote_ctn) >= 2**48:
+            # TODO: upstream this check to lightning/bolts spec
+            self.logger.error(f"channel_reestablish ({chan.get_id_for_log()}): ctn overflow")
+            self.schedule_force_closing(chan.channel_id)
+            raise RemoteMisbehaving("channel_reestablish: ctn overflow")
         # ctns
         oldest_unrevoked_local_ctn = chan.get_oldest_unrevoked_ctn(LOCAL)
         latest_remote_ctn = chan.get_latest_ctn(REMOTE)

electrum/lnutil.py

@@ -3,7 +3,10 @@
 # file LICENCE or http://www.opensource.org/licenses/mit-license.php
 from enum import IntFlag, IntEnum
 import enum
-from typing import NamedTuple, List, Tuple, Mapping, Optional, TYPE_CHECKING, Union, Dict, Set, Sequence, FrozenSet
+from typing import (
+    NamedTuple, List, Tuple, Mapping, Optional, TYPE_CHECKING, Union, Dict, Set, Sequence, FrozenSet,
+    TypedDict,
+)
 import sys
 import time
 from functools import lru_cache
@@ -532,19 +535,24 @@ class LNProtocolWarning(Exception):
 TIME_FOR_OFFERED_HTLCS_TO_GET_FAILED_OFFCHAIN_ON_RESTART = 30
 
 
+class RevStoreStorage(TypedDict):
+    index: int
+    buckets: dict[int, 'ShachainElement']
+
 class RevocationStore:
     # closely based on code in lightningnetwork/lnd
 
-    START_INDEX = 2 ** 48 - 1
+    NUM_BITS = 48
+    START_INDEX = 2 ** NUM_BITS - 1
 
-    def __init__(self, storage):
+    def __init__(self, storage: RevStoreStorage):
         if len(storage) == 0:
             storage['index'] = self.START_INDEX
             storage['buckets'] = {}
         self.storage = storage
         self.buckets = storage['buckets']
 
-    def add_next_entry(self, hsh):
+    def add_next_entry(self, hsh: bytes) -> None:
         index = self.storage['index']
         new_element = ShachainElement(index=index, secret=hsh)
         bucket = count_trailing_zeros(index)
@@ -553,12 +561,15 @@ def add_next_entry(self, hsh):
             e = shachain_derive(new_element, this_bucket.index)
             if e != this_bucket:
                 raise Exception("hash is not derivable: {} {} {}".format(e.secret.hex(), this_bucket.secret.hex(), this_bucket.index))
+        # update state
+        new_index = index - 1
+        assert new_index > 3  # arbitrary small positive int. could not hurt to fail a bit early, before underflow
         self.buckets[bucket] = new_element
-        self.storage['index'] = index - 1
+        self.storage['index'] = new_index
 
     def retrieve_secret(self, index: int) -> bytes:
-        assert index <= self.START_INDEX, index
-        for i in range(0, 49):
+        assert 0 < index <= self.START_INDEX, index
+        for i in range(0, self.NUM_BITS + 1):
             bucket = self.buckets.get(i)
             if bucket is None:
                 raise UnableToDeriveSecret()
@@ -570,17 +581,24 @@ def retrieve_secret(self, index: int) -> bytes:
         raise UnableToDeriveSecret()
 
 
-def count_trailing_zeros(index):
+def count_trailing_zeros(index: int) -> int:
     """ BOLT-03 (where_to_put_secret) """
-    try:
-        return list(reversed(bin(index)[2:])).index("1")
-    except ValueError:
-        return 48
-
-
-def shachain_derive(element, to_index):
-    def get_prefix(index, pos):
-        mask = (1 << 64) - 1 - ((1 << pos) - 1)
+    assert isinstance(index, int)
+    assert 0 < index <= RevocationStore.START_INDEX, f"{index=}"
+    tz = list(reversed(bin(index)[2:])).index("1")
+    assert 0 <= tz < RevocationStore.NUM_BITS
+    return tz
+
+
+def shachain_derive(element: 'ShachainElement', to_index: int) -> 'ShachainElement':
+    assert isinstance(to_index, int)
+    assert 0 < to_index <= RevocationStore.START_INDEX, f"{to_index=}"
+    def get_prefix(index: int, pos: int) -> int:
+        assert isinstance(index, int)
+        assert isinstance(pos, int)
+        max_mask_len = 64  # TODO just use RevocationStore.NUM_BITS + 1 ?
+        assert max_mask_len > RevocationStore.NUM_BITS
+        mask = (1 << max_mask_len) - 1 - ((1 << pos) - 1)
         return index & mask
     from_index = element.index
     zeros = count_trailing_zeros(from_index)
@@ -603,8 +621,12 @@ def read(*x):
         return ShachainElement(bfh(x[0]), int(x[1]))
 
 
-def get_per_commitment_secret_from_seed(seed: bytes, i: int, bits: int = 48) -> bytes:
+def get_per_commitment_secret_from_seed(seed: bytes, i: int, bits: int = RevocationStore.NUM_BITS) -> bytes:
     """Generate per commitment secret."""
+    assert isinstance(seed, bytes) and len(seed) == 32
+    assert isinstance(bits, int) and (0 <= bits <= RevocationStore.NUM_BITS)
+    assert isinstance(i, int)
+    assert 0 < i <= RevocationStore.START_INDEX, f"{i=}"
     per_commitment_secret = bytearray(seed)
     for bitindex in range(bits - 1, -1, -1):
         mask = 1 << bitindex

tests/test_lnpeer.py

@@ -253,7 +253,7 @@ def prepare_peers(
         w1, w2 = graph.workers.values()
         return p1, p2, w1, w2
 
-    async def test_reestablish(self):
+    async def test_reestablish_happycase(self):
         graph = self.prepare_chans_and_peers_in_graph(self.GRAPH_DEFINITIONS['single_chan'])
         p1, p2 = graph.peers.values()
         alice_channel = graph.channels[('alice', 'bob')][0]
@@ -267,6 +267,7 @@ async def reestablish():
                 p2.reestablish_channel(bob_channel))
             self.assertEqual(alice_channel.peer_state, PeerState.GOOD)
             self.assertEqual(bob_channel.peer_state, PeerState.GOOD)
+            self.assertEqual((alice_channel._state, bob_channel._state), (ChannelState.OPEN, ChannelState.OPEN))
             gath.cancel()
         gath = asyncio.gather(reestablish(), p1._message_loop(), p2._message_loop(), p1.htlc_switch(), p2.htlc_switch())
         with self.assertRaises(asyncio.CancelledError):
@@ -354,10 +355,11 @@ async def alice_sends_reest():
                 oldest_unrevoked_remote_ctn = chan.get_oldest_unrevoked_ctn(REMOTE) + revnum_delta
                 assert oldest_unrevoked_remote_ctn >= 0, oldest_unrevoked_remote_ctn
                 if last_rev_secret is None:
+                    revnum_for_secret = oldest_unrevoked_remote_ctn % (2**48)
                     if revnum_delta <= 0:
-                        last_rev_secret = chan.revocation_store.retrieve_secret(RevocationStore.START_INDEX - oldest_unrevoked_remote_ctn + 1)
+                        last_rev_secret = chan.revocation_store.retrieve_secret(RevocationStore.START_INDEX - revnum_for_secret + 1)
                     else:  # Alice is using *magic* here, i.e. cheating: she uses Bob's channel to learn future unrevealed secrets
-                        last_rev_secret, _point = bob_channel.get_secret_and_point(LOCAL, oldest_unrevoked_remote_ctn - 1)
+                        last_rev_secret, _point = bob_channel.get_secret_and_point(LOCAL, revnum_for_secret - 1)
                 p1.send_message(
                     "channel_reestablish",
                     channel_id=chan.channel_id,
@@ -411,6 +413,18 @@ async def exit_after_bob_receives_reest():
             with self.subTest(msg="invalid last_rev_secret", **kwargs):
                 a_chan, b_chan = await f(last_rev_secret=sha256("fake_data"), **kwargs)
                 self.assertEqual((a_chan._state, b_chan._state), (cs.OPEN, cs.FORCE_CLOSING))
+            with self.subTest(msg="overflow of next_local_ctn", **kwargs):
+                with self.assertLogs('electrum', level='INFO') as logs:
+                    a_chan, b_chan = await f(ctn_delta=2**48, **kwargs)
+                self.assertEqual((a_chan._state, b_chan._state), (cs.OPEN, cs.FORCE_CLOSING))
+                self.assertTrue(any(("bob->alice" in msg and "channel_reestablish" in msg and "ctn overflow" in msg)
+                                    for msg in logs.output))
+            with self.subTest(msg="overflow of oldest_unrevoked_remote_ctn", **kwargs):
+                with self.assertLogs('electrum', level='INFO') as logs:
+                    a_chan, b_chan = await f(revnum_delta=2**48, **kwargs)
+                self.assertEqual((a_chan._state, b_chan._state), (cs.OPEN, cs.FORCE_CLOSING))
+                self.assertTrue(any(("bob->alice" in msg and "channel_reestablish" in msg and "ctn overflow" in msg)
+                                    for msg in logs.output))
 
     @staticmethod
     def _send_fake_htlc(peer: Peer, chan: Channel) -> UpdateAddHtlc: