Skip to content

Commit c8227a6

Browse files
committed
feat(core): add socket module
Signed-off-by: Ronan Abhamon <ronan.abhamon@vates.tech>
1 parent 57397c3 commit c8227a6

2 files changed

Lines changed: 659 additions & 0 deletions

File tree

‎src/xcp_storage/network/socket.py‎

Lines changed: 325 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,325 @@
1+
# Copyright (C) 2026 Vates SAS
2+
#
3+
# This program is free software: you can redistribute it and/or modify
4+
# it under the terms of the GNU General Public License as published by
5+
# the Free Software Foundation, either version 3 of the License, or
6+
# (at your option) any later version.
7+
# This program is distributed in the hope that it will be useful,
8+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
9+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
10+
# GNU General Public License for more details.
11+
#
12+
# You should have received a copy of the GNU General Public License
13+
# along with this program. If not, see <https://www.gnu.org/licenses/>.
14+
15+
import contextlib
16+
import errno
17+
import ipaddress
18+
import select
19+
import socket
20+
import ssl
21+
from types import TracebackType
22+
23+
from xcp_storage.typing import (
24+
Any,
25+
Dict,
26+
Final,
27+
Optional,
28+
override,
29+
Tuple,
30+
Type,
31+
Union,
32+
)
33+
34+
# ==============================================================================
35+
36+
_FAMILY_TO_STR: Final = {
37+
socket.AF_INET: "IPv4",
38+
socket.AF_INET6: "IPv6",
39+
socket.AF_UNIX: "Unix"
40+
}
41+
42+
_IP_VERSION_TO_FAMILY: Final = {
43+
4: socket.AF_INET,
44+
6: socket.AF_INET6,
45+
0: socket.AF_UNSPEC
46+
}
47+
48+
_SERVER_BACKLOG: Final = 128
49+
50+
# ------------------------------------------------------------------------------
51+
52+
class SocketError(Exception):
53+
def __init__(self, message: str) -> None:
54+
super().__init__(message)
55+
56+
class SocketTimeoutError(SocketError):
57+
def __init__(self) -> None:
58+
super().__init__("Socket timeout.")
59+
60+
class SocketDisconnectedError(SocketError):
61+
def __init__(self, message: str) -> None:
62+
super().__init__(message)
63+
64+
# ------------------------------------------------------------------------------
65+
66+
def get_ip_address(address: str, ip_version: int = 0) -> Union[ipaddress.IPv4Address, ipaddress.IPv6Address]:
67+
# Note: `address` can be a hostname or IP.
68+
with contextlib.suppress(ValueError):
69+
return ipaddress.ip_address(address)
70+
71+
family = _IP_VERSION_TO_FAMILY.get(ip_version)
72+
if family is None:
73+
raise SocketError("Unknown IP version.")
74+
75+
try:
76+
info = socket.getaddrinfo(address or socket.gethostname(), 80, family, socket.SOCK_STREAM, socket.SOL_TCP)
77+
return ipaddress.ip_address(info[0][4][0])
78+
except socket.gaierror as e:
79+
raise SocketError("Cannot resolve IP.") from e
80+
except IndexError:
81+
raise SocketError("Cannot resolve IP: no valid address.") from None
82+
83+
# ------------------------------------------------------------------------------
84+
85+
def format_address(address: str, port: int) -> Tuple[socket.AddressFamily, Union[
86+
Tuple[str, int],
87+
Tuple[str, int, int, int]
88+
]]:
89+
if not address:
90+
raise SocketError("No hostname/IP.")
91+
92+
ip_address = get_ip_address(address)
93+
if ip_address.version == 4:
94+
family = socket.AF_INET
95+
return (family, (str(ip_address), port))
96+
if ip_address.version == 6:
97+
family = socket.AF_INET6
98+
return (family, (str(ip_address), port, 0, 0))
99+
100+
raise SocketError("Unknown IP version.")
101+
102+
# ------------------------------------------------------------------------------
103+
104+
def _create_stream_sock(
105+
address: str,
106+
family: socket.AddressFamily,
107+
*,
108+
bind: bool,
109+
reuse_address: bool,
110+
ssl_context: Optional[ssl.SSLContext]
111+
) -> socket.socket:
112+
sock = socket.socket(family, socket.SOCK_STREAM)
113+
if ssl_context:
114+
if bind:
115+
sock = ssl_context.wrap_socket(sock, server_side=True)
116+
else:
117+
sock = ssl_context.wrap_socket(sock, server_side=False, server_hostname=address or None)
118+
119+
if reuse_address:
120+
set_socket_reuseaddr(sock)
121+
122+
return sock
123+
124+
def _normalize_timeout(timeout: Optional[float]) -> Optional[float]:
125+
if timeout is not None and timeout < 0:
126+
timeout = None
127+
return timeout
128+
129+
# ------------------------------------------------------------------------------
130+
131+
def create_server_sock(
132+
address: str,
133+
port: int,
134+
*,
135+
reuse_address: bool = True,
136+
keep_alive: bool = True,
137+
timeout: Optional[float] = None,
138+
ssl_context: Optional[ssl.SSLContext] = None
139+
) -> socket.socket:
140+
family, bind = format_address(address, port)
141+
sock = _create_stream_sock(address, family, bind=True, reuse_address=reuse_address, ssl_context=ssl_context)
142+
143+
timeout = _normalize_timeout(timeout)
144+
if timeout is not None:
145+
sock.settimeout(timeout)
146+
147+
try:
148+
sock.bind(bind)
149+
sock.listen(_SERVER_BACKLOG)
150+
except OSError as e:
151+
with contextlib.suppress(Exception):
152+
sock.close()
153+
raise SocketError("Failed to bind server sock.") from e
154+
155+
if keep_alive:
156+
set_socket_keepalive(sock)
157+
158+
return sock
159+
160+
def create_client_sock(
161+
address: str,
162+
port: int,
163+
*,
164+
reuse_address: bool = True,
165+
keep_alive: bool = True,
166+
timeout: Optional[float] = None,
167+
ssl_context: Optional[ssl.SSLContext] = None
168+
) -> socket.socket:
169+
family, connect = format_address(address, port)
170+
sock = _create_stream_sock(address, family, bind=False, reuse_address=reuse_address, ssl_context=ssl_context)
171+
172+
timeout = _normalize_timeout(timeout)
173+
if timeout is not None:
174+
sock.settimeout(timeout)
175+
176+
while True:
177+
try:
178+
sock.connect(connect)
179+
break
180+
except OSError as e:
181+
if e.errno in (errno.EAGAIN, errno.EINPROGRESS):
182+
_, ready, _ = select.select([], [sock], [], timeout)
183+
if sock in ready:
184+
continue
185+
186+
with contextlib.suppress(Exception):
187+
sock.close()
188+
raise SocketError("Failed to connect client sock.") from e
189+
190+
if keep_alive:
191+
set_socket_keepalive(sock)
192+
193+
return sock
194+
195+
# ------------------------------------------------------------------------------
196+
197+
def set_socket_reuseaddr(sock: socket.socket) -> None:
198+
with contextlib.suppress(Exception):
199+
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
200+
201+
def set_socket_keepalive(sock: socket.socket) -> None:
202+
with contextlib.suppress(Exception):
203+
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
204+
205+
# ------------------------------------------------------------------------------
206+
207+
def _get_buffer_size(buffer: Union[bytes, bytearray], size: Optional[int] = None) -> int:
208+
if size is None:
209+
size = len(buffer)
210+
else:
211+
assert len(buffer) >= size, "Buffer size must be greater than or equal to size."
212+
return size
213+
214+
# ------------------------------------------------------------------------------
215+
216+
def socket_send(sock: socket.socket, buffer: bytes, size: Optional[int] = None) -> None:
217+
size = _get_buffer_size(buffer, size)
218+
view = memoryview(buffer)
219+
220+
pos = 0
221+
while pos < size:
222+
try:
223+
n = sock.send(view[pos:size])
224+
if not n:
225+
break
226+
pos += n
227+
# TODO(XCPNG-3032): `socket.timeout` is a workaround for python 3.6. Must be removed later.
228+
except (TimeoutError, socket.timeout):
229+
raise SocketTimeoutError() from None
230+
except OSError as e:
231+
if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
232+
select.select([], [sock], [])
233+
continue
234+
raise SocketDisconnectedError("Unable to send data.") from e
235+
236+
if pos != size:
237+
raise SocketDisconnectedError("Not enough data sent.") from None
238+
239+
# ------------------------------------------------------------------------------
240+
241+
def socket_receive(sock: socket.socket, buffer: bytearray, size: Optional[int] = None) -> None:
242+
size = _get_buffer_size(buffer, size)
243+
view = memoryview(buffer)
244+
245+
pos = 0
246+
while pos < size:
247+
try:
248+
n = sock.recv_into(view[pos:], min(size - pos, 8192))
249+
if not n:
250+
break
251+
pos += n
252+
# TODO(XCPNG-3032): `socket.timeout` is a workaround for python 3.6. Must be removed later.
253+
except (TimeoutError, socket.timeout):
254+
raise SocketTimeoutError() from None
255+
except OSError as e:
256+
if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
257+
select.select([sock], [], [])
258+
continue
259+
raise SocketDisconnectedError("Unable to receive data.") from e
260+
261+
if pos != size:
262+
raise SocketDisconnectedError("Not enough data received.") from None
263+
264+
# ------------------------------------------------------------------------------
265+
266+
def get_socket_family_str(sock: socket.socket) -> str:
267+
return _FAMILY_TO_STR.get(sock.family, "Unknown")
268+
269+
# ------------------------------------------------------------------------------
270+
271+
class Socket(contextlib.AbstractContextManager):
272+
def __init__(self, sock: socket.socket, *, keep_open: bool = False) -> None:
273+
self.sock = sock
274+
self.keep_open = keep_open
275+
276+
def __del__(self) -> None:
277+
self.close()
278+
279+
@override
280+
def __enter__(self) -> "Socket":
281+
# `Socket` is a socket manager/wrapper that does not create sockets directly.
282+
# It's simply built around a `socket.socket` instance; this is why `__enter__` does nothing,
283+
# whereas `__exit__` can close a connection.
284+
return self
285+
286+
@override
287+
def __exit__(
288+
self,
289+
exc_type: Optional[Type[BaseException]],
290+
exc_value: Optional[BaseException],
291+
traceback: Optional[TracebackType]
292+
) -> None:
293+
self.close()
294+
295+
def send(self, buffer: bytes, size: Optional[int] = None) -> None:
296+
socket_send(self.sock, buffer, size)
297+
298+
def receive(self, buffer: bytearray, size: Optional[int] = None) -> None:
299+
socket_receive(self.sock, buffer, size)
300+
301+
def close(self) -> None:
302+
if self.keep_open:
303+
return
304+
with contextlib.suppress(Exception):
305+
self.sock.shutdown(socket.SHUT_RDWR)
306+
with contextlib.suppress(Exception):
307+
self.sock.close()
308+
309+
@property
310+
def family_str(self) -> str:
311+
return get_socket_family_str(self.sock)
312+
313+
@property
314+
def timeout(self) -> Optional[float]:
315+
return self.sock.gettimeout()
316+
317+
@timeout.setter
318+
def timeout(self, value: Optional[float]) -> None:
319+
self.sock.settimeout(value)
320+
321+
@property
322+
def peer_certificate(self) -> Optional[Dict[str, Any]]:
323+
if isinstance(self.sock, ssl.SSLSocket):
324+
return self.sock.getpeercert()
325+
return None

0 commit comments

Comments
 (0)