ideal/electrum/lnmsg.py

527 lines
24 KiB
Python

import os
import csv
import io
from typing import Callable, Tuple, Any, Dict, List, Sequence, Union, Optional
from collections import OrderedDict
from .lnutil import OnionFailureCodeMetaFlag
class FailedToParseMsg(Exception): pass
class UnknownMsgType(FailedToParseMsg): pass
class UnknownOptionalMsgType(UnknownMsgType): pass
class UnknownMandatoryMsgType(UnknownMsgType): pass
class MalformedMsg(FailedToParseMsg): pass
class UnknownMsgFieldType(MalformedMsg): pass
class UnexpectedEndOfStream(MalformedMsg): pass
class FieldEncodingNotMinimal(MalformedMsg): pass
class UnknownMandatoryTLVRecordType(MalformedMsg): pass
class MsgTrailingGarbage(MalformedMsg): pass
class MsgInvalidFieldOrder(MalformedMsg): pass
class UnexpectedFieldSizeForEncoder(MalformedMsg): pass
def _num_remaining_bytes_to_read(fd: io.BytesIO) -> int:
cur_pos = fd.tell()
end_pos = fd.seek(0, io.SEEK_END)
fd.seek(cur_pos)
return end_pos - cur_pos
def _assert_can_read_at_least_n_bytes(fd: io.BytesIO, n: int) -> None:
# note: it's faster to read n bytes and then check if we read n, than
# to assert we can read at least n and then read n bytes.
nremaining = _num_remaining_bytes_to_read(fd)
if nremaining < n:
raise UnexpectedEndOfStream(f"wants to read {n} bytes but only {nremaining} bytes left")
def write_bigsize_int(i: int) -> bytes:
assert i >= 0, i
if i < 0xfd:
return int.to_bytes(i, length=1, byteorder="big", signed=False)
elif i < 0x1_0000:
return b"\xfd" + int.to_bytes(i, length=2, byteorder="big", signed=False)
elif i < 0x1_0000_0000:
return b"\xfe" + int.to_bytes(i, length=4, byteorder="big", signed=False)
else:
return b"\xff" + int.to_bytes(i, length=8, byteorder="big", signed=False)
def read_bigsize_int(fd: io.BytesIO) -> Optional[int]:
try:
first = fd.read(1)[0]
except IndexError:
return None # end of file
if first < 0xfd:
return first
elif first == 0xfd:
buf = fd.read(2)
if len(buf) != 2:
raise UnexpectedEndOfStream()
val = int.from_bytes(buf, byteorder="big", signed=False)
if not (0xfd <= val < 0x1_0000):
raise FieldEncodingNotMinimal()
return val
elif first == 0xfe:
buf = fd.read(4)
if len(buf) != 4:
raise UnexpectedEndOfStream()
val = int.from_bytes(buf, byteorder="big", signed=False)
if not (0x1_0000 <= val < 0x1_0000_0000):
raise FieldEncodingNotMinimal()
return val
elif first == 0xff:
buf = fd.read(8)
if len(buf) != 8:
raise UnexpectedEndOfStream()
val = int.from_bytes(buf, byteorder="big", signed=False)
if not (0x1_0000_0000 <= val):
raise FieldEncodingNotMinimal()
return val
raise Exception()
# TODO: maybe if field_type is not "byte", we could return a list of type_len sized chunks?
# if field_type is a numeric, we could return a list of ints?
def _read_field(*, fd: io.BytesIO, field_type: str, count: Union[int, str]) -> Union[bytes, int]:
if not fd: raise Exception()
if isinstance(count, int):
assert count >= 0, f"{count!r} must be non-neg int"
elif count == "...":
pass
else:
raise Exception(f"unexpected field count: {count!r}")
if count == 0:
return b""
type_len = None
if field_type == 'byte':
type_len = 1
elif field_type in ('u8', 'u16', 'u32', 'u64'):
if field_type == 'u8':
type_len = 1
elif field_type == 'u16':
type_len = 2
elif field_type == 'u32':
type_len = 4
else:
assert field_type == 'u64'
type_len = 8
assert count == 1, count
buf = fd.read(type_len)
if len(buf) != type_len:
raise UnexpectedEndOfStream()
return int.from_bytes(buf, byteorder="big", signed=False)
elif field_type in ('tu16', 'tu32', 'tu64'):
if field_type == 'tu16':
type_len = 2
elif field_type == 'tu32':
type_len = 4
else:
assert field_type == 'tu64'
type_len = 8
assert count == 1, count
raw = fd.read(type_len)
if len(raw) > 0 and raw[0] == 0x00:
raise FieldEncodingNotMinimal()
return int.from_bytes(raw, byteorder="big", signed=False)
elif field_type == 'varint':
assert count == 1, count
val = read_bigsize_int(fd)
if val is None:
raise UnexpectedEndOfStream()
return val
elif field_type == 'chain_hash':
type_len = 32
elif field_type == 'channel_id':
type_len = 32
elif field_type == 'sha256':
type_len = 32
elif field_type == 'signature':
type_len = 64
elif field_type == 'point':
type_len = 33
elif field_type == 'short_channel_id':
type_len = 8
if count == "...":
total_len = -1 # read all
else:
if type_len is None:
raise UnknownMsgFieldType(f"unknown field type: {field_type!r}")
total_len = count * type_len
buf = fd.read(total_len)
if total_len >= 0 and len(buf) != total_len:
raise UnexpectedEndOfStream()
return buf
# TODO: maybe for "value" we could accept a list with len "count" of appropriate items
def _write_field(*, fd: io.BytesIO, field_type: str, count: Union[int, str],
value: Union[bytes, int]) -> None:
if not fd: raise Exception()
if isinstance(count, int):
assert count >= 0, f"{count!r} must be non-neg int"
elif count == "...":
pass
else:
raise Exception(f"unexpected field count: {count!r}")
if count == 0:
return
type_len = None
if field_type == 'byte':
type_len = 1
elif field_type == 'u8':
type_len = 1
elif field_type == 'u16':
type_len = 2
elif field_type == 'u32':
type_len = 4
elif field_type == 'u64':
type_len = 8
elif field_type in ('tu16', 'tu32', 'tu64'):
if field_type == 'tu16':
type_len = 2
elif field_type == 'tu32':
type_len = 4
else:
assert field_type == 'tu64'
type_len = 8
assert count == 1, count
if isinstance(value, int):
value = int.to_bytes(value, length=type_len, byteorder="big", signed=False)
if not isinstance(value, (bytes, bytearray)):
raise Exception(f"can only write bytes into fd. got: {value!r}")
while len(value) > 0 and value[0] == 0x00:
value = value[1:]
nbytes_written = fd.write(value)
if nbytes_written != len(value):
raise Exception(f"tried to write {len(value)} bytes, but only wrote {nbytes_written}!?")
return
elif field_type == 'varint':
assert count == 1, count
if isinstance(value, int):
value = write_bigsize_int(value)
if not isinstance(value, (bytes, bytearray)):
raise Exception(f"can only write bytes into fd. got: {value!r}")
nbytes_written = fd.write(value)
if nbytes_written != len(value):
raise Exception(f"tried to write {len(value)} bytes, but only wrote {nbytes_written}!?")
return
elif field_type == 'chain_hash':
type_len = 32
elif field_type == 'channel_id':
type_len = 32
elif field_type == 'sha256':
type_len = 32
elif field_type == 'signature':
type_len = 64
elif field_type == 'point':
type_len = 33
elif field_type == 'short_channel_id':
type_len = 8
total_len = -1
if count != "...":
if type_len is None:
raise UnknownMsgFieldType(f"unknown field type: {field_type!r}")
total_len = count * type_len
if isinstance(value, int) and (count == 1 or field_type == 'byte'):
value = int.to_bytes(value, length=total_len, byteorder="big", signed=False)
if not isinstance(value, (bytes, bytearray)):
raise Exception(f"can only write bytes into fd. got: {value!r}")
if count != "..." and total_len != len(value):
raise UnexpectedFieldSizeForEncoder(f"expected: {total_len}, got {len(value)}")
nbytes_written = fd.write(value)
if nbytes_written != len(value):
raise Exception(f"tried to write {len(value)} bytes, but only wrote {nbytes_written}!?")
def _read_tlv_record(*, fd: io.BytesIO) -> Tuple[int, bytes]:
if not fd: raise Exception()
tlv_type = _read_field(fd=fd, field_type="varint", count=1)
tlv_len = _read_field(fd=fd, field_type="varint", count=1)
tlv_val = _read_field(fd=fd, field_type="byte", count=tlv_len)
return tlv_type, tlv_val
def _write_tlv_record(*, fd: io.BytesIO, tlv_type: int, tlv_val: bytes) -> None:
if not fd: raise Exception()
tlv_len = len(tlv_val)
_write_field(fd=fd, field_type="varint", count=1, value=tlv_type)
_write_field(fd=fd, field_type="varint", count=1, value=tlv_len)
_write_field(fd=fd, field_type="byte", count=tlv_len, value=tlv_val)
def _resolve_field_count(field_count_str: str, *, vars_dict: dict, allow_any=False) -> Union[int, str]:
"""Returns an evaluated field count, typically an int.
If allow_any is True, the return value can be a str with value=="...".
"""
if field_count_str == "":
field_count = 1
elif field_count_str == "...":
if not allow_any:
raise Exception("field count is '...' but allow_any is False")
return field_count_str
else:
try:
field_count = int(field_count_str)
except ValueError:
field_count = vars_dict[field_count_str]
if isinstance(field_count, (bytes, bytearray)):
field_count = int.from_bytes(field_count, byteorder="big")
assert isinstance(field_count, int)
return field_count
def _parse_msgtype_intvalue_for_onion_wire(value: str) -> int:
msg_type_int = 0
for component in value.split("|"):
try:
msg_type_int |= int(component)
except ValueError:
msg_type_int |= OnionFailureCodeMetaFlag[component]
return msg_type_int
class LNSerializer:
def __init__(self, *, for_onion_wire: bool = False):
# TODO msg_type could be 'int' everywhere...
self.msg_scheme_from_type = {} # type: Dict[bytes, List[Sequence[str]]]
self.msg_type_from_name = {} # type: Dict[str, bytes]
self.in_tlv_stream_get_tlv_record_scheme_from_type = {} # type: Dict[str, Dict[int, List[Sequence[str]]]]
self.in_tlv_stream_get_record_type_from_name = {} # type: Dict[str, Dict[str, int]]
self.in_tlv_stream_get_record_name_from_type = {} # type: Dict[str, Dict[int, str]]
if for_onion_wire:
path = os.path.join(os.path.dirname(__file__), "lnwire", "onion_wire.csv")
else:
path = os.path.join(os.path.dirname(__file__), "lnwire", "peer_wire.csv")
with open(path, newline='') as f:
csvreader = csv.reader(f)
for row in csvreader:
#print(f">>> {row!r}")
if row[0] == "msgtype":
# msgtype,<msgname>,<value>[,<option>]
msg_type_name = row[1]
if for_onion_wire:
msg_type_int = _parse_msgtype_intvalue_for_onion_wire(str(row[2]))
else:
msg_type_int = int(row[2])
msg_type_bytes = msg_type_int.to_bytes(2, 'big')
assert msg_type_bytes not in self.msg_scheme_from_type, f"type collision? for {msg_type_name}"
assert msg_type_name not in self.msg_type_from_name, f"type collision? for {msg_type_name}"
row[2] = msg_type_int
self.msg_scheme_from_type[msg_type_bytes] = [tuple(row)]
self.msg_type_from_name[msg_type_name] = msg_type_bytes
elif row[0] == "msgdata":
# msgdata,<msgname>,<fieldname>,<typename>,[<count>][,<option>]
assert msg_type_name == row[1]
self.msg_scheme_from_type[msg_type_bytes].append(tuple(row))
elif row[0] == "tlvtype":
# tlvtype,<tlvstreamname>,<tlvname>,<value>[,<option>]
tlv_stream_name = row[1]
tlv_record_name = row[2]
tlv_record_type = int(row[3])
row[3] = tlv_record_type
if tlv_stream_name not in self.in_tlv_stream_get_tlv_record_scheme_from_type:
self.in_tlv_stream_get_tlv_record_scheme_from_type[tlv_stream_name] = OrderedDict()
self.in_tlv_stream_get_record_type_from_name[tlv_stream_name] = {}
self.in_tlv_stream_get_record_name_from_type[tlv_stream_name] = {}
assert tlv_record_type not in self.in_tlv_stream_get_tlv_record_scheme_from_type[tlv_stream_name], f"type collision? for {tlv_stream_name}/{tlv_record_name}"
assert tlv_record_name not in self.in_tlv_stream_get_record_type_from_name[tlv_stream_name], f"type collision? for {tlv_stream_name}/{tlv_record_name}"
assert tlv_record_type not in self.in_tlv_stream_get_record_type_from_name[tlv_stream_name], f"type collision? for {tlv_stream_name}/{tlv_record_name}"
self.in_tlv_stream_get_tlv_record_scheme_from_type[tlv_stream_name][tlv_record_type] = [tuple(row)]
self.in_tlv_stream_get_record_type_from_name[tlv_stream_name][tlv_record_name] = tlv_record_type
self.in_tlv_stream_get_record_name_from_type[tlv_stream_name][tlv_record_type] = tlv_record_name
if max(self.in_tlv_stream_get_tlv_record_scheme_from_type[tlv_stream_name].keys()) > tlv_record_type:
raise Exception(f"tlv record types must be listed in monotonically increasing order for stream. "
f"stream={tlv_stream_name}")
elif row[0] == "tlvdata":
# tlvdata,<tlvstreamname>,<tlvname>,<fieldname>,<typename>,[<count>][,<option>]
assert tlv_stream_name == row[1]
assert tlv_record_name == row[2]
self.in_tlv_stream_get_tlv_record_scheme_from_type[tlv_stream_name][tlv_record_type].append(tuple(row))
else:
pass # TODO
def write_tlv_stream(self, *, fd: io.BytesIO, tlv_stream_name: str, **kwargs) -> None:
scheme_map = self.in_tlv_stream_get_tlv_record_scheme_from_type[tlv_stream_name]
for tlv_record_type, scheme in scheme_map.items(): # note: tlv_record_type is monotonically increasing
tlv_record_name = self.in_tlv_stream_get_record_name_from_type[tlv_stream_name][tlv_record_type]
if tlv_record_name not in kwargs:
continue
with io.BytesIO() as tlv_record_fd:
for row in scheme:
if row[0] == "tlvtype":
pass
elif row[0] == "tlvdata":
# tlvdata,<tlvstreamname>,<tlvname>,<fieldname>,<typename>,[<count>][,<option>]
assert tlv_stream_name == row[1]
assert tlv_record_name == row[2]
field_name = row[3]
field_type = row[4]
field_count_str = row[5]
field_count = _resolve_field_count(field_count_str,
vars_dict=kwargs[tlv_record_name],
allow_any=True)
field_value = kwargs[tlv_record_name][field_name]
_write_field(fd=tlv_record_fd,
field_type=field_type,
count=field_count,
value=field_value)
else:
raise Exception(f"unexpected row in scheme: {row!r}")
_write_tlv_record(fd=fd, tlv_type=tlv_record_type, tlv_val=tlv_record_fd.getvalue())
def read_tlv_stream(self, *, fd: io.BytesIO, tlv_stream_name: str) -> Dict[str, Dict[str, Any]]:
parsed = {} # type: Dict[str, Dict[str, Any]]
scheme_map = self.in_tlv_stream_get_tlv_record_scheme_from_type[tlv_stream_name]
last_seen_tlv_record_type = -1 # type: int
while _num_remaining_bytes_to_read(fd) > 0:
tlv_record_type, tlv_record_val = _read_tlv_record(fd=fd)
if not (tlv_record_type > last_seen_tlv_record_type):
raise MsgInvalidFieldOrder(f"TLV records must be monotonically increasing by type. "
f"cur: {tlv_record_type}. prev: {last_seen_tlv_record_type}")
last_seen_tlv_record_type = tlv_record_type
try:
scheme = scheme_map[tlv_record_type]
except KeyError:
if tlv_record_type % 2 == 0:
# unknown "even" type: hard fail
raise UnknownMandatoryTLVRecordType(f"{tlv_stream_name}/{tlv_record_type}") from None
else:
# unknown "odd" type: skip it
continue
tlv_record_name = self.in_tlv_stream_get_record_name_from_type[tlv_stream_name][tlv_record_type]
parsed[tlv_record_name] = {}
with io.BytesIO(tlv_record_val) as tlv_record_fd:
for row in scheme:
#print(f"row: {row!r}")
if row[0] == "tlvtype":
pass
elif row[0] == "tlvdata":
# tlvdata,<tlvstreamname>,<tlvname>,<fieldname>,<typename>,[<count>][,<option>]
assert tlv_stream_name == row[1]
assert tlv_record_name == row[2]
field_name = row[3]
field_type = row[4]
field_count_str = row[5]
field_count = _resolve_field_count(field_count_str,
vars_dict=parsed[tlv_record_name],
allow_any=True)
#print(f">> count={field_count}. parsed={parsed}")
parsed[tlv_record_name][field_name] = _read_field(fd=tlv_record_fd,
field_type=field_type,
count=field_count)
else:
raise Exception(f"unexpected row in scheme: {row!r}")
if _num_remaining_bytes_to_read(tlv_record_fd) > 0:
raise MsgTrailingGarbage(f"TLV record ({tlv_stream_name}/{tlv_record_name}) has extra trailing garbage")
return parsed
def encode_msg(self, msg_type: str, **kwargs) -> bytes:
"""
Encode kwargs into a Lightning message (bytes)
of the type given in the msg_type string
"""
#print(f">>> encode_msg. msg_type={msg_type}, payload={kwargs!r}")
msg_type_bytes = self.msg_type_from_name[msg_type]
scheme = self.msg_scheme_from_type[msg_type_bytes]
with io.BytesIO() as fd:
fd.write(msg_type_bytes)
for row in scheme:
if row[0] == "msgtype":
pass
elif row[0] == "msgdata":
# msgdata,<msgname>,<fieldname>,<typename>,[<count>][,<option>]
field_name = row[2]
field_type = row[3]
field_count_str = row[4]
#print(f">>> encode_msg. msgdata. field_name={field_name!r}. field_type={field_type!r}. field_count_str={field_count_str!r}")
field_count = _resolve_field_count(field_count_str, vars_dict=kwargs)
if field_name == "tlvs":
tlv_stream_name = field_type
if tlv_stream_name in kwargs:
self.write_tlv_stream(fd=fd, tlv_stream_name=tlv_stream_name, **(kwargs[tlv_stream_name]))
continue
try:
field_value = kwargs[field_name]
except KeyError:
if len(row) > 5:
break # optional feature field not present
else:
field_value = 0 # default mandatory fields to zero
#print(f">>> encode_msg. writing field: {field_name}. value={field_value!r}. field_type={field_type!r}. count={field_count!r}")
_write_field(fd=fd,
field_type=field_type,
count=field_count,
value=field_value)
#print(f">>> encode_msg. so far: {fd.getvalue().hex()}")
else:
raise Exception(f"unexpected row in scheme: {row!r}")
return fd.getvalue()
def decode_msg(self, data: bytes) -> Tuple[str, dict]:
"""
Decode Lightning message by reading the first
two bytes to determine message type.
Returns message type string and parsed message contents dict,
or raises FailedToParseMsg.
"""
#print(f"decode_msg >>> {data.hex()}")
assert len(data) >= 2
msg_type_bytes = data[:2]
msg_type_int = int.from_bytes(msg_type_bytes, byteorder="big", signed=False)
try:
scheme = self.msg_scheme_from_type[msg_type_bytes]
except KeyError:
if msg_type_int % 2 == 0: # even types must be understood: "mandatory"
raise UnknownMandatoryMsgType(f"msg_type={msg_type_int}")
else: # odd types are ok not to understand: "optional"
raise UnknownOptionalMsgType(f"msg_type={msg_type_int}")
assert scheme[0][2] == msg_type_int
msg_type_name = scheme[0][1]
parsed = {}
with io.BytesIO(data[2:]) as fd:
for row in scheme:
#print(f"row: {row!r}")
if row[0] == "msgtype":
pass
elif row[0] == "msgdata":
field_name = row[2]
field_type = row[3]
field_count_str = row[4]
field_count = _resolve_field_count(field_count_str, vars_dict=parsed)
if field_name == "tlvs":
tlv_stream_name = field_type
d = self.read_tlv_stream(fd=fd, tlv_stream_name=tlv_stream_name)
parsed[tlv_stream_name] = d
continue
#print(f">> count={field_count}. parsed={parsed}")
try:
parsed[field_name] = _read_field(fd=fd,
field_type=field_type,
count=field_count)
except UnexpectedEndOfStream as e:
if len(row) > 5:
break # optional feature field not present
else:
raise
else:
raise Exception(f"unexpected row in scheme: {row!r}")
return msg_type_name, parsed
_inst = LNSerializer()
encode_msg = _inst.encode_msg
decode_msg = _inst.decode_msg
OnionWireSerializer = LNSerializer(for_onion_wire=True)