]>
Commit | Line | Data |
---|---|---|
4b84d874 RF |
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 | """ | |
0fc8f660 RF |
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 | # | |
4b84d874 | 15 | |
932ca4bb | 16 | from collections import deque |
0fc8f660 RF |
17 | import socket |
18 | import threading | |
0fc8f660 | 19 | import time |
e35c1382 | 20 | from typing import Deque, Optional |
4b84d874 | 21 | |
0fc8f660 | 22 | |
80ded8e9 | 23 | class ConsoleSocket(socket.socket): |
4b84d874 RF |
24 | """ |
25 | ConsoleSocket represents a socket attached to a char device. | |
0fc8f660 | 26 | |
80ded8e9 RF |
27 | Optionally (if drain==True), drains the socket and places the bytes |
28 | into an in memory buffer for later processing. | |
4b84d874 RF |
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 | """ | |
e35c1382 JS |
33 | def __init__(self, address: str, file: Optional[str] = None, |
34 | drain: bool = False): | |
6cf4cce7 | 35 | self._recv_timeout_sec = 300.0 |
4b84d874 | 36 | self._sleep_time = 0.5 |
af0db882 | 37 | self._buffer: Deque[int] = deque() |
80ded8e9 RF |
38 | socket.socket.__init__(self, socket.AF_UNIX, socket.SOCK_STREAM) |
39 | self.connect(address) | |
0fc8f660 RF |
40 | self._logfile = None |
41 | if file: | |
8825fed8 | 42 | # pylint: disable=consider-using-with |
af0db882 | 43 | self._logfile = open(file, "bw") |
0fc8f660 | 44 | self._open = True |
714ac05a | 45 | self._drain_thread = None |
80ded8e9 RF |
46 | if drain: |
47 | self._drain_thread = self._thread_start() | |
0fc8f660 | 48 | |
afded359 | 49 | def __repr__(self) -> str: |
ee1a2723 JS |
50 | tmp = super().__repr__() |
51 | tmp = tmp.rstrip(">") | |
52 | tmp = "%s, logfile=%s, drain_thread=%s>" % (tmp, self._logfile, | |
53 | self._drain_thread) | |
54 | return tmp | |
afded359 | 55 | |
e35c1382 | 56 | def _drain_fn(self) -> None: |
80ded8e9 RF |
57 | """Drains the socket and runs while the socket is open.""" |
58 | while self._open: | |
59 | try: | |
60 | self._drain_socket() | |
61 | except socket.timeout: | |
62 | # The socket is expected to timeout since we set a | |
63 | # short timeout to allow the thread to exit when | |
64 | # self._open is set to False. | |
65 | time.sleep(self._sleep_time) | |
0fc8f660 | 66 | |
e35c1382 | 67 | def _thread_start(self) -> threading.Thread: |
80ded8e9 RF |
68 | """Kick off a thread to drain the socket.""" |
69 | # Configure socket to not block and timeout. | |
70 | # This allows our drain thread to not block | |
71 | # on recieve and exit smoothly. | |
72 | socket.socket.setblocking(self, False) | |
73 | socket.socket.settimeout(self, 1) | |
74 | drain_thread = threading.Thread(target=self._drain_fn) | |
75 | drain_thread.daemon = True | |
76 | drain_thread.start() | |
77 | return drain_thread | |
0fc8f660 | 78 | |
e35c1382 | 79 | def close(self) -> None: |
0fc8f660 RF |
80 | """Close the base object and wait for the thread to terminate""" |
81 | if self._open: | |
82 | self._open = False | |
80ded8e9 RF |
83 | if self._drain_thread is not None: |
84 | thread, self._drain_thread = self._drain_thread, None | |
0fc8f660 | 85 | thread.join() |
80ded8e9 | 86 | socket.socket.close(self) |
0fc8f660 RF |
87 | if self._logfile: |
88 | self._logfile.close() | |
89 | self._logfile = None | |
90 | ||
e35c1382 | 91 | def _drain_socket(self) -> None: |
0fc8f660 | 92 | """process arriving characters into in memory _buffer""" |
80ded8e9 | 93 | data = socket.socket.recv(self, 1) |
0fc8f660 | 94 | if self._logfile: |
af0db882 | 95 | self._logfile.write(data) |
0fc8f660 | 96 | self._logfile.flush() |
af0db882 | 97 | self._buffer.extend(data) |
0fc8f660 | 98 | |
ff3513e6 | 99 | def recv(self, bufsize: int = 1, flags: int = 0) -> bytes: |
4b84d874 RF |
100 | """Return chars from in memory buffer. |
101 | Maintains the same API as socket.socket.recv. | |
102 | """ | |
80ded8e9 RF |
103 | if self._drain_thread is None: |
104 | # Not buffering the socket, pass thru to socket. | |
ff3513e6 JS |
105 | return socket.socket.recv(self, bufsize, flags) |
106 | assert not flags, "Cannot pass flags to recv() in drained mode" | |
0fc8f660 | 107 | start_time = time.time() |
80ded8e9 | 108 | while len(self._buffer) < bufsize: |
4b84d874 | 109 | time.sleep(self._sleep_time) |
0fc8f660 RF |
110 | elapsed_sec = time.time() - start_time |
111 | if elapsed_sec > self._recv_timeout_sec: | |
112 | raise socket.timeout | |
af0db882 | 113 | return bytes((self._buffer.popleft() for i in range(bufsize))) |
0fc8f660 | 114 | |
e35c1382 | 115 | def setblocking(self, value: bool) -> None: |
80ded8e9 RF |
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) | |
0fc8f660 | 121 | |
6cf4cce7 | 122 | def settimeout(self, value: Optional[float]) -> None: |
80ded8e9 RF |
123 | """When not draining we pass thru to the socket, |
124 | since when draining we control the timeout. | |
125 | """ | |
6cf4cce7 JS |
126 | if value is not None: |
127 | self._recv_timeout_sec = value | |
80ded8e9 | 128 | if self._drain_thread is None: |
6cf4cce7 | 129 | socket.socket.settimeout(self, value) |