async def init_conn(bridge, iface=""):
*sarg, _, rbridge = (await socket.getaddrinfo(*bridge,
type=socket.SOCK_STREAM))[0]
retry = 3
for _ in range(retry):
sa = socket.socket(*sarg)
sa.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
await sa.connect(rbridge)
priv_host, priv_port = sa.getsockname() # self, private
if iface:
priv_host, priv_port = get_priv_host(iface), 0
try:
s_cona = await bound_socket(
(priv_host, priv_port)) # for connect, a
if iface:
_, priv_port = s_cona.getsockname()
s_conb = await bound_socket(
(priv_host, priv_port)) # for connect, b
s_acp = await bound_socket(('', priv_port)) # for accept
break
except OSError: # NOTE: albeit SO_REUSEADDR set, bind still occassionally fails
logger.debug(
f"cannot bind to private addr {priv_host}:{priv_port}, retry")
sa.close()
await trio.sleep(0.1)
continue
else:
raise OSError("failed to bind to a local address")
logger.debug(f"self's private addr {priv_host}:{priv_port}")
return trio.SocketStream(
sa), f"{priv_host}:{priv_port}", s_acp, s_cona, s_conb
async def make_socket(self, af, socktype, proto=0, source=None,
destination=None, timeout=None,
ssl_context=None, server_hostname=None):
s = trio.socket.socket(af, socktype, proto)
stream = None
try:
if source:
await s.bind(_lltuple(source, af))
if socktype == socket.SOCK_STREAM:
with _maybe_timeout(timeout):
await s.connect(_lltuple(destination, af))
except Exception: # pragma: no cover
s.close()
raise
if socktype == socket.SOCK_DGRAM:
return DatagramSocket(s)
elif socktype == socket.SOCK_STREAM:
stream = trio.SocketStream(s)
s = None
tls = False
if ssl_context:
tls = True
try:
stream = trio.SSLStream(stream, ssl_context,
server_hostname=server_hostname)
except Exception: # pragma: no cover
await stream.aclose()
raise
return StreamSocket(af, stream, tls)
raise NotImplementedError('unsupported socket ' +
f'type {socktype}') # pragma: no cover
async def open_unix_socket(filename):
"""Opens a connection to the specified
`Unix domain socket <https://en.wikipedia.org/wiki/Unix_domain_socket>`__.
You must have read/write permission on the specified file to connect.
Args:
filename (str or bytes): The filename to open the connection to.
Returns:
SocketStream: a :class:`~trio.abc.Stream` connected to the given file.
Raises:
OSError: If the socket file could not be connected to.
RuntimeError: If AF_UNIX sockets are not supported.
"""
if not has_unix:
raise RuntimeError("Unix sockets are not supported on this platform")
# much more simplified logic vs tcp sockets - one socket type and only one
# possible location to connect to
sock = socket(AF_UNIX, SOCK_STREAM)
with close_on_error(sock):
await sock.connect(os.fspath(filename))
return trio.SocketStream(sock)
async def _get(data):
result = None
async def sender(client_sock, data):
# print("sender: started!")
# print("sender: sending {!r}".format(data))
print('sent:', data.strip())
await client_sock.send_all(bytes(data, ENCODING))
async def receiver(client_sock:trio.SocketStream):
_data = await client_sock.receive_some(BUFSIZE)
expcols, exprows = stream_size(_data)
if exprows and expcols:
while not _data.endswith(b'\r\n\r\n'):
_data += await client_sock.receive_some(BUFSIZE)
if not data:
sys.exit()
nonlocal result
result = _data
print('received:', str(_data, ENCODING).strip())
with trio.socket.socket() as client_sock:
await client_sock.connect((SERVER_ADDRESS, PORT))
client_sock = trio.SocketStream(client_sock)
async with trio.open_nursery() as nursery:
nursery.start_soon(sender, client_sock, data)
nursery.start_soon(receiver, client_sock)
return result
async def _get(data):
result = None
async def sender(client_sock, data):
# print("sender: started!")
# print("sender: sending {!r}".format(data))
await client_sock.send_all(bytes(data, ENCODING))
async def receiver(client_sock):
# print("receiver: started!")
_data = await client_sock.receive_some(BUFSIZE)
# print("receiver: got data {!r}".format(_data))
if not data:
# print("receiver: connection closed")
sys.exit()
nonlocal result
result = _data
print('received:',_data)
# print("parent: connecting to 127.0.0.1:{}".format(PORT))
with trio.socket.socket() as client_sock:
await client_sock.connect((SERVER_ADDRESS, PORT))
client_sock = trio.SocketStream(client_sock)
async with trio.open_nursery() as nursery:
# print("parent: spawning sender...")
nursery.start_soon(sender, client_sock, data)
# print("parent: spawning receiver...")
nursery.start_soon(receiver, client_sock)
return result
async def startup_socket(self):
if not self._socket:
self._socket = trio.socket.socket()
await self._socket.connect((alsdac.SERVER_ADDRESS, alsdac.PORT))
self._socket_stream = trio.SocketStream(self._socket)
self._socket_stream.setsockopt(socket.SOL_SOCKET,
socket.SO_KEEPALIVE, 1)
self._socket_stream.setsockopt(socket.IPPROTO_TCP,
socket.TCP_KEEPIDLE, 1)
self._socket_stream.setsockopt(socket.IPPROTO_TCP,
socket.TCP_KEEPINTVL, 1)
self._socket_stream.setsockopt(socket.IPPROTO_TCP,
socket.TCP_KEEPCNT, 10)
self._socket_stream.setsockopt(socket.SOL_SOCKET,
socket.SO_REUSEADDR, 1)
async def make_request(proxy: AsyncProxy,
url: str,
resolve_host=False,
timeout=None):
url = URL(url)
dest_host = url.host
if resolve_host:
resolver = Resolver()
_, dest_host = await resolver.resolve(url.host)
sock: socket.socket = await proxy.connect(dest_host=dest_host,
dest_port=url.port,
timeout=timeout)
ssl_context = None
if url.scheme == 'https':
ssl_context = ssl.create_default_context()
stream = trio.SocketStream(sock)
if ssl_context is not None:
stream = trio.SSLStream(stream, ssl_context, server_hostname=url.host)
await stream.do_handshake()
request = ('GET {rel_url} HTTP/1.1\r\n'
'Host: {host}\r\n'
'Connection: close\r\n\r\n')
request = request.format(rel_url=url.path_qs, host=url.host)
request = request.encode('ascii')
await stream.send_all(request)
response = await stream.receive_some(1024)
status_line = response.split(b'\r\n', 1)[0]
status_line = status_line.decode('utf-8', 'surrogateescape')
version, status_code, *reason = status_line.split()
return int(status_code)
async def test_defer_connect(self, set_me_up):
await set_me_up(self)
for db in self.databases:
d = db.copy()
try:
sock = trio.socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
await sock.connect(d['unix_socket'])
sock = trio.SocketStream(sock)
except KeyError:
sock = await trio.open_tcp_stream \
(d.get('host', 'localhost'), d.get('port', 3306))
for k in ['unix_socket', 'host', 'port']:
try:
del d[k]
except KeyError:
pass
c = trio_mysql.connect(**d)
self.assertFalse(c.open)
await c.connect(sock)
await c.aclose()
await sock.aclose()
async def hole_punching(peer, s_acp, s_cona, s_conb):
peer_publ_addr, peer_priv_addr = peer
peer_publ_addr, peer_priv_addr = parse_addr(peer_publ_addr), parse_addr(
peer_priv_addr)
logger.debug(
f"peer's public addr {peer_publ_addr}, private addr {peer_priv_addr}")
win_s, win_r = trio.open_memory_channel(1)
async with trio.open_nursery() as _n:
_n.start_soon(accept, s_acp, win_s)
_n.start_soon(connect, s_cona, peer_publ_addr, win_s)
if peer_priv_addr != peer_publ_addr:
_n.start_soon(connect, s_conb, peer_priv_addr, win_s)
else:
s_conb.close()
conn = None
with trio.move_on_after(RENDEZVOUS_TIMEOUT):
conn = trio.SocketStream(await win_r.receive())
if conn is None:
raise OSError(
f"failed to established a P2P connection with {peer}")
_n.cancel_scope.cancel()
return conn
async def open_tcp_stream(
host,
port,
*,
# No trailing comma b/c bpo-9232 (fixed in py36)
happy_eyeballs_delay=DEFAULT_DELAY,
):
"""Connect to the given host and port over TCP.
If the given ``host`` has multiple IP addresses associated with it, then
we have a problem: which one do we use?
One approach would be to attempt to connect to the first one, and then if
that fails, attempt to connect to the second one ... until we've tried all
of them. But the problem with this is that if the first IP address is
unreachable (for example, because it's an IPv6 address and our network
discards IPv6 packets), then we might end up waiting tens of seconds for
the first connection attempt to timeout before we try the second address.
Another approach would be to attempt to connect to all of the addresses at
the same time, in parallel, and then use whichever connection succeeds
first, abandoning the others. This would be fast, but create a lot of
unnecessary load on the network and the remote server.
This function strikes a balance between these two extremes: it works its
way through the available addresses one at a time, like the first
approach; but, if ``happy_eyeballs_delay`` seconds have passed and it's
still waiting for an attempt to succeed or fail, then it gets impatient
and starts the next connection attempt in parallel. As soon as any one
connection attempt succeeds, all the other attempts are cancelled. This
avoids unnecessary load because most connections will succeed after just
one or two attempts, but if one of the addresses is unreachable then it
doesn't slow us down too much.
This is known as a "happy eyeballs" algorithm, and our particular variant
is modelled after how Chrome connects to webservers; see `RFC 6555
<https://tools.ietf.org/html/rfc6555>`__ for more details.
Args:
host (str or bytes): The host to connect to. Can be an IPv4 address,
IPv6 address, or a hostname.
port (int): The port to connect to.
happy_eyeballs_delay (float): How many seconds to wait for each
connection attempt to succeed or fail before getting impatient and
starting another one in parallel. Set to `math.inf` if you want
to limit to only one connection attempt at a time (like
:func:`socket.create_connection`). Default: 0.25 (250 ms).
Returns:
SocketStream: a :class:`~trio.abc.Stream` connected to the given server.
Raises:
OSError: if the connection fails.
See also:
open_ssl_over_tcp_stream
"""
# To keep our public API surface smaller, rule out some cases that
# getaddrinfo will accept in some circumstances, but that act weird or
# have non-portable behavior or are just plain not useful.
# No type check on host though b/c we want to allow bytes-likes.
if host is None:
raise ValueError("host cannot be None")
if not isinstance(port, int):
raise TypeError("port must be int, not {!r}".format(port))
if happy_eyeballs_delay is None:
happy_eyeballs_delay = DEFAULT_DELAY
targets = await getaddrinfo(host, port, type=SOCK_STREAM)
# I don't think this can actually happen -- if there are no results,
# getaddrinfo should have raised OSError instead of returning an empty
# list. But let's be paranoid and handle it anyway:
if not targets:
msg = "no results found for hostname lookup: {}".format(
format_host_port(host, port))
raise OSError(msg)
reorder_for_rfc_6555_section_5_4(targets)
# This list records all the connection failures that we ignored.
oserrors = []
# Keeps track of the socket that we're going to complete with,
# need to make sure this isn't automatically closed
winning_socket = None
# Try connecting to the specified address. Possible outcomes:
# - success: record connected socket in winning_socket and cancel
# concurrent attempts
# - failure: record exception in oserrors, set attempt_failed allowing
# the next connection attempt to start early
# code needs to ensure sockets can be closed appropriately in the
# face of crash or cancellation
async def attempt_connect(socket_args, sockaddr, attempt_failed):
nonlocal winning_socket
try:
sock = socket(*socket_args)
open_sockets.add(sock)
await sock.connect(sockaddr)
# Success! Save the winning socket and cancel all outstanding
# connection attempts.
winning_socket = sock
nursery.cancel_scope.cancel()
except OSError as exc:
# This connection attempt failed, but the next one might
# succeed. Save the error for later so we can report it if
# everything fails, and tell the next attempt that it should go
# ahead (if it hasn't already).
oserrors.append(exc)
attempt_failed.set()
with close_all() as open_sockets:
# nursery spawns a task for each connection attempt, will be
# cancelled by the task that gets a successful connection
async with trio.open_nursery() as nursery:
for *sa, _, addr in targets:
# create an event to indicate connection failure,
# allowing the next target to be tried early
attempt_failed = trio.Event()
nursery.start_soon(attempt_connect, sa, addr, attempt_failed)
# give this attempt at most this time before moving on
with trio.move_on_after(happy_eyeballs_delay):
await attempt_failed.wait()
# nothing succeeded
if winning_socket is None:
assert len(oserrors) == len(targets)
msg = "all attempts to connect to {} failed".format(
format_host_port(host, port))
raise OSError(msg) from trio.MultiError(oserrors)
else:
stream = trio.SocketStream(winning_socket)
open_sockets.remove(winning_socket)
return stream
async def open_tcp_stream(
host,
port,
*,
# No trailing comma b/c bpo-9232 (fixed in py36)
happy_eyeballs_delay=DEFAULT_DELAY):
"""Connect to the given host and port over TCP.
If the given ``host`` has multiple IP addresses associated with it, then
we have a problem: which one do we use?
One approach would be to attempt to connect to the first one, and then if
that fails, attempt to connect to the second one ... until we've tried all
of them. But the problem with this is that if the first IP address is
unreachable (for example, because it's an IPv6 address and our network
discards IPv6 packets), then we might end up waiting tens of seconds for
the first connection attempt to timeout before we try the second address.
Another approach would be to attempt to connect to all of the addresses at
the same time, in parallel, and then use whichever connection succeeds
first, abandoning the others. This would be fast, but create a lot of
unnecessary load on the network and the remote server.
This function strikes a balance between these two extremes: it works its
way through the available addresses one at a time, like the first
approach; but, if ``happy_eyeballs_delay`` seconds have passed and it's
still waiting for an attempt to succeed or fail, then it gets impatient
and starts the next connection attempt in parallel. As soon as any one
connection attempt succeeds, all the other attempts are cancelled. This
avoids unnecessary load because most connections will succeed after just
one or two attempts, but if one of the addresses is unreachable then it
doesn't slow us down too much.
This is known as a "happy eyeballs" algorithm, and our particular variant
is modelled after how Chrome connects to webservers; see `RFC 6555
<https://tools.ietf.org/html/rfc6555>`__ for more details.
Args:
host (str or bytes): The host to connect to. Can be an IPv4 address,
IPv6 address, or a hostname.
port (int): The port to connect to.
happy_eyeballs_delay (float): How many seconds to wait for each
connection attempt to succeed or fail before getting impatient and
starting another one in parallel. Set to :obj:`math.inf` if you want
to limit to only one connection attempt at a time (like
:func:`socket.create_connection`). Default: 0.3 (300 ms).
Returns:
SocketStream: a :class:`~trio.abc.Stream` connected to the given server.
Raises:
OSError: if the connection fails.
See also:
open_ssl_over_tcp_stream
"""
# To keep our public API surface smaller, rule out some cases that
# getaddrinfo will accept in some circumstances, but that act weird or
# have non-portable behavior or are just plain not useful.
# No type check on host though b/c we want to allow bytes-likes.
if host is None:
raise ValueError("host cannot be None")
if not isinstance(port, int):
raise TypeError("port must be int, not {!r}".format(port))
if happy_eyeballs_delay is None:
happy_eyeballs_delay = DEFAULT_DELAY
targets = await getaddrinfo(host, port, type=SOCK_STREAM)
# I don't think this can actually happen -- if there are no results,
# getaddrinfo should have raised OSError instead of returning an empty
# list. But let's be paranoid and handle it anyway:
if not targets:
msg = "no results found for hostname lookup: {}".format(
format_host_port(host, port))
raise OSError(msg)
reorder_for_rfc_6555_section_5_4(targets)
targets_iter = iter(targets)
# This list records all the connection failures that we ignored.
oserrors = []
# It's possible for multiple connection attempts to succeed at the ~same
# time; this list records all successful connections.
winning_sockets = []
# Sleep for the given amount of time, then kick off the next task and
# start a connection attempt. On failure, expedite the next task; on
# success, kill everything. Possible outcomes:
# - records a failure in oserrors, returns None
# - records a connected socket in winning_sockets, returns None
# - crash (raises an unexpected exception)
async def attempt_connect(nursery, previous_attempt_failed):
# Wait until either the previous attempt failed, or the timeout
# expires (unless this is the first invocation, in which case we just
# go ahead).
if previous_attempt_failed is not None:
with trio.move_on_after(happy_eyeballs_delay):
await previous_attempt_failed.wait()
# Claim our target.
try:
*socket_args, _, target_sockaddr = next(targets_iter)
except StopIteration:
return
# Then kick off the next attempt.
this_attempt_failed = trio.Event()
nursery.start_soon(attempt_connect, nursery, this_attempt_failed)
# Then make this invocation's attempt
try:
with close_on_error(socket(*socket_args)) as sock:
await sock.connect(target_sockaddr)
except OSError as exc:
# This connection attempt failed, but the next one might
# succeed. Save the error for later so we can report it if
# everything fails, and tell the next attempt that it should go
# ahead (if it hasn't already).
oserrors.append(exc)
this_attempt_failed.set()
else:
# Success! Save the winning socket and cancel all outstanding
# connection attempts.
winning_sockets.append(sock)
nursery.cancel_scope.cancel()
# Kick off the chain of connection attempts.
async with trio.open_nursery() as nursery:
nursery.start_soon(attempt_connect, nursery, None)
# All connection attempts complete, and no unexpected errors escaped. So
# at this point the oserrors and winning_sockets lists are filled in.
if winning_sockets:
first_prize = winning_sockets.pop(0)
for sock in winning_sockets:
sock.close()
return trio.SocketStream(first_prize)
else:
assert len(oserrors) == len(targets)
msg = "all attempts to connect to {} failed".format(
format_host_port(host, port))
raise OSError(msg) from trio.MultiError(oserrors)
async def open_tcp_stream(
host,
port,
*,
happy_eyeballs_delay=DEFAULT_DELAY,
local_address=None,
):
"""Connect to the given host and port over TCP.
If the given ``host`` has multiple IP addresses associated with it, then
we have a problem: which one do we use?
One approach would be to attempt to connect to the first one, and then if
that fails, attempt to connect to the second one ... until we've tried all
of them. But the problem with this is that if the first IP address is
unreachable (for example, because it's an IPv6 address and our network
discards IPv6 packets), then we might end up waiting tens of seconds for
the first connection attempt to timeout before we try the second address.
Another approach would be to attempt to connect to all of the addresses at
the same time, in parallel, and then use whichever connection succeeds
first, abandoning the others. This would be fast, but create a lot of
unnecessary load on the network and the remote server.
This function strikes a balance between these two extremes: it works its
way through the available addresses one at a time, like the first
approach; but, if ``happy_eyeballs_delay`` seconds have passed and it's
still waiting for an attempt to succeed or fail, then it gets impatient
and starts the next connection attempt in parallel. As soon as any one
connection attempt succeeds, all the other attempts are cancelled. This
avoids unnecessary load because most connections will succeed after just
one or two attempts, but if one of the addresses is unreachable then it
doesn't slow us down too much.
This is known as a "happy eyeballs" algorithm, and our particular variant
is modelled after how Chrome connects to webservers; see `RFC 6555
<https://tools.ietf.org/html/rfc6555>`__ for more details.
Args:
host (str or bytes): The host to connect to. Can be an IPv4 address,
IPv6 address, or a hostname.
port (int): The port to connect to.
happy_eyeballs_delay (float): How many seconds to wait for each
connection attempt to succeed or fail before getting impatient and
starting another one in parallel. Set to `math.inf` if you want
to limit to only one connection attempt at a time (like
:func:`socket.create_connection`). Default: 0.25 (250 ms).
local_address (None or str): The local IP address or hostname to use as
the source for outgoing connections. If ``None``, we let the OS pick
the source IP.
This is useful in some exotic networking configurations where your
host has multiple IP addresses, and you want to force the use of a
specific one.
Note that if you pass an IPv4 ``local_address``, then you won't be
able to connect to IPv6 hosts, and vice-versa. If you want to take
advantage of this to force the use of IPv4 or IPv6 without
specifying an exact source address, you can use the IPv4 wildcard
address ``local_address="0.0.0.0"``, or the IPv6 wildcard address
``local_address="::"``.
Returns:
SocketStream: a :class:`~trio.abc.Stream` connected to the given server.
Raises:
OSError: if the connection fails.
See also:
open_ssl_over_tcp_stream
"""
# To keep our public API surface smaller, rule out some cases that
# getaddrinfo will accept in some circumstances, but that act weird or
# have non-portable behavior or are just plain not useful.
# No type check on host though b/c we want to allow bytes-likes.
if host is None:
raise ValueError("host cannot be None")
if not isinstance(port, int):
raise TypeError("port must be int, not {!r}".format(port))
if happy_eyeballs_delay is None:
happy_eyeballs_delay = DEFAULT_DELAY
targets = await getaddrinfo(host, port, type=SOCK_STREAM)
# I don't think this can actually happen -- if there are no results,
# getaddrinfo should have raised OSError instead of returning an empty
# list. But let's be paranoid and handle it anyway:
if not targets:
msg = "no results found for hostname lookup: {}".format(
format_host_port(host, port))
raise OSError(msg)
reorder_for_rfc_6555_section_5_4(targets)
# This list records all the connection failures that we ignored.
oserrors = []
# Keeps track of the socket that we're going to complete with,
# need to make sure this isn't automatically closed
winning_socket = None
# Try connecting to the specified address. Possible outcomes:
# - success: record connected socket in winning_socket and cancel
# concurrent attempts
# - failure: record exception in oserrors, set attempt_failed allowing
# the next connection attempt to start early
# code needs to ensure sockets can be closed appropriately in the
# face of crash or cancellation
async def attempt_connect(socket_args, sockaddr, attempt_failed):
nonlocal winning_socket
try:
sock = socket(*socket_args)
open_sockets.add(sock)
if local_address is not None:
# TCP connections are identified by a 4-tuple:
#
# (local IP, local port, remote IP, remote port)
#
# So if a single local IP wants to make multiple connections
# to the same (remote IP, remote port) pair, then those
# connections have to use different local ports, or else TCP
# won't be able to tell them apart. OTOH, if you have multiple
# connections to different remote IP/ports, then those
# connections can share a local port.
#
# Normally, when you call bind(), the kernel will immediately
# assign a specific local port to your socket. At this point
# the kernel doesn't know which (remote IP, remote port)
# you're going to use, so it has to pick a local port that
# *no* other connection is using. That's the only way to
# guarantee that this local port will be usable later when we
# call connect(). (Alternatively, you can set SO_REUSEADDR to
# allow multiple nascent connections to share the same port,
# but then connect() might fail with EADDRNOTAVAIL if we get
# unlucky and our TCP 4-tuple ends up colliding with another
# unrelated connection.)
#
# So calling bind() before connect() works, but it disables
# sharing of local ports. This is inefficient: it makes you
# more likely to run out of local ports.
#
# But on some versions of Linux, we can re-enable sharing of
# local ports by setting a special flag. This flag tells
# bind() to only bind the IP, and not the port. That way,
# connect() is allowed to pick the the port, and it can do a
# better job of it because it knows the remote IP/port.
try:
sock.setsockopt(trio.socket.IPPROTO_IP,
trio.socket.IP_BIND_ADDRESS_NO_PORT, 1)
except (OSError, AttributeError):
pass
try:
await sock.bind((local_address, 0))
except OSError:
raise OSError(
f"local_address={local_address!r} is incompatible "
f"with remote address {sockaddr}")
await sock.connect(sockaddr)
# Success! Save the winning socket and cancel all outstanding
# connection attempts.
winning_socket = sock
nursery.cancel_scope.cancel()
except OSError as exc:
# This connection attempt failed, but the next one might
# succeed. Save the error for later so we can report it if
# everything fails, and tell the next attempt that it should go
# ahead (if it hasn't already).
oserrors.append(exc)
attempt_failed.set()
with close_all() as open_sockets:
# nursery spawns a task for each connection attempt, will be
# cancelled by the task that gets a successful connection
async with trio.open_nursery() as nursery:
for *sa, _, addr in targets:
# create an event to indicate connection failure,
# allowing the next target to be tried early
attempt_failed = trio.Event()
nursery.start_soon(attempt_connect, sa, addr, attempt_failed)
# give this attempt at most this time before moving on
with trio.move_on_after(happy_eyeballs_delay):
await attempt_failed.wait()
# nothing succeeded
if winning_socket is None:
assert len(oserrors) == len(targets)
msg = "all attempts to connect to {} failed".format(
format_host_port(host, port))
raise OSError(msg) from trio.MultiError(oserrors)
else:
stream = trio.SocketStream(winning_socket)
open_sockets.remove(winning_socket)
return stream
async def run_with_trio(
executable: pathlib.PosixPath,
argv: List[bytes],
environ: Dict[bytes, bytes],
stdin: Tuple[io.FileIO, io.FileIO],
stdout: Tuple[io.FileIO, io.FileIO],
stderr: Tuple[io.FileIO, io.FileIO],
server: socket.socket,
reply: bytes,
wmem_default: int,
optmem_max: int,
uid: int,
gid: int,
) -> Result:
stdout_content = io.BytesIO()
stderr_content = io.BytesIO()
messages = None
sent = None
with trio.move_on_after(TIMEOUT):
process = await trio.open_process(
argv,
executable=executable,
env=environ,
bufsize=0,
text=False,
encoding=None,
stdin=os.dup(stdin[0].fileno()),
stdout=os.dup(stdout[1].fileno()),
stderr=os.dup(stderr[1].fileno()),
)
try:
async with trio.open_nursery() as nursery:
nursery: trio.Nursery = nursery
# Read stdout/stderr in the background
nursery.start_soon(track_process, process)
nursery.start_soon(
read_pipe,
trio.lowlevel.FdStream(os.dup(stdout[0].fileno())),
stdout_content)
nursery.start_soon(
read_pipe,
trio.lowlevel.FdStream(os.dup(stderr[0].fileno())),
stderr_content)
server = trio.socket.from_stdlib_socket(server)
client, _ = await server.accept()
client = trio.SocketStream(client)
credentials = ucred.unpack(
client.getsockopt(
socket.SOL_SOCKET,
socket.SO_PEERCRED,
ucred.size,
))
assert (process.pid, uid, gid) == credentials
messages = []
await receive_messages(client, wmem_default, optmem_max,
messages)
sent = await send_reply(client, reply)
await client.send_eof()
except ChildStopped:
pass
else:
process.kill()
drain_pipe(stdout[0], stdout_content)
drain_pipe(stderr[0], stderr_content)
return process.returncode, stdout_content.getvalue(
), stderr_content.getvalue(), messages, sent
async def stream(q,
where,
tls=False,
port=None,
source=None,
source_port=0,
one_rr_per_rrset=False,
ignore_trailing=False,
ssl_context=None,
server_hostname=None):
"""Return the response obtained after sending a query using TCP or TLS.
*q*, a ``dns.message.Message``, the query to send.
*where*, a ``str`` containing an IPv4 or IPv6 address, where
to send the message.
*tls*, a ``bool``. If ``False``, the default, the query will be
sent using TCP and *port* will default to 53. If ``True``, the
query is sent using TLS, and *port* will default to 853.
*port*, an ``int``, the port send the message to. The default is as
specified in the description for *tls*.
*source*, a ``str`` containing an IPv4 or IPv6 address, specifying
the source address. The default is the wildcard address.
*source_port*, an ``int``, the port from which to send the message.
The default is 0.
*one_rr_per_rrset*, a ``bool``. If ``True``, put each RR into its own
RRset.
*ignore_trailing*, a ``bool``. If ``True``, ignore trailing
junk at end of the received message.
*ssl_context*, an ``ssl.SSLContext``, the context to use when establishing
a TLS connection. If ``None``, the default, creates one with the default
configuration. If this value is not ``None``, then the *tls* parameter
is treated as if it were ``True`` regardless of its value.
*server_hostname*, a ``str`` containing the server's hostname. The
default is ``None``, which means that no hostname is known, and if an
SSL context is created, hostname checking will be disabled.
Returns a ``dns.message.Message``.
"""
if ssl_context is not None:
tls = True
if port is None:
if tls:
port = 853
else:
port = 53
wire = q.to_wire()
(af, destination, source) = \
dns.query._destination_and_source(None, where, port, source,
source_port)
with socket_factory(af, socket.SOCK_STREAM, 0) as s:
begin_time = time.time()
if source is not None:
await s.bind(source)
await s.connect(destination)
stream = trio.SocketStream(s)
if tls and ssl_context is None:
ssl_context = ssl.create_default_context()
if server_hostname is None:
ssl_context.check_hostname = False
if ssl_context:
stream = trio.SSLStream(stream,
ssl_context,
server_hostname=server_hostname)
async with stream:
await send_stream(stream, wire)
(r, received_time) = await receive_stream(stream, one_rr_per_rrset,
q.keyring, q.mac,
ignore_trailing)
if not q.is_response(r):
raise BadResponse
r.time = received_time - begin_time
return r
async def interact(
*,
prompt: str = "> ",
greeting: str = "",
terminal: Optional[trio.abc.Stream] = None,
use_pty: Optional[bool] = None,
completer: Callable[[str, str, int], Sequence[str]] = lambda x, y, z: (),
) -> AsyncIterator[trio.abc.ReceiveChannel[
AsyncContextManager[InteractiveRequest]]]:
async def proxy_one_way(source: trio.abc.Stream,
sink: trio.abc.Stream) -> None:
while True:
chunk = await source.receive_some(4096)
if not chunk:
break
await sink.send_all(chunk)
async with AsyncExitStack() as stack, trio.open_nursery() as nursery:
preexec_fn: Optional[Callable[[], None]] = None
if use_pty is None and terminal is None:
use_pty = os.isatty(0) or os.isatty(1)
if use_pty:
if terminal is not None:
our_fd, their_fd = pty.openpty()
our_stream = await stack.enter_async_context(
trio.hazmat.FdStream(our_fd))
nursery.start_soon(proxy_one_way, terminal, our_stream)
nursery.start_soon(proxy_one_way, our_stream, terminal)
else:
try:
their_fd = os.open("/dev/tty", os.O_RDWR)
except OSError:
raise RuntimeError(
"You must pass a 'terminal' argument when running "
"with no controlling terminal and use_pty=True.")
preexec_fn = os.setpgrp
stack.callback(os.close, their_fd)
else:
our_sock, their_sock = socket.socketpair()
their_fd = their_sock.detach()
our_stream = await stack.enter_async_context(
trio.SocketStream(trio.socket.from_stdlib_socket(our_sock)))
if terminal is None:
in_ = trio.hazmat.FdStream(
os.open("/proc/self/fd/0", os.O_RDONLY))
out = trio.hazmat.FdStream(
os.open("/proc/self/fd/1", os.O_WRONLY))
terminal = await stack.enter_async_context(
trio.StapledStream(out, in_))
assert terminal is not None
nursery.start_soon(proxy_one_way, terminal, our_stream)
nursery.start_soon(proxy_one_way, our_stream, terminal)
with trio.socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as ssock:
with tempfile.TemporaryDirectory(
prefix="trio-readline.") as tmpdir:
sockname = os.path.join(tmpdir, "control.sock")
await ssock.bind(sockname)
ssock.listen(1)
process = await nursery.start( # type: ignore
partial(
run_process,
[sys.executable, "-m", rlproxy.__name__, sockname],
stdin=their_fd,
stdout=their_fd,
preexec_fn=preexec_fn,
),
name="<rlproxy process>",
)
# Make the readline-proxy process part of the foreground process
# group so that it gets Ctrl+C rather than us.
os.tcsetpgrp(their_fd, process.pid)
sock, _ = await ssock.accept()
control = await stack.enter_async_context(
trio.SocketStream(sock))
if terminal is None and use_pty:
@stack.callback
def restore_pgrp() -> None:
old_handler = signal.getsignal(signal.SIGTTOU)
signal.signal(signal.SIGTTOU, signal.SIG_IGN)
try:
# This will raise SIGTTOU if we're not in the
# foreground process group, which we have to ignore
os.tcsetpgrp(their_fd, os.getpgrp())
finally:
signal.signal(signal.SIGTTOU, old_handler)
async def send_control(words: Sequence[str]) -> None:
await control.send_all(json.dumps(words).encode("utf-8") + b"\n")
send_channel, receive_channel = trio.open_memory_channel[
AsyncContextManager[InteractiveRequest]](1)
active_request_scopes = MultiCancelScope()
async def handle_control() -> None:
nonlocal active_request_scopes
try:
async for line in TextReceiveStream(control, encoding="utf-8"):
request = json.loads(line)
if request[0] == "complete":
await send_control(
completer(request[1], request[2], int(request[3])))
elif request[0] == "run":
await send_channel.send(
handle_interactive_request(
request[1],
send_control,
active_request_scopes.open_child(),
))
elif request[0] == "eof":
await send_channel.send(
handle_interactive_request(
"quit", send_control,
active_request_scopes.open_child()))
elif request[0] == "interrupt":
active_request_scopes.cancel()
active_request_scopes = MultiCancelScope()
else:
raise ValueError(
f"readline proxy sent invalid request {request!r}")
except trio.BrokenResourceError:
pass
finally:
await send_channel.aclose()
await send_control(["prompt", prompt])
if greeting:
if greeting[-1] != "\n":
greeting += "\n"
await send_control(["output", greeting])
await send_control(["done"])
nursery.start_soon(handle_control)
async with receive_channel:
yield receive_channel
nursery.cancel_scope.cancel()
