Skip to content

Commit 755d35c

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

2 files changed

Lines changed: 660 additions & 0 deletions

File tree

‎src/xcp_storage/network/socket.py‎

Lines changed: 326 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,326 @@
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+
assert family is not None
77+
info = socket.getaddrinfo(address or socket.gethostname(), 80, family, socket.SOCK_STREAM, socket.SOL_TCP)
78+
return ipaddress.ip_address(info[0][4][0])
79+
except socket.gaierror as e:
80+
raise SocketError("Cannot resolve IP.") from e
81+
except IndexError:
82+
raise SocketError("Cannot resolve IP: no valid address.") from None
83+
84+
# ------------------------------------------------------------------------------
85+
86+
def format_address(address: str, port: int) -> Tuple[socket.AddressFamily, Union[
87+
Tuple[str, int],
88+
Tuple[str, int, int, int]
89+
]]:
90+
if not address:
91+
raise SocketError("No hostname/IP.")
92+
93+
ip_address = get_ip_address(address)
94+
if ip_address.version == 4:
95+
family = socket.AF_INET
96+
return (family, (str(ip_address), port))
97+
if ip_address.version == 6:
98+
family = socket.AF_INET6
99+
return (family, (str(ip_address), port, 0, 0))
100+
101+
raise SocketError("Unknown IP version.")
102+
103+
# ------------------------------------------------------------------------------
104+
105+
def _create_stream_sock(
106+
address: str,
107+
family: socket.AddressFamily,
108+
*,
109+
bind: bool,
110+
reuse_address: bool,
111+
ssl_context: Optional[ssl.SSLContext]
112+
) -> socket.socket:
113+
sock = socket.socket(family, socket.SOCK_STREAM)
114+
if ssl_context:
115+
if bind:
116+
sock = ssl_context.wrap_socket(sock, server_side=True)
117+
else:
118+
sock = ssl_context.wrap_socket(sock, server_side=False, server_hostname=address or None)
119+
120+
if reuse_address:
121+
set_socket_reuseaddr(sock)
122+
123+
return sock
124+
125+
def _normalize_timeout(timeout: Optional[float]) -> Optional[float]:
126+
if timeout is not None and timeout < 0:
127+
timeout = None
128+
return timeout
129+
130+
# ------------------------------------------------------------------------------
131+
132+
def create_server_sock(
133+
address: str,
134+
port: int,
135+
*,
136+
reuse_address: bool = True,
137+
keep_alive: bool = True,
138+
timeout: Optional[float] = None,
139+
ssl_context: Optional[ssl.SSLContext] = None
140+
) -> socket.socket:
141+
family, bind = format_address(address, port)
142+
sock = _create_stream_sock(address, family, bind=True, reuse_address=reuse_address, ssl_context=ssl_context)
143+
144+
timeout = _normalize_timeout(timeout)
145+
if timeout is not None:
146+
sock.settimeout(timeout)
147+
148+
try:
149+
sock.bind(bind)
150+
sock.listen(_SERVER_BACKLOG)
151+
except OSError as e:
152+
with contextlib.suppress(Exception):
153+
sock.close()
154+
raise SocketError("Failed to bind server sock.") from e
155+
156+
if keep_alive:
157+
set_socket_keepalive(sock)
158+
159+
return sock
160+
161+
def create_client_sock(
162+
address: str,
163+
port: int,
164+
*,
165+
reuse_address: bool = True,
166+
keep_alive: bool = True,
167+
timeout: Optional[float] = None,
168+
ssl_context: Optional[ssl.SSLContext] = None
169+
) -> socket.socket:
170+
family, connect = format_address(address, port)
171+
sock = _create_stream_sock(address, family, bind=False, reuse_address=reuse_address, ssl_context=ssl_context)
172+
173+
timeout = _normalize_timeout(timeout)
174+
if timeout is not None:
175+
sock.settimeout(timeout)
176+
177+
while True:
178+
try:
179+
sock.connect(connect)
180+
break
181+
except OSError as e:
182+
if e.errno in (errno.EAGAIN, errno.EINPROGRESS):
183+
_, ready, _ = select.select([], [sock], [], timeout)
184+
if sock in ready:
185+
continue
186+
187+
with contextlib.suppress(Exception):
188+
sock.close()
189+
raise SocketError("Failed to connect client sock.") from e
190+
191+
if keep_alive:
192+
set_socket_keepalive(sock)
193+
194+
return sock
195+
196+
# ------------------------------------------------------------------------------
197+
198+
def set_socket_reuseaddr(sock: socket.socket) -> None:
199+
with contextlib.suppress(Exception):
200+
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
201+
202+
def set_socket_keepalive(sock: socket.socket) -> None:
203+
with contextlib.suppress(Exception):
204+
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
205+
206+
# ------------------------------------------------------------------------------
207+
208+
def _get_buffer_size(buffer: Union[bytes, bytearray], size: Optional[int] = None) -> int:
209+
if size is None:
210+
size = len(buffer)
211+
else:
212+
assert len(buffer) >= size, "Buffer size must be greater than or equal to size."
213+
return size
214+
215+
# ------------------------------------------------------------------------------
216+
217+
def socket_send(sock: socket.socket, buffer: bytes, size: Optional[int] = None) -> None:
218+
size = _get_buffer_size(buffer, size)
219+
view = memoryview(buffer)
220+
221+
pos = 0
222+
while pos < size:
223+
try:
224+
n = sock.send(view[pos:size])
225+
if not n:
226+
break
227+
pos += n
228+
# TODO(XCPNG-3032): `socket.timeout` is a workaround for python 3.6. Must be removed later.
229+
except (TimeoutError, socket.timeout):
230+
raise SocketTimeoutError() from None
231+
except OSError as e:
232+
if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
233+
select.select([], [sock], [])
234+
continue
235+
raise SocketDisconnectedError("Unable to send data.") from e
236+
237+
if pos != size:
238+
raise SocketDisconnectedError("Not enough data sent.") from None
239+
240+
# ------------------------------------------------------------------------------
241+
242+
def socket_receive(sock: socket.socket, buffer: bytearray, size: Optional[int] = None) -> None:
243+
size = _get_buffer_size(buffer, size)
244+
view = memoryview(buffer)
245+
246+
pos = 0
247+
while pos < size:
248+
try:
249+
n = sock.recv_into(view[pos:], min(size - pos, 8192))
250+
if not n:
251+
break
252+
pos += n
253+
# TODO(XCPNG-3032): `socket.timeout` is a workaround for python 3.6. Must be removed later.
254+
except (TimeoutError, socket.timeout):
255+
raise SocketTimeoutError() from None
256+
except OSError as e:
257+
if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
258+
select.select([sock], [], [])
259+
continue
260+
raise SocketDisconnectedError("Unable to receive data.") from e
261+
262+
if pos != size:
263+
raise SocketDisconnectedError("Not enough data received.") from None
264+
265+
# ------------------------------------------------------------------------------
266+
267+
def get_socket_family_str(sock: socket.socket) -> str:
268+
return _FAMILY_TO_STR.get(sock.family, "Unknown")
269+
270+
# ------------------------------------------------------------------------------
271+
272+
class Socket(contextlib.AbstractContextManager):
273+
def __init__(self, sock: socket.socket, *, keep_open: bool = False) -> None:
274+
self.sock = sock
275+
self.keep_open = keep_open
276+
277+
def __del__(self) -> None:
278+
self.close()
279+
280+
@override
281+
def __enter__(self) -> "Socket":
282+
# `Socket` is a socket manager/wrapper that does not create sockets directly.
283+
# It's simply built around a `socket.socket` instance; this is why `__enter__` does nothing,
284+
# whereas `__exit__` can close a connection.
285+
return self
286+
287+
@override
288+
def __exit__(
289+
self,
290+
exc_type: Optional[Type[BaseException]],
291+
exc_value: Optional[BaseException],
292+
traceback: Optional[TracebackType]
293+
) -> None:
294+
self.close()
295+
296+
def send(self, buffer: bytes, size: Optional[int] = None) -> None:
297+
socket_send(self.sock, buffer, size)
298+
299+
def receive(self, buffer: bytearray, size: Optional[int] = None) -> None:
300+
socket_receive(self.sock, buffer, size)
301+
302+
def close(self) -> None:
303+
if self.keep_open:
304+
return
305+
with contextlib.suppress(Exception):
306+
self.sock.shutdown(socket.SHUT_RDWR)
307+
with contextlib.suppress(Exception):
308+
self.sock.close()
309+
310+
@property
311+
def family_str(self) -> str:
312+
return get_socket_family_str(self.sock)
313+
314+
@property
315+
def timeout(self) -> Optional[float]:
316+
return self.sock.gettimeout()
317+
318+
@timeout.setter
319+
def timeout(self, value: Optional[float]) -> None:
320+
self.sock.settimeout(value)
321+
322+
@property
323+
def peer_certificate(self) -> Optional[Dict[str, Any]]:
324+
if isinstance(self.sock, ssl.SSLSocket):
325+
return self.sock.getpeercert()
326+
return None

0 commit comments

Comments
 (0)