Inside Bitcoin

1# Copyright (c) 2022 Pieter Wuille
2# Distributed under the MIT software license, see the accompanying
3# file LICENSE or http://www.opensource.org/licenses/mit-license.php.
4
5"""
6This module provides the ASNEntry and ASMap classes.
7"""
8
9import copy
10import ipaddress
11import random
12import unittest
13from collections.abc import Callable, Iterable
14from enum import Enum
15from functools import total_ordering
16from typing import Optional, Union, overload
17
18def net_to_prefix(net: Union[ipaddress.IPv4Network,ipaddress.IPv6Network]) -> list[bool]:
19    """
20    Convert an IPv4 or IPv6 network to a prefix represented as a list of bits.
21
22    IPv4 ranges are remapped to their IPv4-mapped IPv6 range (::ffff:0:0/96).
23    """
24    num_bits = net.prefixlen
25    netrange = int.from_bytes(net.network_address.packed, 'big')
26
27    # Map an IPv4 prefix into IPv6 space.
28    if isinstance(net, ipaddress.IPv4Network):
29        num_bits += 96
30        netrange += 0xffff00000000
31
32    # Strip unused bottom bits.
33    assert (netrange & ((1 << (128 - num_bits)) - 1)) == 0
34    return [((netrange >> (127 - i)) & 1) != 0 for i in range(num_bits)]
35
36def prefix_to_net(prefix: list[bool]) -> Union[ipaddress.IPv4Network,ipaddress.IPv6Network]:
37    """The reverse operation of net_to_prefix."""
38    # Convert to number
39    netrange = sum(b << (127 - i) for i, b in enumerate(prefix))
40    num_bits = len(prefix)
41    assert num_bits <= 128
42
43    # Return IPv4 range if in ::ffff:0:0/96
44    if num_bits >= 96 and (netrange >> 32) == 0xffff:
45        return ipaddress.IPv4Network((netrange & 0xffffffff, num_bits - 96), True)
46
47    # Return IPv6 range otherwise.
48    return ipaddress.IPv6Network((netrange, num_bits), True)
49
50# Shortcut for (prefix, ASN) entries.
51ASNEntry = tuple[list[bool], int]
52
53# Shortcut for (prefix, old ASN, new ASN) entries.
54ASNDiff = tuple[list[bool], int, int]
55
56class _VarLenCoder:
57    """
58    A class representing a custom variable-length binary encoder/decoder for
59    integers. Each object represents a different coder, with different parameters
60    minval and clsbits.
61
62    The encoding is easiest to describe using an example. Let's say minval=100 and
63    clsbits=[4,2,2,3]. In that case:
64    - x in [100..115]: encoded as [0] + [4-bit BE encoding of (x-100)].
65    - x in [116..119]: encoded as [1,0] + [2-bit BE encoding of (x-116)].
66    - x in [120..123]: encoded as [1,1,0] + [2-bit BE encoding of (x-120)].
67    - x in [124..131]: encoded as [1,1,1] + [3-bit BE encoding of (x-124)].
68
69    In general, every number is encoded as:
70    - First, k "1"-bits, where k is the class the number falls in (there is one class
71      per element of clsbits).
72    - Then, a "0"-bit, unless k is the highest class, in which case there is nothing.
73    - Lastly, clsbits[k] bits encoding in big endian the position in its class that
74      number falls into.
75    - Every class k consists of 2^clsbits[k] consecutive integers. k=0 starts at minval,
76      other classes start one past the last element of the class before it.
77    """
78
79    def __init__(self, minval: int, clsbits: list[int]):
80        """Construct a new _VarLenCoder."""
81        self._minval = minval
82        self._clsbits = clsbits
83        self._maxval = minval + sum(1 << b for b in clsbits) - 1
84
85    def can_encode(self, val: int) -> bool:
86        """Check whether value val is in the range this coder supports."""
87        return self._minval <= val <= self._maxval
88
89    def encode(self, val: int, ret: list[int]) -> None:
90        """Append encoding of val onto integer list ret."""
91
92        assert self._minval <= val <= self._maxval
93        val -= self._minval
94        bits = 0
95        for k, bits in enumerate(self._clsbits):
96            if val >> bits:
97                # If the value will not fit in class k, subtract its range from v,
98                # emit a "1" bit and continue with the next class.
99                val -= 1 << bits
100                ret.append(1)
101            else:
102                if k + 1 < len(self._clsbits):
103                    # Unless we're in the last class, emit a "0" bit.
104                    ret.append(0)
105                break
106        # And then encode v (now the position within the class) in big endian.
107        ret.extend((val >> (bits - 1 - b)) & 1 for b in range(bits))
108
109    def encode_size(self, val: int) -> int:
110        """Compute how many bits are needed to encode val."""
111        assert self._minval <= val <= self._maxval
112        val -= self._minval
113        ret = 0
114        bits = 0
115        for k, bits in enumerate(self._clsbits):
116            if val >> bits:
117                val -= 1 << bits
118                ret += 1
119            else:
120                ret += k + 1 < len(self._clsbits)
121                break
122        return ret + bits
123
124    def decode(self, stream, bitpos) -> tuple[int,int]:
125        """Decode a number starting at bitpos in stream, returning value and new bitpos."""
126        val = self._minval
127        bits = 0
128        for k, bits in enumerate(self._clsbits):
129            bit = 0
130            if k + 1 < len(self._clsbits):
131                bit = stream[bitpos]
132                bitpos += 1
133            if not bit:
134                break
135            val += 1 << bits
136        for i in range(bits):
137            bit = stream[bitpos]
138            bitpos += 1
139            val += bit << (bits - 1 - i)
140        return val, bitpos
141
142# Variable-length encoders used in the binary asmap format.
143_CODER_INS = _VarLenCoder(0, [0, 0, 1])
144_CODER_ASN = _VarLenCoder(1, list(range(15, 25)))
145_CODER_MATCH = _VarLenCoder(2, list(range(1, 9)))
146_CODER_JUMP = _VarLenCoder(17, list(range(5, 31)))
147
148class _Instruction(Enum):
149    """One instruction in the binary asmap format."""
150    # A return instruction, encoded as [0], returns a constant ASN. It is followed by
151    # an integer using the ASN encoding.
152    RETURN = 0
153    # A jump instruction, encoded as [1,0] inspects the next unused bit in the input
154    # and either continues execution (if 0), or skips a specified number of bits (if 1).
155    # It is followed by an integer, and then two subprograms. The integer uses jump encoding
156    # and corresponds to the length of the first subprogram (so it can be skipped).
157    JUMP = 1
158    # A match instruction, encoded as [1,1,0] inspects 1 or more of the next unused bits
159    # in the input with its argument. If they all match, execution continues. If they do
160    # not, failure (represented by 0) is returned. If a default instruction has been executed before, instead
161    # of failure the default instruction's argument is returned. It is followed by an
162    # integer in match encoding, and a subprogram. That value is at least 2 bits and at
163    # most 9 bits. An n-bit value signifies matching (n-1) bits in the input with the lower
164    # (n-1) bits in the match value.
165    MATCH = 2
166    # A default instruction, encoded as [1,1,1] sets the default variable to its argument,
167    # and continues execution. It is followed by an integer in ASN encoding, and a subprogram.
168    DEFAULT = 3
169    # Not an actual instruction, but a way to encode the empty program that fails. In the
170    # encoder, it is used more generally to represent the failure case inside MATCH instructions,
171    # which may (if used inside the context of a DEFAULT instruction) actually correspond to
172    # a successful return. In this usage, they're always converted to an actual MATCH or RETURN
173    # before the top level is reached (see make_default below).
174    END = 4
175
176class _BinNode:
177    """A class representing a (node of) the parsed binary asmap format."""
178
179    @overload
180    def __init__(self, ins: _Instruction): ...
181    @overload
182    def __init__(self, ins: _Instruction, arg1: int): ...
183    @overload
184    def __init__(self, ins: _Instruction, arg1: "_BinNode", arg2: "_BinNode"): ...
185    @overload
186    def __init__(self, ins: _Instruction, arg1: int, arg2: "_BinNode"): ...
187
188    def __init__(self, ins: _Instruction, arg1=None, arg2=None):
189        """
190        Construct a new asmap node. Possibilities are:
191        - _BinNode(_Instruction.RETURN, asn)
192        - _BinNode(_Instruction.JUMP, node_0, node_1)
193        - _BinNode(_Instruction.MATCH, val, node)
194        - _BinNode(_Instruction.DEFAULT, asn, node)
195        - _BinNode(_Instruction.END)
196        """
197        self.ins = ins
198        self.arg1 = arg1
199        self.arg2 = arg2
200        if ins == _Instruction.RETURN:
201            assert isinstance(arg1, int)
202            assert arg2 is None
203            self.size = _CODER_INS.encode_size(ins.value) + _CODER_ASN.encode_size(arg1)
204        elif ins == _Instruction.JUMP:
205            assert isinstance(arg1, _BinNode)
206            assert isinstance(arg2, _BinNode)
207            self.size = (_CODER_INS.encode_size(ins.value) + _CODER_JUMP.encode_size(arg1.size) +
208                         arg1.size + arg2.size)
209        elif ins == _Instruction.DEFAULT:
210            assert isinstance(arg1, int)
211            assert isinstance(arg2, _BinNode)
212            self.size = _CODER_INS.encode_size(ins.value) + _CODER_ASN.encode_size(arg1) + arg2.size
213        elif ins == _Instruction.MATCH:
214            assert isinstance(arg1, int)
215            assert isinstance(arg2, _BinNode)
216            self.size = (_CODER_INS.encode_size(ins.value) + _CODER_MATCH.encode_size(arg1)
217                         + arg2.size)
218        elif ins == _Instruction.END:
219            assert arg1 is None
220            assert arg2 is None
221            self.size = 0
222        else:
223            assert False
224
225    @staticmethod
226    def make_end() -> "_BinNode":
227        """Constructor for a _BinNode with just an END instruction."""
228        return _BinNode(_Instruction.END)
229
230    @staticmethod
231    def make_leaf(val: int) -> "_BinNode":
232        """Constructor for a _BinNode of just a RETURN instruction."""
233        assert val is not None and val > 0
234        return _BinNode(_Instruction.RETURN, val)
235
236    @staticmethod
237    def make_branch(node0: "_BinNode", node1: "_BinNode") -> "_BinNode":
238        """
239        Construct a _BinNode corresponding to running either the node0 or node1 subprogram,
240        based on the next input bit. It exploits shortcuts that are possible in the encoding,
241        and uses either a JUMP, MATCH, or END instruction.
242        """
243        if node0.ins == _Instruction.END and node1.ins == _Instruction.END:
244            return node0
245        if node0.ins == _Instruction.END:
246            if node1.ins == _Instruction.MATCH and node1.arg1 <= 0xFF:
247                return _BinNode(node1.ins, node1.arg1 + (1 << node1.arg1.bit_length()), node1.arg2)
248            return _BinNode(_Instruction.MATCH, 3, node1)
249        if node1.ins == _Instruction.END:
250            if node0.ins == _Instruction.MATCH and node0.arg1 <= 0xFF: