]>
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 |
6cf4cce7 | 20 | from typing import 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 | """ | |
80ded8e9 | 33 | def __init__(self, address, file=None, drain=False): |
0fc8f660 | 34 | self._recv_timeout_sec = 300 |
6cf4cce7 | 35 | self._recv_timeout_sec = 300.0 |
4b84d874 | 36 | self._sleep_time = 0.5 |
0fc8f660 | 37 | self._buffer = 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: | |
42 | self._logfile = open(file, "w") | |
0fc8f660 | 43 | self._open = True |
80ded8e9 RF |
44 | if drain: |
45 | self._drain_thread = self._thread_start() | |
46 | else: | |
47 | self._drain_thread = None | |
0fc8f660 | 48 | |
80ded8e9 RF |
49 | def _drain_fn(self): |
50 | """Drains the socket and runs while the socket is open.""" | |
51 | while self._open: | |
52 | try: | |
53 | self._drain_socket() | |
54 | except socket.timeout: | |
55 | # The socket is expected to timeout since we set a | |
56 | # short timeout to allow the thread to exit when | |
57 | # self._open is set to False. | |
58 | time.sleep(self._sleep_time) | |
0fc8f660 | 59 | |
80ded8e9 RF |
60 | def _thread_start(self): |
61 | """Kick off a thread to drain the socket.""" | |
62 | # Configure socket to not block and timeout. | |
63 | # This allows our drain thread to not block | |
64 | # on recieve and exit smoothly. | |
65 | socket.socket.setblocking(self, False) | |
66 | socket.socket.settimeout(self, 1) | |
67 | drain_thread = threading.Thread(target=self._drain_fn) | |
68 | drain_thread.daemon = True | |
69 | drain_thread.start() | |
70 | return drain_thread | |
0fc8f660 RF |
71 | |
72 | def close(self): | |
73 | """Close the base object and wait for the thread to terminate""" | |
74 | if self._open: | |
75 | self._open = False | |
80ded8e9 RF |
76 | if self._drain_thread is not None: |
77 | thread, self._drain_thread = self._drain_thread, None | |
0fc8f660 | 78 | thread.join() |
80ded8e9 | 79 | socket.socket.close(self) |
0fc8f660 RF |
80 | if self._logfile: |
81 | self._logfile.close() | |
82 | self._logfile = None | |
83 | ||
80ded8e9 | 84 | def _drain_socket(self): |
0fc8f660 | 85 | """process arriving characters into in memory _buffer""" |
80ded8e9 | 86 | data = socket.socket.recv(self, 1) |
4b84d874 RF |
87 | # latin1 is needed since there are some chars |
88 | # we are receiving that cannot be encoded to utf-8 | |
89 | # such as 0xe2, 0x80, 0xA6. | |
90 | string = data.decode("latin1") | |
0fc8f660 RF |
91 | if self._logfile: |
92 | self._logfile.write("{}".format(string)) | |
93 | self._logfile.flush() | |
94 | for c in string: | |
95 | self._buffer.extend(c) | |
96 | ||
ff3513e6 | 97 | def recv(self, bufsize: int = 1, flags: int = 0) -> bytes: |
4b84d874 RF |
98 | """Return chars from in memory buffer. |
99 | Maintains the same API as socket.socket.recv. | |
100 | """ | |
80ded8e9 RF |
101 | if self._drain_thread is None: |
102 | # Not buffering the socket, pass thru to socket. | |
ff3513e6 JS |
103 | return socket.socket.recv(self, bufsize, flags) |
104 | assert not flags, "Cannot pass flags to recv() in drained mode" | |
0fc8f660 | 105 | start_time = time.time() |
80ded8e9 | 106 | while len(self._buffer) < bufsize: |
4b84d874 | 107 | time.sleep(self._sleep_time) |
0fc8f660 RF |
108 | elapsed_sec = time.time() - start_time |
109 | if elapsed_sec > self._recv_timeout_sec: | |
110 | raise socket.timeout | |
80ded8e9 | 111 | chars = ''.join([self._buffer.popleft() for i in range(bufsize)]) |
0fc8f660 RF |
112 | # We choose to use latin1 to remain consistent with |
113 | # handle_read() and give back the same data as the user would | |
114 | # receive if they were reading directly from the | |
115 | # socket w/o our intervention. | |
116 | return chars.encode("latin1") | |
117 | ||
80ded8e9 RF |
118 | def setblocking(self, value): |
119 | """When not draining we pass thru to the socket, | |
120 | since when draining we control socket blocking. | |
121 | """ | |
122 | if self._drain_thread is None: | |
123 | socket.socket.setblocking(self, value) | |
0fc8f660 | 124 | |
6cf4cce7 | 125 | def settimeout(self, value: Optional[float]) -> None: |
80ded8e9 RF |
126 | """When not draining we pass thru to the socket, |
127 | since when draining we control the timeout. | |
128 | """ | |
6cf4cce7 JS |
129 | if value is not None: |
130 | self._recv_timeout_sec = value | |
80ded8e9 | 131 | if self._drain_thread is None: |
6cf4cce7 | 132 | socket.socket.settimeout(self, value) |