python/qemu/console_socket.py

   1 """
   2 QEMU Console Socket Module:
   3
   4 This python module implements a ConsoleSocket object,
   5 which can drain a socket and optionally dump the bytes to file.
   6 """
   7 # Copyright 2020 Linaro
   8 #
   9 # Authors:
  10 #  Robert Foley <robert.foley@linaro.org>
  11 #
  12 # This code is licensed under the GPL version 2 or later.  See
  13 # the COPYING file in the top-level directory.
  14 #
  15
  16 from collections import deque
  17 import socket
  18 import threading
  19 import time
  20 from typing import Deque, Optional
  21
  22
  23 class ConsoleSocket(socket.socket):
  24     """
  25     ConsoleSocket represents a socket attached to a char device.
  26
  27     Optionally (if drain==True), drains the socket and places the bytes
  28     into an in memory buffer for later processing.
  29
  30     Optionally a file path can be passed in and we will also
  31     dump the characters to this file for debugging purposes.
  32     """
  33     def __init__(self, address: str, file: Optional[str] = None,
  34                  drain: bool = False):
  35         self._recv_timeout_sec = 300.0
  36         self._sleep_time = 0.5
  37         self._buffer: Deque[int] = deque()
  38         socket.socket.__init__(self, socket.AF_UNIX, socket.SOCK_STREAM)
  39         self.connect(address)
  40         self._logfile = None
  41         if file:
  42             self._logfile = open(file, "bw")
  43         self._open = True
  44         self._drain_thread = None
  45         if drain:
  46             self._drain_thread = self._thread_start()
  47
  48     def __repr__(self) -> str:
  49         s = super().__repr__()
  50         s = s.rstrip(">")
  51         s = "%s,  logfile=%s, drain_thread=%s>" % (s, self._logfile,
  52                                                    self._drain_thread)
  53         return s
  54
  55     def _drain_fn(self) -> None:
  56         """Drains the socket and runs while the socket is open."""
  57         while self._open:
  58             try:
  59                 self._drain_socket()
  60             except socket.timeout:
  61                 # The socket is expected to timeout since we set a
  62                 # short timeout to allow the thread to exit when
  63                 # self._open is set to False.
  64                 time.sleep(self._sleep_time)
  65
  66     def _thread_start(self) -> threading.Thread:
  67         """Kick off a thread to drain the socket."""
  68         # Configure socket to not block and timeout.
  69         # This allows our drain thread to not block
  70         # on recieve and exit smoothly.
  71         socket.socket.setblocking(self, False)
  72         socket.socket.settimeout(self, 1)
  73         drain_thread = threading.Thread(target=self._drain_fn)
  74         drain_thread.daemon = True
  75         drain_thread.start()
  76         return drain_thread
  77
  78     def close(self) -> None:
  79         """Close the base object and wait for the thread to terminate"""
  80         if self._open:
  81             self._open = False
  82             if self._drain_thread is not None:
  83                 thread, self._drain_thread = self._drain_thread, None
  84                 thread.join()
  85             socket.socket.close(self)
  86             if self._logfile:
  87                 self._logfile.close()
  88                 self._logfile = None
  89
  90     def _drain_socket(self) -> None:
  91         """process arriving characters into in memory _buffer"""
  92         data = socket.socket.recv(self, 1)
  93         if self._logfile:
  94             self._logfile.write(data)
  95             self._logfile.flush()
  96         self._buffer.extend(data)
  97
  98     def recv(self, bufsize: int = 1, flags: int = 0) -> bytes:
  99         """Return chars from in memory buffer.
 100            Maintains the same API as socket.socket.recv.
 101         """
 102         if self._drain_thread is None:
 103             # Not buffering the socket, pass thru to socket.
 104             return socket.socket.recv(self, bufsize, flags)
 105         assert not flags, "Cannot pass flags to recv() in drained mode"
 106         start_time = time.time()
 107         while len(self._buffer) < bufsize:
 108             time.sleep(self._sleep_time)
 109             elapsed_sec = time.time() - start_time
 110             if elapsed_sec > self._recv_timeout_sec:
 111                 raise socket.timeout
 112         return bytes((self._buffer.popleft() for i in range(bufsize)))
 113
 114     def setblocking(self, value: bool) -> None:
 115         """When not draining we pass thru to the socket,
 116            since when draining we control socket blocking.
 117         """
 118         if self._drain_thread is None:
 119             socket.socket.setblocking(self, value)
 120
 121     def settimeout(self, value: Optional[float]) -> None:
 122         """When not draining we pass thru to the socket,
 123            since when draining we control the timeout.
 124         """
 125         if value is not None:
 126             self._recv_timeout_sec = value
 127         if self._drain_thread is None:
 128             socket.socket.settimeout(self, value)