3fd864e5a9
Optional types were provided in places but were not always correct. Add mypy checking and fix those that fail(ed).
418 lines
14 KiB
Python
418 lines
14 KiB
Python
"""JSON Communication
|
|
|
|
This module implements a client/server communication method based on JSON
|
|
serialization. It uses unix-domain-datagram-sockets and provides a simple
|
|
unicast message transmission.
|
|
"""
|
|
|
|
|
|
import array
|
|
import contextlib
|
|
import errno
|
|
import json
|
|
import os
|
|
import socket
|
|
from typing import Any
|
|
from typing import Optional
|
|
from .types import PathLike
|
|
|
|
|
|
class FdSet:
|
|
"""File-Descriptor Set
|
|
|
|
This object wraps an array of file-descriptors. Unlike a normal integer
|
|
array, this object owns the file-descriptors and therefore closes them once
|
|
the object is released.
|
|
|
|
File-descriptor sets are initialized once. From then one, the only allowed
|
|
operation is to query it for information, or steal file-descriptors from
|
|
it. If you close a set, all remaining file-descriptors are closed and
|
|
removed from the set. It will then be an empty set.
|
|
"""
|
|
|
|
_fds = array.array("i")
|
|
|
|
def __init__(self, *, rawfds):
|
|
for i in rawfds:
|
|
if not isinstance(i, int) or i < 0:
|
|
raise ValueError()
|
|
|
|
self._fds = rawfds
|
|
|
|
def __del__(self):
|
|
self.close()
|
|
|
|
def close(self):
|
|
"""Close All Entries
|
|
|
|
This closes all stored file-descriptors and clears the set. Once this
|
|
returns, the set will be empty. It is safe to call this multiple times.
|
|
Note that a set is automatically closed when it is garbage collected.
|
|
"""
|
|
|
|
for i in self._fds:
|
|
if i >= 0:
|
|
os.close(i)
|
|
|
|
self._fds = array.array("i")
|
|
|
|
@classmethod
|
|
def from_list(cls, l: list):
|
|
"""Create new Set from List
|
|
|
|
This creates a new file-descriptor set initialized to the same entries
|
|
as in the given list. This consumes the file-descriptors. The caller
|
|
must not assume ownership anymore.
|
|
"""
|
|
|
|
fds = array.array("i")
|
|
fds.fromlist(l)
|
|
return cls(rawfds=fds)
|
|
|
|
def __len__(self):
|
|
return len(self._fds)
|
|
|
|
def __getitem__(self, key: Any):
|
|
if self._fds[key] < 0:
|
|
raise IndexError
|
|
return self._fds[key]
|
|
|
|
def steal(self, key: Any):
|
|
"""Steal Entry
|
|
|
|
Retrieve the entry at the given position, but drop it from the internal
|
|
file-descriptor set. The caller will now own the file-descriptor and it
|
|
can no longer be accessed through the set.
|
|
|
|
Note that this does not reshuffle the set. All indices stay constant.
|
|
"""
|
|
|
|
v = self[key]
|
|
self._fds[key] = -1
|
|
return v
|
|
|
|
|
|
class Socket(contextlib.AbstractContextManager):
|
|
"""Communication Socket
|
|
|
|
This socket object represents a communication channel. It allows sending
|
|
and receiving JSON-encoded messages. It uses unix-domain sequenced-packet
|
|
sockets as underlying transport.
|
|
"""
|
|
|
|
_socket = None
|
|
_unlink = None
|
|
|
|
def __init__(self, sock, unlink):
|
|
self._socket = sock
|
|
self._unlink = unlink
|
|
|
|
def __del__(self):
|
|
self.close()
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_tb):
|
|
self.close()
|
|
return False
|
|
|
|
@property
|
|
def blocking(self):
|
|
"""Get the current blocking mode of the socket.
|
|
|
|
This is related to the socket's timeout, i.e. if no time out is set
|
|
the socket is in blocking mode; otherwise it is non-blocking.
|
|
"""
|
|
timeout = self._socket.gettimeout()
|
|
return timeout is not None
|
|
|
|
@blocking.setter
|
|
def blocking(self, value: bool):
|
|
"""Set the blocking mode of the socket."""
|
|
if self._socket:
|
|
self._socket.setblocking(value)
|
|
else:
|
|
raise RuntimeError("Tried to set blocking mode without socket.")
|
|
|
|
def accept(self) -> Optional["Socket"]:
|
|
"""Accept a new connection on the socket.
|
|
|
|
See python's `socket.accept` for more information.
|
|
"""
|
|
|
|
if not self._socket:
|
|
raise RuntimeError("Tried to accept without socket.")
|
|
|
|
# Since, in the kernel, for AF_UNIX, new connection requests,
|
|
# i.e. clients connecting, are directly put on the receive
|
|
# queue of the listener socket, accept here *should* always
|
|
# return a socket and not block, even if the client meanwhile
|
|
# disconnected; we don't rely on that kernel behavior though
|
|
try:
|
|
conn, _ = self._socket.accept()
|
|
except (socket.timeout, BlockingIOError):
|
|
return None
|
|
return Socket(conn, None)
|
|
|
|
def listen(self, backlog: Optional[int] = 2**16):
|
|
"""Enable accepting of incoming connections.
|
|
|
|
See python's `socket.listen` for details.
|
|
"""
|
|
|
|
if not self._socket:
|
|
raise RuntimeError("Tried to listen without socket.")
|
|
|
|
# `Socket.listen` accepts an `int` or no argument, but not `None`
|
|
args = [backlog] if backlog is not None else []
|
|
self._socket.listen(*args)
|
|
|
|
def close(self):
|
|
"""Close Socket
|
|
|
|
Close the socket and all underlying resources. This can be called
|
|
multiple times.
|
|
"""
|
|
|
|
# close the socket if it is set
|
|
if self._socket is not None:
|
|
self._socket.close()
|
|
self._socket = None
|
|
|
|
# unlink the file-system entry, if pinned
|
|
if self._unlink is not None:
|
|
try:
|
|
os.unlink(self._unlink[1], dir_fd=self._unlink[0])
|
|
except OSError as e:
|
|
if e.errno != errno.ENOENT:
|
|
raise
|
|
|
|
os.close(self._unlink[0])
|
|
self._unlink = None
|
|
|
|
@classmethod
|
|
def new_client(cls, connect_to: Optional[PathLike] = None):
|
|
"""Create Client
|
|
|
|
Create a new client socket.
|
|
|
|
Parameters
|
|
----------
|
|
connect_to
|
|
If not `None`, the client will use the specified address as the
|
|
default destination for all send operations.
|
|
"""
|
|
|
|
sock = None
|
|
|
|
try:
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
|
|
|
|
# Trigger an auto-bind. If you do not do this, you might end up with
|
|
# an unbound unix socket, which cannot receive messages.
|
|
# Alternatively, you can also set `SO_PASSCRED`, but this has
|
|
# side-effects.
|
|
sock.bind("")
|
|
|
|
# Connect the socket. This has no effect other than specifying the
|
|
# default destination for send operations.
|
|
if connect_to is not None:
|
|
sock.connect(os.fspath(connect_to))
|
|
except:
|
|
if sock is not None:
|
|
sock.close()
|
|
raise
|
|
|
|
return cls(sock, None)
|
|
|
|
@classmethod
|
|
def new_server(cls, bind_to: PathLike):
|
|
"""Create Server
|
|
|
|
Create a new listener socket. Returned socket is in non-blocking
|
|
mode by default. See `blocking` property.
|
|
|
|
Parameters
|
|
----------
|
|
bind_to
|
|
The socket-address to listen on for incoming client requests.
|
|
"""
|
|
|
|
sock = None
|
|
unlink = None
|
|
path = os.path.split(bind_to)
|
|
|
|
try:
|
|
# We bind the socket and then open a directory-fd on the target
|
|
# socket. This allows us to properly unlink the socket when the
|
|
# server is closed. Note that sockets are never automatically
|
|
# cleaned up on linux, nor can you bind to existing sockets.
|
|
# We use a dirfd to guarantee this works even when you change
|
|
# your mount points in-between.
|
|
# Yeah, this is racy when mount-points change between the socket
|
|
# creation and open. But then your entire socket creation is racy
|
|
# as well. We do not guarantee atomicity, so you better make sure
|
|
# you do not rely on it.
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
|
|
sock.bind(os.fspath(bind_to))
|
|
unlink = os.open(os.path.join(".", path[0]), os.O_CLOEXEC | os.O_PATH)
|
|
sock.setblocking(False)
|
|
except:
|
|
if unlink is not None:
|
|
os.close(unlink)
|
|
if sock is not None:
|
|
sock.close()
|
|
raise
|
|
|
|
return cls(sock, (unlink, path[1]))
|
|
|
|
@classmethod
|
|
def new_pair(cls, *, blocking=True):
|
|
"""Create a connected socket pair
|
|
|
|
Create a pair of connected sockets and return both as a tuple.
|
|
|
|
Parameters
|
|
----------
|
|
blocking
|
|
The blocking mode for the socket pair.
|
|
"""
|
|
a, b = socket.socketpair(socket.AF_UNIX, socket.SOCK_SEQPACKET)
|
|
|
|
a.setblocking(blocking)
|
|
b.setblocking(blocking)
|
|
|
|
return cls(a, None), cls(b, None)
|
|
|
|
@classmethod
|
|
def new_from_fd(cls, fd: int, *, blocking=True, close_fd=True):
|
|
"""Create a socket for an existing file descriptor
|
|
|
|
Duplicate the file descriptor and return a `Socket` for it.
|
|
The blocking mode can be set via `blocking`. If `close_fd`
|
|
is True (the default) `fd` will be closed.
|
|
|
|
Parameters
|
|
----------
|
|
fd
|
|
The file descriptor to use.
|
|
blocking
|
|
The blocking mode for the socket pair.
|
|
"""
|
|
sock = socket.fromfd(fd, socket.AF_UNIX, socket.SOCK_SEQPACKET)
|
|
sock.setblocking(blocking)
|
|
if close_fd:
|
|
os.close(fd)
|
|
return cls(sock, None)
|
|
|
|
def fileno(self) -> int:
|
|
assert self._socket is not None
|
|
return self._socket.fileno()
|
|
|
|
def recv(self):
|
|
"""Receive a Message
|
|
|
|
This receives the next pending message from the socket. This operation
|
|
is synchronous.
|
|
|
|
A tuple consisting of the deserialized message payload, the auxiliary
|
|
file-descriptor set, and the socket-address of the sender is returned.
|
|
|
|
In case the peer closed the connection, A tuple of `None` values is
|
|
returned.
|
|
"""
|
|
|
|
# On `SOCK_SEQPACKET`, packets might be arbitrarily sized. There is no
|
|
# hard-coded upper limit, since it is only restricted by the size of
|
|
# the kernel write buffer on sockets (which itself can be modified via
|
|
# sysctl). The only real maximum is probably something like 2^31-1,
|
|
# since that is the maximum of that sysctl datatype.
|
|
# Anyway, `MSG_TRUNC+MSG_PEEK` usually allows us to easily peek at the
|
|
# incoming buffer. Unfortunately, the python `recvmsg()` wrapper
|
|
# discards the return code and we cannot use that. Instead, we simply
|
|
# loop until we know the size. This is slightly awkward, but seems fine
|
|
# as long as you do not put this into a hot-path.
|
|
size = 4096
|
|
while True:
|
|
peek = self._socket.recvmsg(size, 0, socket.MSG_PEEK)
|
|
if not peek[0]:
|
|
# Connection was closed
|
|
return None, None, None
|
|
if not (peek[2] & socket.MSG_TRUNC):
|
|
break
|
|
size *= 2
|
|
|
|
# Fetch a packet from the socket. On linux, the maximum SCM_RIGHTS array
|
|
# size is hard-coded to 253. This allows us to size the ancillary buffer
|
|
# big enough to receive any possible message.
|
|
fds = array.array("i")
|
|
msg = self._socket.recvmsg(size, socket.CMSG_LEN(253 * fds.itemsize))
|
|
|
|
# First thing we do is always to fetch the CMSG FDs into an FdSet. This
|
|
# guarantees that we do not leak FDs in case the message handling fails
|
|
# for other reasons.
|
|
for level, ty, data in msg[1]:
|
|
if level == socket.SOL_SOCKET and ty == socket.SCM_RIGHTS:
|
|
assert len(data) % fds.itemsize == 0
|
|
fds.frombytes(data)
|
|
fdset = FdSet(rawfds=fds)
|
|
|
|
# Check the returned message flags. If the message was truncated, we
|
|
# have to discard it. This shouldn't happen, but there is no harm in
|
|
# handling it. However, `CTRUNC` can happen, since it is also triggered
|
|
# when LSMs reject FD transmission. Treat it the same as a parser error.
|
|
flags = msg[2]
|
|
if flags & (socket.MSG_TRUNC | socket.MSG_CTRUNC):
|
|
raise BufferError
|
|
|
|
try:
|
|
payload = json.loads(msg[0])
|
|
except json.JSONDecodeError as e:
|
|
raise BufferError from e
|
|
|
|
return (payload, fdset, msg[3])
|
|
|
|
def send(self, payload: object, *, fds: Optional[list] = None):
|
|
"""Send Message
|
|
|
|
Send a new message via this socket. This operation is synchronous. The
|
|
maximum message size depends on the configured send-buffer on the
|
|
socket. An `OSError` with `EMSGSIZE` is raised when it is exceeded.
|
|
|
|
Parameters
|
|
----------
|
|
payload
|
|
A python object to serialize as JSON and send via this socket. See
|
|
`json.dump()` for details about the serialization involved.
|
|
destination
|
|
The destination to send to. If `None`, the default destination is
|
|
used (if none is set, this will raise an `OSError`).
|
|
fds
|
|
A list of file-descriptors to send with the message.
|
|
|
|
Raises
|
|
------
|
|
OSError
|
|
If the socket cannot be written, a matching `OSError` is raised.
|
|
TypeError
|
|
If the payload cannot be serialized, a type error is raised.
|
|
"""
|
|
|
|
if not self._socket:
|
|
raise RuntimeError("Tried to send without socket.")
|
|
|
|
serialized = json.dumps(payload).encode()
|
|
cmsg = []
|
|
if fds:
|
|
cmsg.append((socket.SOL_SOCKET, socket.SCM_RIGHTS, array.array("i", fds)))
|
|
|
|
n = self._socket.sendmsg([serialized], cmsg, 0)
|
|
assert n == len(serialized)
|
|
|
|
def send_and_recv(self, payload: object, *, fds: Optional[list] = None):
|
|
"""Send a message and wait for a reply
|
|
|
|
This is a convenience helper that combines `send` and `recv`.
|
|
See the individual methods for details about the parameters.
|
|
"""
|
|
|
|
self.send(payload, fds=fds)
|
|
return self.recv()
|