# This should eventually be cleaned up and become public, but for right now I'm just
|
# implementing enough to test DTLS.
|
|
# TODO:
|
# - user-defined routers
|
# - TCP
|
# - UDP broadcast
|
|
import trio
|
import attr
|
import ipaddress
|
from collections import deque
|
import errno
|
import os
|
from typing import Union, List, Optional
|
import enum
|
from contextlib import contextmanager
|
|
from trio._util import Final, NoPublicConstructor
|
|
IPAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
|
|
|
def _family_for(ip: IPAddress) -> int:
|
if isinstance(ip, ipaddress.IPv4Address):
|
return trio.socket.AF_INET
|
elif isinstance(ip, ipaddress.IPv6Address):
|
return trio.socket.AF_INET6
|
assert False # pragma: no cover
|
|
|
def _wildcard_ip_for(family: int) -> IPAddress:
|
if family == trio.socket.AF_INET:
|
return ipaddress.ip_address("0.0.0.0")
|
elif family == trio.socket.AF_INET6:
|
return ipaddress.ip_address("::")
|
else:
|
assert False
|
|
|
def _localhost_ip_for(family: int) -> IPAddress:
|
if family == trio.socket.AF_INET:
|
return ipaddress.ip_address("127.0.0.1")
|
elif family == trio.socket.AF_INET6:
|
return ipaddress.ip_address("::1")
|
else:
|
assert False
|
|
|
def _fake_err(code):
|
raise OSError(code, os.strerror(code))
|
|
|
def _scatter(data, buffers):
|
written = 0
|
for buf in buffers:
|
next_piece = data[written : written + len(buf)]
|
with memoryview(buf) as mbuf:
|
mbuf[: len(next_piece)] = next_piece
|
written += len(next_piece)
|
if written == len(data):
|
break
|
return written
|
|
|
@attr.frozen
|
class UDPEndpoint:
|
ip: IPAddress
|
port: int
|
|
def as_python_sockaddr(self):
|
sockaddr = (self.ip.compressed, self.port)
|
if isinstance(self.ip, ipaddress.IPv6Address):
|
sockaddr += (0, 0)
|
return sockaddr
|
|
@classmethod
|
def from_python_sockaddr(cls, sockaddr):
|
ip, port = sockaddr[:2]
|
return cls(ip=ipaddress.ip_address(ip), port=port)
|
|
|
@attr.frozen
|
class UDPBinding:
|
local: UDPEndpoint
|
|
|
@attr.frozen
|
class UDPPacket:
|
source: UDPEndpoint
|
destination: UDPEndpoint
|
payload: bytes = attr.ib(repr=lambda p: p.hex())
|
|
def reply(self, payload):
|
return UDPPacket(
|
source=self.destination, destination=self.source, payload=payload
|
)
|
|
|
@attr.frozen
|
class FakeSocketFactory(trio.abc.SocketFactory):
|
fake_net: "FakeNet"
|
|
def socket(self, family: int, type: int, proto: int) -> "FakeSocket":
|
return FakeSocket._create(self.fake_net, family, type, proto)
|
|
|
@attr.frozen
|
class FakeHostnameResolver(trio.abc.HostnameResolver):
|
fake_net: "FakeNet"
|
|
async def getaddrinfo(
|
self, host: str, port: Union[int, str], family=0, type=0, proto=0, flags=0
|
):
|
raise NotImplementedError("FakeNet doesn't do fake DNS yet")
|
|
async def getnameinfo(self, sockaddr, flags: int):
|
raise NotImplementedError("FakeNet doesn't do fake DNS yet")
|
|
|
class FakeNet(metaclass=Final):
|
def __init__(self):
|
# When we need to pick an arbitrary unique ip address/port, use these:
|
self._auto_ipv4_iter = ipaddress.IPv4Network("1.0.0.0/8").hosts()
|
self._auto_ipv4_iter = ipaddress.IPv6Network("1::/16").hosts()
|
self._auto_port_iter = iter(range(50000, 65535))
|
|
self._bound: Dict[UDPBinding, FakeSocket] = {}
|
|
self.route_packet = None
|
|
def _bind(self, binding: UDPBinding, socket: "FakeSocket") -> None:
|
if binding in self._bound:
|
_fake_err(errno.EADDRINUSE)
|
self._bound[binding] = socket
|
|
def enable(self) -> None:
|
trio.socket.set_custom_socket_factory(FakeSocketFactory(self))
|
trio.socket.set_custom_hostname_resolver(FakeHostnameResolver(self))
|
|
def send_packet(self, packet) -> None:
|
if self.route_packet is None:
|
self.deliver_packet(packet)
|
else:
|
self.route_packet(packet)
|
|
def deliver_packet(self, packet) -> None:
|
binding = UDPBinding(local=packet.destination)
|
if binding in self._bound:
|
self._bound[binding]._deliver_packet(packet)
|
else:
|
# No valid destination, so drop it
|
pass
|
|
|
class FakeSocket(trio.socket.SocketType, metaclass=NoPublicConstructor):
|
def __init__(self, fake_net: FakeNet, family: int, type: int, proto: int):
|
self._fake_net = fake_net
|
|
if not family:
|
family = trio.socket.AF_INET
|
if not type:
|
type = trio.socket.SOCK_STREAM
|
|
if family not in (trio.socket.AF_INET, trio.socket.AF_INET6):
|
raise NotImplementedError(f"FakeNet doesn't (yet) support family={family}")
|
if type != trio.socket.SOCK_DGRAM:
|
raise NotImplementedError(f"FakeNet doesn't (yet) support type={type}")
|
|
self.family = family
|
self.type = type
|
self.proto = proto
|
|
self._closed = False
|
|
self._packet_sender, self._packet_receiver = trio.open_memory_channel(
|
float("inf")
|
)
|
|
# This is the source-of-truth for what port etc. this socket is bound to
|
self._binding: Optional[UDPBinding] = None
|
|
def _check_closed(self):
|
if self._closed:
|
_fake_err(errno.EBADF)
|
|
def close(self):
|
# breakpoint()
|
if self._closed:
|
return
|
self._closed = True
|
if self._binding is not None:
|
del self._fake_net._bound[self._binding]
|
self._packet_receiver.close()
|
|
async def _resolve_address_nocp(self, address, *, local):
|
return await trio._socket._resolve_address_nocp(
|
self.type,
|
self.family,
|
self.proto,
|
address=address,
|
ipv6_v6only=False,
|
local=local,
|
)
|
|
def _deliver_packet(self, packet: UDPPacket):
|
try:
|
self._packet_sender.send_nowait(packet)
|
except trio.BrokenResourceError:
|
# sending to a closed socket -- UDP packets get dropped
|
pass
|
|
################################################################
|
# Actual IO operation implementations
|
################################################################
|
|
async def bind(self, addr):
|
self._check_closed()
|
if self._binding is not None:
|
_fake_error(errno.EINVAL)
|
await trio.lowlevel.checkpoint()
|
ip_str, port = await self._resolve_address_nocp(addr, local=True)
|
ip = ipaddress.ip_address(ip_str)
|
assert _family_for(ip) == self.family
|
# We convert binds to INET_ANY into binds to localhost
|
if ip == ipaddress.ip_address("0.0.0.0"):
|
ip = ipaddress.ip_address("127.0.0.1")
|
elif ip == ipaddress.ip_address("::"):
|
ip = ipaddress.ip_address("::1")
|
if port == 0:
|
port = next(self._fake_net._auto_port_iter)
|
binding = UDPBinding(local=UDPEndpoint(ip, port))
|
self._fake_net._bind(binding, self)
|
self._binding = binding
|
|
async def connect(self, peer):
|
raise NotImplementedError("FakeNet does not (yet) support connected sockets")
|
|
async def sendmsg(self, *args):
|
self._check_closed()
|
ancdata = []
|
flags = 0
|
address = None
|
if len(args) == 1:
|
(buffers,) = args
|
elif len(args) == 2:
|
buffers, address = args
|
elif len(args) == 3:
|
buffers, flags, address = args
|
elif len(args) == 4:
|
buffers, ancdata, flags, address = args
|
else:
|
raise TypeError("wrong number of arguments")
|
|
await trio.lowlevel.checkpoint()
|
|
if address is not None:
|
address = await self._resolve_address_nocp(address, local=False)
|
if ancdata:
|
raise NotImplementedError("FakeNet doesn't support ancillary data")
|
if flags:
|
raise NotImplementedError(f"FakeNet send flags must be 0, not {flags}")
|
|
if address is None:
|
_fake_err(errno.ENOTCONN)
|
|
destination = UDPEndpoint.from_python_sockaddr(address)
|
|
if self._binding is None:
|
await self.bind((_wildcard_ip_for(self.family).compressed, 0))
|
|
payload = b"".join(buffers)
|
|
packet = UDPPacket(
|
source=self._binding.local,
|
destination=destination,
|
payload=payload,
|
)
|
|
self._fake_net.send_packet(packet)
|
|
return len(payload)
|
|
async def recvmsg_into(self, buffers, ancbufsize=0, flags=0):
|
if ancbufsize != 0:
|
raise NotImplementedError("FakeNet doesn't support ancillary data")
|
if flags != 0:
|
raise NotImplementedError("FakeNet doesn't support any recv flags")
|
|
self._check_closed()
|
|
ancdata = []
|
msg_flags = 0
|
|
packet = await self._packet_receiver.receive()
|
address = packet.source.as_python_sockaddr()
|
written = _scatter(packet.payload, buffers)
|
if written < len(packet.payload):
|
msg_flags |= trio.socket.MSG_TRUNC
|
return written, ancdata, msg_flags, address
|
|
################################################################
|
# Simple state query stuff
|
################################################################
|
|
def getsockname(self):
|
self._check_closed()
|
if self._binding is not None:
|
return self._binding.local.as_python_sockaddr()
|
elif self.family == trio.socket.AF_INET:
|
return ("0.0.0.0", 0)
|
else:
|
assert self.family == trio.socket.AF_INET6
|
return ("::", 0)
|
|
def getpeername(self):
|
self._check_closed()
|
if self._binding is not None:
|
if self._binding.remote is not None:
|
return self._binding.remote.as_python_sockaddr()
|
_fake_err(errno.ENOTCONN)
|
|
def getsockopt(self, level, item):
|
self._check_closed()
|
raise OSError(f"FakeNet doesn't implement getsockopt({level}, {item})")
|
|
def setsockopt(self, level, item, value):
|
self._check_closed()
|
|
if (level, item) == (trio.socket.IPPROTO_IPV6, trio.socket.IPV6_V6ONLY):
|
if not value:
|
raise NotImplementedError("FakeNet always has IPV6_V6ONLY=True")
|
|
raise OSError(f"FakeNet doesn't implement setsockopt({level}, {item}, ...)")
|
|
################################################################
|
# Various boilerplate and trivial stubs
|
################################################################
|
|
def __enter__(self):
|
return self
|
|
def __exit__(self, *exc_info):
|
self.close()
|
|
async def send(self, data, flags=0):
|
return await self.sendto(data, flags, None)
|
|
async def sendto(self, *args):
|
if len(args) == 2:
|
data, address = args
|
flags = 0
|
elif len(args) == 3:
|
data, flags, address = args
|
else:
|
raise TypeError("wrong number of arguments")
|
return await self.sendmsg([data], [], flags, address)
|
|
async def recv(self, bufsize, flags=0):
|
data, address = await self.recvfrom(bufsize, flags)
|
return data
|
|
async def recv_into(self, buf, nbytes=0, flags=0):
|
got_bytes, address = await self.recvfrom_into(buf, nbytes, flags)
|
return got_bytes
|
|
async def recvfrom(self, bufsize, flags=0):
|
data, ancdata, msg_flags, address = await self.recvmsg(bufsize, flags)
|
return data, address
|
|
async def recvfrom_into(self, buf, nbytes=0, flags=0):
|
if nbytes != 0 and nbytes != len(buf):
|
raise NotImplementedError("partial recvfrom_into")
|
got_nbytes, ancdata, msg_flags, address = await self.recvmsg_into(
|
[buf], 0, flags
|
)
|
return got_nbytes, address
|
|
async def recvmsg(self, bufsize, ancbufsize=0, flags=0):
|
buf = bytearray(bufsize)
|
got_nbytes, ancdata, msg_flags, address = await self.recvmsg_into(
|
[buf], ancbufsize, flags
|
)
|
return (bytes(buf[:got_nbytes]), ancdata, msg_flags, address)
|
|
def fileno(self):
|
raise NotImplementedError("can't get fileno() for FakeNet sockets")
|
|
def detach(self):
|
raise NotImplementedError("can't detach() a FakeNet socket")
|
|
def get_inheritable(self):
|
return False
|
|
def set_inheritable(self, inheritable):
|
if inheritable:
|
raise NotImplementedError("FakeNet can't make inheritable sockets")
|
|
def share(self, process_id):
|
raise NotImplementedError("FakeNet can't share sockets")
|
