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 import socket
  17 import threading
  18 from collections import deque
  19 import time
  20
  21
  22 class ConsoleSocket(socket.socket):
  23     """
  24     ConsoleSocket represents a socket attached to a char device.
  25
  26     Optionally (if drain==True), drains the socket and places the bytes
  27     into an in memory buffer for later processing.
  28
  29     Optionally a file path can be passed in and we will also
  30     dump the characters to this file for debugging purposes.
  31     """
  32     def __init__(self, address, file=None, drain=False):
  33         self._recv_timeout_sec = 300
  34         self._sleep_time = 0.5
  35         self._buffer = deque()
  36         socket.socket.__init__(self, socket.AF_UNIX, socket.SOCK_STREAM)
  37         self.connect(address)
  38         self._logfile = None
  39         if file:
  40             self._logfile = open(file, "w")
  41         self._open = True
  42         if drain:
  43             self._drain_thread = self._thread_start()
  44         else:
  45             self._drain_thread = None
  46
  47     def _drain_fn(self):
  48         """Drains the socket and runs while the socket is open."""
  49         while self._open:
  50             try:
  51                 self._drain_socket()
  52             except socket.timeout:
  53                 # The socket is expected to timeout since we set a
  54                 # short timeout to allow the thread to exit when
  55                 # self._open is set to False.
  56                 time.sleep(self._sleep_time)
  57
  58     def _thread_start(self):
  59         """Kick off a thread to drain the socket."""
  60         # Configure socket to not block and timeout.
  61         # This allows our drain thread to not block
  62         # on recieve and exit smoothly.
  63         socket.socket.setblocking(self, False)
  64         socket.socket.settimeout(self, 1)
  65         drain_thread = threading.Thread(target=self._drain_fn)
  66         drain_thread.daemon = True
  67         drain_thread.start()
  68         return drain_thread
  69
  70     def close(self):
  71         """Close the base object and wait for the thread to terminate"""
  72         if self._open:
  73             self._open = False
  74             if self._drain_thread is not None:
  75                 thread, self._drain_thread = self._drain_thread, None
  76                 thread.join()
  77             socket.socket.close(self)
  78             if self._logfile:
  79                 self._logfile.close()
  80                 self._logfile = None
  81
  82     def _drain_socket(self):
  83         """process arriving characters into in memory _buffer"""
  84         data = socket.socket.recv(self, 1)
  85         # latin1 is needed since there are some chars
  86         # we are receiving that cannot be encoded to utf-8
  87         # such as 0xe2, 0x80, 0xA6.
  88         string = data.decode("latin1")
  89         if self._logfile:
  90             self._logfile.write("{}".format(string))
  91             self._logfile.flush()
  92         for c in string:
  93             self._buffer.extend(c)
  94
  95     def recv(self, bufsize=1):
  96         """Return chars from in memory buffer.
  97            Maintains the same API as socket.socket.recv.
  98         """
  99         if self._drain_thread is None:
 100             # Not buffering the socket, pass thru to socket.
 101             return socket.socket.recv(self, bufsize)
 102         start_time = time.time()
 103         while len(self._buffer) < bufsize:
 104             time.sleep(self._sleep_time)
 105             elapsed_sec = time.time() - start_time
 106             if elapsed_sec > self._recv_timeout_sec:
 107                 raise socket.timeout
 108         chars = ''.join([self._buffer.popleft() for i in range(bufsize)])
 109         # We choose to use latin1 to remain consistent with
 110         # handle_read() and give back the same data as the user would
 111         # receive if they were reading directly from the
 112         # socket w/o our intervention.
 113         return chars.encode("latin1")
 114
 115     def setblocking(self, value):
 116         """When not draining we pass thru to the socket,
 117            since when draining we control socket blocking.
 118         """
 119         if self._drain_thread is None:
 120             socket.socket.setblocking(self, value)
 121
 122     def settimeout(self, seconds):
 123         """When not draining we pass thru to the socket,
 124            since when draining we control the timeout.
 125         """
 126         if seconds is not None:
 127             self._recv_timeout_sec = seconds
 128         if self._drain_thread is None:
 129             socket.socket.settimeout(self, seconds)