1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
|
# Copyright 2023 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Wrapers for socket clients to log read and write data."""
from __future__ import annotations
from typing import Callable, Optional, TYPE_CHECKING, Tuple, Union
import errno
import re
import socket
from pw_console.plugins.bandwidth_toolbar import SerialBandwidthTracker
if TYPE_CHECKING:
from _typeshed import ReadableBuffer
class SocketClient:
"""Socket transport implementation."""
FILE_SOCKET_SERVER = 'file'
DEFAULT_SOCKET_SERVER = 'localhost'
DEFAULT_SOCKET_PORT = 33000
PW_RPC_MAX_PACKET_SIZE = 256
_InitArgsType = Tuple[
socket.AddressFamily, int # pylint: disable=no-member
]
# Can be a string, (address, port) for AF_INET or (address, port, flowinfo,
# scope_id) AF_INET6.
_AddressType = Union[str, Tuple[str, int], Tuple[str, int, int, int]]
def __init__(
self,
config: str,
on_disconnect: Optional[Callable[[SocketClient], None]] = None,
):
"""Creates a socket connection.
Args:
config: The socket configuration. Accepted values and formats are:
'default' - uses the default configuration (localhost:33000)
'address:port' - An IPv4 address and port.
'address' - An IPv4 address. Uses default port 33000.
'[address]:port' - An IPv6 address and port.
'[address]' - An IPv6 address. Uses default port 33000.
'file:path_to_file' - A Unix socket at ``path_to_file``.
In the formats above,``address`` can be an actual address or a name
that resolves to an address through name-resolution.
on_disconnect: An optional callback called when the socket
disconnects.
Raises:
TypeError: The type of socket is not supported.
ValueError: The socket configuration is invalid.
"""
self.socket: socket.socket
(
self._socket_init_args,
self._address,
) = SocketClient._parse_socket_config(config)
self._on_disconnect = on_disconnect
self._connected = False
self.connect()
@staticmethod
def _parse_socket_config(
config: str,
) -> Tuple[SocketClient._InitArgsType, SocketClient._AddressType]:
"""Sets the variables used to create a socket given a config string.
Raises:
TypeError: The type of socket is not supported.
ValueError: The socket configuration is invalid.
"""
init_args: SocketClient._InitArgsType
address: SocketClient._AddressType
# Check if this is using the default settings.
if config == 'default':
init_args = socket.AF_INET6, socket.SOCK_STREAM
address = (
SocketClient.DEFAULT_SOCKET_SERVER,
SocketClient.DEFAULT_SOCKET_PORT,
)
return init_args, address
# Check if this is a UNIX socket.
unix_socket_file_setting = f'{SocketClient.FILE_SOCKET_SERVER}:'
if config.startswith(unix_socket_file_setting):
# Unix socket support is available on Windows 10 since April
# 2018. However, there is no Python support on Windows yet.
# See https://bugs.python.org/issue33408 for more information.
if not hasattr(socket, 'AF_UNIX'):
raise TypeError(
'Unix sockets are not supported in this environment.'
)
init_args = (
socket.AF_UNIX, # pylint: disable=no-member
socket.SOCK_STREAM,
)
address = config[len(unix_socket_file_setting) :]
return init_args, address
# Search for IPv4 or IPv6 address or name and port.
# First, try to capture an IPv6 address as anything inside []. If there
# are no [] capture the IPv4 address. Lastly, capture the port as the
# numbers after :, if any.
match = re.match(
r'(\[(?P<ipv6_addr>.+)\]:?|(?P<ipv4_addr>[a-zA-Z0-9\._\/]+):?)'
r'(?P<port>[0-9]+)?',
config,
)
invalid_config_message = (
f'Invalid socket configuration "{config}"'
'Accepted values are "default", "file:<file_path>", '
'"<name_or_ipv4_address>" with optional ":<port>", and '
'"[<name_or_ipv6_address>]" with optional ":<port>".'
)
if match is None:
raise ValueError(invalid_config_message)
info = match.groupdict()
if info['port']:
port = int(info['port'])
else:
port = SocketClient.DEFAULT_SOCKET_PORT
if info['ipv4_addr']:
ip_addr = info['ipv4_addr']
elif info['ipv6_addr']:
ip_addr = info['ipv6_addr']
else:
raise ValueError(invalid_config_message)
sock_family, sock_type, _, _, address = socket.getaddrinfo(
ip_addr, port, type=socket.SOCK_STREAM
)[0]
init_args = sock_family, sock_type
return init_args, address
def __del__(self):
if self._connected:
self.socket.close()
def write(self, data: ReadableBuffer) -> None:
"""Writes data and detects disconnects."""
if not self._connected:
raise Exception('Socket is not connected.')
try:
self.socket.sendall(data)
except socket.error as exc:
if isinstance(exc.args, tuple) and exc.args[0] == errno.EPIPE:
self._handle_disconnect()
else:
raise exc
def read(self, num_bytes: int = PW_RPC_MAX_PACKET_SIZE) -> bytes:
"""Blocks until data is ready and reads up to num_bytes."""
if not self._connected:
raise Exception('Socket is not connected.')
data = self.socket.recv(num_bytes)
# Since this is a blocking read, no data returned means the socket is
# closed.
if not data:
self._handle_disconnect()
return data
def connect(self) -> None:
"""Connects to socket."""
self.socket = socket.socket(*self._socket_init_args)
# Enable reusing address and port for reconnections.
self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
if hasattr(socket, 'SO_REUSEPORT'):
self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
self.socket.connect(self._address)
self._connected = True
def _handle_disconnect(self):
"""Escalates a socket disconnect to the user."""
self.socket.close()
self._connected = False
if self._on_disconnect:
self._on_disconnect(self)
def fileno(self) -> int:
return self.socket.fileno()
class SocketClientWithLogging(SocketClient):
"""Socket with read and write wrappers for logging."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._bandwidth_tracker = SerialBandwidthTracker()
def read(
self, num_bytes: int = SocketClient.PW_RPC_MAX_PACKET_SIZE
) -> bytes:
data = super().read(num_bytes)
self._bandwidth_tracker.track_read_data(data)
return data
def write(self, data: ReadableBuffer) -> None:
self._bandwidth_tracker.track_write_data(data)
super().write(data)
|
