]>
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 |
4b84d874 | 20 | |
0fc8f660 | 21 | |
80ded8e9 | 22 | class ConsoleSocket(socket.socket): |
4b84d874 RF |
23 | """ |
24 | ConsoleSocket represents a socket attached to a char device. | |
0fc8f660 | 25 | |
80ded8e9 RF |
26 | Optionally (if drain==True), drains the socket and places the bytes |
27 | into an in memory buffer for later processing. | |
4b84d874 RF |
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 | """ | |
80ded8e9 | 32 | def __init__(self, address, file=None, drain=False): |
0fc8f660 | 33 | self._recv_timeout_sec = 300 |
4b84d874 | 34 | self._sleep_time = 0.5 |
0fc8f660 | 35 | self._buffer = deque() |
80ded8e9 RF |
36 | socket.socket.__init__(self, socket.AF_UNIX, socket.SOCK_STREAM) |
37 | self.connect(address) | |
0fc8f660 RF |
38 | self._logfile = None |
39 | if file: | |
40 | self._logfile = open(file, "w") | |
0fc8f660 | 41 | self._open = True |
80ded8e9 RF |
42 | if drain: |
43 | self._drain_thread = self._thread_start() | |
44 | else: | |
45 | self._drain_thread = None | |
0fc8f660 | 46 | |
80ded8e9 RF |
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) | |
0fc8f660 | 57 | |
80ded8e9 RF |
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 | |
0fc8f660 RF |
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 | |
80ded8e9 RF |
74 | if self._drain_thread is not None: |
75 | thread, self._drain_thread = self._drain_thread, None | |
0fc8f660 | 76 | thread.join() |
80ded8e9 | 77 | socket.socket.close(self) |
0fc8f660 RF |
78 | if self._logfile: |
79 | self._logfile.close() | |
80 | self._logfile = None | |
81 | ||
80ded8e9 | 82 | def _drain_socket(self): |
0fc8f660 | 83 | """process arriving characters into in memory _buffer""" |
80ded8e9 | 84 | data = socket.socket.recv(self, 1) |
4b84d874 RF |
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") | |
0fc8f660 RF |
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 | ||
ff3513e6 | 95 | def recv(self, bufsize: int = 1, flags: int = 0) -> bytes: |
4b84d874 RF |
96 | """Return chars from in memory buffer. |
97 | Maintains the same API as socket.socket.recv. | |
98 | """ | |
80ded8e9 RF |
99 | if self._drain_thread is None: |
100 | # Not buffering the socket, pass thru to socket. | |
ff3513e6 JS |
101 | return socket.socket.recv(self, bufsize, flags) |
102 | assert not flags, "Cannot pass flags to recv() in drained mode" | |
0fc8f660 | 103 | start_time = time.time() |
80ded8e9 | 104 | while len(self._buffer) < bufsize: |
4b84d874 | 105 | time.sleep(self._sleep_time) |
0fc8f660 RF |
106 | elapsed_sec = time.time() - start_time |
107 | if elapsed_sec > self._recv_timeout_sec: | |
108 | raise socket.timeout | |
80ded8e9 | 109 | chars = ''.join([self._buffer.popleft() for i in range(bufsize)]) |
0fc8f660 RF |
110 | # We choose to use latin1 to remain consistent with |
111 | # handle_read() and give back the same data as the user would | |
112 | # receive if they were reading directly from the | |
113 | # socket w/o our intervention. | |
114 | return chars.encode("latin1") | |
115 | ||
80ded8e9 RF |
116 | def setblocking(self, value): |
117 | """When not draining we pass thru to the socket, | |
118 | since when draining we control socket blocking. | |
119 | """ | |
120 | if self._drain_thread is None: | |
121 | socket.socket.setblocking(self, value) | |
0fc8f660 RF |
122 | |
123 | def settimeout(self, seconds): | |
80ded8e9 RF |
124 | """When not draining we pass thru to the socket, |
125 | since when draining we control the timeout. | |
126 | """ | |
127 | if seconds is not None: | |
128 | self._recv_timeout_sec = seconds | |
129 | if self._drain_thread is None: | |
130 | socket.socket.settimeout(self, seconds) |