webvirtcloud/console/socketiod

#!/usr/bin/env python3
# The serial console daemon. The page connects with its signed console token
# (Socket.IO auth); each accepted connection gets its own console worker
# (conf/daemon/consolecallback) in a pty, which opens the VM's console with the
# compute's credentials. This process never opens libvirt. Connections are
# refused while SERIAL_CONSOLE_ENABLED is off.
import argparse
import codecs
import fcntl
import os
import pty
import select
import signal
import struct
import sys
import termios
import threading

import django

ROOT_PATH = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
if ROOT_PATH not in sys.path:
    sys.path.append(ROOT_PATH)
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "webvirtcloud.settings")
django.setup()

import socketio  # noqa: E402
from gunicorn.app.base import BaseApplication  # noqa: E402
from django import db  # noqa: E402
from django.conf import settings  # noqa: E402

from console.tokens import authorize  # noqa: E402

CONSOLECALLBACK = os.path.join(ROOT_PATH, "conf", "daemon", "consolecallback")
MAX_SESSIONS = 20

sio = socketio.Server(async_mode="threading")  # the same origin only
# Handlers run on their own threads. sessions (sid -> {"pid", "fd", "vm",
# "lock", "open"}) changes only under sessions_lock; a session's fd is used
# and closed only under its own lock, while it is open, so a closed fd's
# number, reused by a new session, never reaches another console.
sessions = {}
sessions_lock = threading.Lock()


@sio.event
def connect(sid, environ, auth):
    # the page shows the reason
    if not settings.SERIAL_CONSOLE_ENABLED:
        raise socketio.exceptions.ConnectionRefusedError("The serial console is turned off on this server")
    try:
        instance = authorize((auth or {}).get("token"))
    except PermissionError:
        raise socketio.exceptions.ConnectionRefusedError("The console link has expired: reload the page")
    finally:
        db.close_old_connections()
    with sessions_lock:
        if any(s["vm"] == instance.pk for s in sessions.values()):
            raise socketio.exceptions.ConnectionRefusedError("The console of this VM is open in another window")
        if len(sessions) >= MAX_SESSIONS:
            raise socketio.exceptions.ConnectionRefusedError("Too many consoles are open on this server: try again later")
        pid, fd = pty.fork()
        if pid == 0:
            try:
                # pty.fork's master is inheritable: without this each worker would
                # hold the other sessions' consoles open
                os.closerange(3, os.sysconf("SC_OPEN_MAX"))
                os.execv(sys.executable, [sys.executable, CONSOLECALLBACK, str(instance.pk)])
            finally:
                os._exit(1)
        sessions[sid] = session = {"pid": pid, "fd": fd, "vm": instance.pk, "lock": threading.Lock(), "open": True}
    sio.start_background_task(forward, sid, session)


def forward(sid, session):
    """The worker's output to its client, until the worker ends."""
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")  # a character may be split
    while session["open"]:
        try:
            # a fd closed meanwhile may wake this up: the session is closed then
            ready = select.select([session["fd"]], [], [], 0.1)[0]
        except (OSError, ValueError):
            ready = True
        if not ready:
            continue
        with session["lock"]:
            if not session["open"]:
                return
            try:
                output = os.read(session["fd"], 20480)
            except OSError:  # EIO: the worker ended
                output = b""
        if not output:
            sio.emit("pty_output", {"output": decoder.decode(b"", final=True) + "\r\n[console closed]\r\n"}, to=sid)
            stop(sid)
            sio.disconnect(sid)
            return
        sio.emit("pty_output", {"output": decoder.decode(output)}, to=sid)


def stop(sid):
    """End the sid's worker; once, whoever calls it first."""
    with sessions_lock:
        session = sessions.pop(sid, None)
    if session is None:
        return
    try:
        os.kill(session["pid"], signal.SIGKILL)  # first: a write blocked on its terminal ends
    except ProcessLookupError:
        pass
    with session["lock"]:
        session["open"] = False
        try:
            os.close(session["fd"])
        except OSError:
            pass
    sio.start_background_task(os.waitpid, session["pid"], 0)


@sio.event
def pty_input(sid, message):
    session = sessions.get(sid)
    if session and isinstance(message, dict) and isinstance(message.get("input"), str):
        data = message["input"].encode(errors="replace")  # a lone surrogate from the browser
        with session["lock"]:
            if session["open"]:
                try:
                    os.write(session["fd"], data)
                except OSError:
                    pass


@sio.event
def resize(sid, message):
    session = sessions.get(sid)
    if not (session and isinstance(message, dict)):
        return
    rows, cols = message.get("rows"), message.get("cols")
    if all(isinstance(v, int) and 1 <= v <= 500 for v in (rows, cols)):
        with session["lock"]:
            if session["open"]:
                fcntl.ioctl(session["fd"], termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))


@sio.event
def disconnect(sid, reason=None):
    stop(sid)


class Server(BaseApplication):
    """gunicorn with one worker (the sessions live in it); a thread per connection."""

    def __init__(self, host, port):
        self.bind = "%s:%d" % (host, port)
        super().__init__()

    def load_config(self):
        self.cfg.set("bind", self.bind)
        self.cfg.set("worker_class", "gthread")
        self.cfg.set("threads", 2 * MAX_SESSIONS)
        self.cfg.set("graceful_timeout", 1)  # open consoles never finish; their workers end with the pty
        self.cfg.set("control_socket_disable", True)

    def load(self):
        return socketio.WSGIApp(sio)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebVirtCloud serial console daemon")
    parser.add_argument("-H", "--host", default=settings.SOCKETIO_HOST, help="Listen host")
    parser.add_argument("-p", "--port", type=int, default=settings.SOCKETIO_PORT, help="Listen port")
    options = parser.parse_args()
    Server(options.host, options.port).run()