cln-plugins/backup/protocol.py
2024-07-10 12:13:42 -05:00

81 lines
2.2 KiB
Python

"""
Socket-based remote backup protocol. This is used to create a connection to a backup backend, and send it incremental database updates.
"""
import socket
import struct
from typing import Tuple
import zlib
from backend import Change
class PacketType:
CHANGE = 0x01
SNAPSHOT = 0x02
REWIND = 0x03
REQ_METADATA = 0x04
RESTORE = 0x05
ACK = 0x06
NACK = 0x07
METADATA = 0x08
DONE = 0x09
COMPACT = 0x0A
COMPACT_RES = 0x0B
PKT_CHANGE_TYPES = {PacketType.CHANGE, PacketType.SNAPSHOT}
def recvall(sock: socket.socket, n: int) -> bytearray:
"""Receive exactly n bytes from a socket."""
buf = bytearray(n)
view = memoryview(buf)
ptr = 0
while ptr < n:
count = sock.recv_into(view[ptr:])
if count == 0:
raise IOError("Premature end of stream")
ptr += count
return buf
def send_packet(sock: socket.socket, typ: int, payload: bytes) -> None:
sock.sendall(struct.pack("!BI", typ, len(payload)))
sock.sendall(payload)
def recv_packet(sock: socket.socket) -> Tuple[int, bytes]:
(typ, length) = struct.unpack("!BI", recvall(sock, 5))
payload = recvall(sock, length)
return (typ, payload)
def change_from_packet(typ, payload):
"""Convert a network packet to a Change object."""
if typ == PacketType.CHANGE:
(version,) = struct.unpack("!I", payload[0:4])
payload = zlib.decompress(payload[4:])
return Change(
version=version,
snapshot=None,
transaction=[t.decode("UTF-8") for t in payload.split(b"\x00")],
)
elif typ == PacketType.SNAPSHOT:
(version,) = struct.unpack("!I", payload[0:4])
payload = zlib.decompress(payload[4:])
return Change(version=version, snapshot=payload, transaction=None)
raise ValueError("Not a change (typ {})".format(typ))
def packet_from_change(entry):
"""Convert a Change object to a network packet."""
if entry.snapshot is None:
typ = PacketType.CHANGE
payload = b"\x00".join([t.encode("UTF-8") for t in entry.transaction])
else:
typ = PacketType.SNAPSHOT
payload = entry.snapshot
version = struct.pack("!I", entry.version)
return typ, version + zlib.compress(payload)