cln-plugins/backup/server.py

import logging
import socket
import struct
import json
import sys
from typing import Tuple

from backend import Backend
from protocol import (
    PacketType,
    PKT_CHANGE_TYPES,
    change_from_packet,
    packet_from_change,
    send_packet,
    recv_packet,
)


class SystemdHandler(logging.Handler):
    PREFIX = {
        # EMERG <0>
        # ALERT <1>
        logging.CRITICAL: "<2>",
        logging.ERROR: "<3>",
        logging.WARNING: "<4>",
        # NOTICE <5>
        logging.INFO: "<6>",
        logging.DEBUG: "<7>",
        logging.NOTSET: "<7>",
    }

    def __init__(self, stream=sys.stdout):
        self.stream = stream
        logging.Handler.__init__(self)

    def emit(self, record):
        try:
            msg = self.PREFIX[record.levelno] + self.format(record) + "\n"
            self.stream.write(msg)
            self.stream.flush()
        except Exception:
            self.handleError(record)


def setup_server_logging(mode, level):
    root_logger = logging.getLogger()
    root_logger.setLevel(level.upper())
    mode = mode.lower()
    if mode == "systemd":
        # replace handler with systemd one
        root_logger.handlers = []
        root_logger.addHandler(SystemdHandler())
    else:
        assert mode == "plain"


class SocketServer:
    def __init__(self, addr: Tuple[str, int], backend: Backend) -> None:
        self.backend = backend
        self.addr = addr
        self.bind = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.bind.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.bind.bind(addr)

    def _send_packet(self, typ: int, payload: bytes) -> None:
        send_packet(self.sock, typ, payload)

    def _recv_packet(self) -> Tuple[int, bytes]:
        return recv_packet(self.sock)

    def _handle_conn(self, conn) -> None:
        # Can only handle one connection at a time
        logging.info("Servicing incoming connection")
        self.sock = conn
        while True:
            try:
                (typ, payload) = self._recv_packet()
            except IOError:
                logging.info("Connection closed")
                break
            if typ in PKT_CHANGE_TYPES:
                change = change_from_packet(typ, payload)
                if typ == PacketType.CHANGE:
                    logging.debug("Received CHANGE {}".format(change.version))
                else:
                    logging.info("Received SNAPSHOT {}".format(change.version))
                self.backend.add_change(change)
                self._send_packet(
                    PacketType.ACK, struct.pack("!I", self.backend.version)
                )
            elif typ == PacketType.REWIND:
                logging.info("Received REWIND")
                (to_version,) = struct.unpack("!I", payload)
                if to_version != self.backend.prev_version:
                    logging.info("Cannot rewind to version {}".format(to_version))
                    self._send_packet(
                        PacketType.NACK, struct.pack("!I", self.backend.version)
                    )
                else:
                    self.backend.rewind()
                    self._send_packet(
                        PacketType.ACK, struct.pack("!I", self.backend.version)
                    )
            elif typ == PacketType.REQ_METADATA:
                logging.debug("Received REQ_METADATA")
                blob = struct.pack(
                    "!IIIQ",
                    0x01,
                    self.backend.version,
                    self.backend.prev_version,
                    self.backend.version_count,
                )
                self._send_packet(PacketType.METADATA, blob)
            elif typ == PacketType.RESTORE:
                logging.info("Received RESTORE")
                for change in self.backend.stream_changes():
                    (typ, payload) = packet_from_change(change)
                    self._send_packet(typ, payload)
                self._send_packet(PacketType.DONE, b"")
            elif typ == PacketType.COMPACT:
                logging.info("Received COMPACT")
                stats = self.backend.compact()
                self._send_packet(PacketType.COMPACT_RES, json.dumps(stats).encode())
            elif typ == PacketType.ACK:
                logging.debug("Received ACK")
            elif typ == PacketType.NACK:
                logging.debug("Received NACK")
            elif typ == PacketType.METADATA:
                logging.debug("Received METADATA")
            elif typ == PacketType.COMPACT_RES:
                logging.debug("Received COMPACT_RES")
            else:
                raise Exception("Unknown or unexpected packet type {}".format(typ))
        self.conn = None

    def run(self) -> None:
        self.bind.listen(1)
        logging.info("Waiting for connection on {}".format(self.addr))
        while True:
            conn, _ = self.bind.accept()
            try:
                self._handle_conn(conn)
            except Exception:
                logging.exception("Got exception")
            finally:
                conn.close()