Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
569 lines (505 sloc) 22.7 KB
"""
Low level socket utilities.
Pyro - Python Remote Objects. Copyright by Irmen de Jong (irmen@razorvine.net).
"""
import os
import platform
import socket
import errno
import time
import select
import ipaddress
import weakref
import contextlib
from typing import Union, Optional, Tuple, Dict, Type, Any
try:
import ssl
except ImportError:
pass
from . import config
from .errors import CommunicationError, TimeoutError, ConnectionClosedError
# Note: other interesting errnos are EPERM, ENOBUFS, EMFILE
# but it seems to me that all these signify an unrecoverable situation.
# So I didn't include them in the list of retryable errors.
ERRNO_RETRIES = [errno.EINTR, errno.EAGAIN, errno.EWOULDBLOCK, errno.EINPROGRESS]
if hasattr(errno, "WSAEINTR"):
ERRNO_RETRIES.append(errno.WSAEINTR)
if hasattr(errno, "WSAEWOULDBLOCK"):
ERRNO_RETRIES.append(errno.WSAEWOULDBLOCK)
if hasattr(errno, "WSAEINPROGRESS"):
ERRNO_RETRIES.append(errno.WSAEINPROGRESS)
ERRNO_BADF = [errno.EBADF]
if hasattr(errno, "WSAEBADF"):
ERRNO_BADF.append(errno.WSAEBADF)
ERRNO_ENOTSOCK = [errno.ENOTSOCK]
if hasattr(errno, "WSAENOTSOCK"):
ERRNO_ENOTSOCK.append(errno.WSAENOTSOCK)
if not hasattr(socket, "SOL_TCP"):
socket.SOL_TCP = socket.IPPROTO_TCP
ERRNO_EADDRNOTAVAIL = [errno.EADDRNOTAVAIL]
if hasattr(errno, "WSAEADDRNOTAVAIL"):
ERRNO_EADDRNOTAVAIL.append(errno.WSAEADDRNOTAVAIL)
ERRNO_EADDRINUSE = [errno.EADDRINUSE]
if hasattr(errno, "WSAEADDRINUSE"):
ERRNO_EADDRINUSE.append(errno.WSAEADDRINUSE)
# msg_waitall has proven to be unreliable on windows
USE_MSG_WAITALL = hasattr(socket, "MSG_WAITALL") and platform.system() != "Windows"
def get_ip_address(hostname: str, workaround127: bool = False, version: int = None) \
-> Union[ipaddress.IPv4Address, ipaddress.IPv6Address]:
"""
Returns the IP address for the given host. If you enable the workaround,
it will use a little hack if the ip address is found to be the loopback address.
The hack tries to discover an externally visible ip address instead (this only works for ipv4 addresses).
Set ipVersion=6 to return ipv6 addresses, 4 to return ipv4, 0 to let OS choose the best one or None to use config.PREFER_IP_VERSION.
"""
if not workaround127:
with contextlib.suppress(ValueError):
addr = ipaddress.ip_address(hostname)
return addr
def getaddr(ip_version):
if ip_version == 6:
family = socket.AF_INET6
elif ip_version == 4:
family = socket.AF_INET
elif ip_version == 0:
family = socket.AF_UNSPEC
else:
raise ValueError("unknown value for argument ipVersion.")
ip = socket.getaddrinfo(hostname or socket.gethostname(), 80, family, socket.SOCK_STREAM, socket.SOL_TCP)[0][4][0]
if workaround127 and (ip.startswith("127.") or ip == "0.0.0.0"):
return get_interface("4.2.2.2").ip
return ipaddress.ip_address(ip)
try:
if hostname and ':' in hostname and version is None:
version = 0
return getaddr(config.PREFER_IP_VERSION) if version is None else getaddr(version)
except socket.gaierror:
if version == 6 or (version is None and config.PREFER_IP_VERSION == 6):
raise socket.error("unable to determine IPV6 address")
return getaddr(0)
def get_interface(ip_address: Union[str, ipaddress.IPv4Address, ipaddress.IPv6Address]) \
-> Union[ipaddress.IPv4Interface, ipaddress.IPv6Interface]:
"""tries to find the network interface that connects to the given host's address"""
if isinstance(ip_address, str):
ip_address = get_ip_address(ip_address)
family = socket.AF_INET if ip_address.version == 4 else socket.AF_INET6
with socket.socket(family, socket.SOCK_DGRAM) as sock:
sock.connect((str(ip_address), 53)) # 53=dns
return ipaddress.ip_interface(sock.getsockname()[0])
def __retrydelays():
# first try a few very short delays,
# if that doesn't work, increase by 0.1 sec every time
yield 0.0001
yield 0.001
yield 0.01
d = 0.1
while True:
yield d
d += 0.1
def receive_data(sock: socket.socket, size: int) -> bytes:
"""Retrieve a given number of bytes from a socket.
It is expected the socket is able to supply that number of bytes.
If it isn't, an exception is raised (you will not get a zero length result
or a result that is smaller than what you asked for). The partial data that
has been received however is stored in the 'partialData' attribute of
the exception object."""
try:
delays = __retrydelays()
msglen = 0
data = bytearray()
if USE_MSG_WAITALL and not hasattr(sock, "getpeercert"): # ssl doesn't support recv flags
while True:
try:
chunk = sock.recv(size, socket.MSG_WAITALL)
if len(chunk) == size:
return chunk
# less data than asked, drop down into normal receive loop to finish
msglen = len(chunk)
data.extend(chunk)
break
except socket.timeout:
raise TimeoutError("receiving: timeout")
except socket.error as x:
err = getattr(x, "errno", x.args[0])
if err not in ERRNO_RETRIES:
raise ConnectionClosedError("receiving: connection lost: " + str(x))
time.sleep(next(delays)) # a slight delay to wait before retrying
# old fashioned recv loop, we gather chunks until the message is complete
while True:
try:
while msglen < size:
# 60k buffer limit avoids problems on certain OSes like VMS, Windows
chunk = sock.recv(min(60000, size - msglen))
if not chunk:
break
data.extend(chunk)
msglen += len(chunk)
if len(data) != size:
err = ConnectionClosedError("receiving: not enough data")
err.partialData = data # store the message that was received until now
raise err
return data # yay, complete
except socket.timeout:
raise TimeoutError("receiving: timeout")
except socket.error as x:
err = getattr(x, "errno", x.args[0])
if err not in ERRNO_RETRIES:
raise ConnectionClosedError("receiving: connection lost: " + str(x))
time.sleep(next(delays)) # a slight delay to wait before retrying
except socket.timeout:
raise TimeoutError("receiving: timeout")
def send_data(sock: socket.socket, data: bytes) -> None:
"""
Send some data over a socket.
Some systems have problems with ``sendall()`` when the socket is in non-blocking mode.
For instance, Mac OS X seems to be happy to throw EAGAIN errors too often.
This function falls back to using a regular send loop if needed.
"""
if sock.gettimeout() is None:
# socket is in blocking mode, we can use sendall normally.
try:
sock.sendall(data)
return
except socket.timeout:
raise TimeoutError("sending: timeout")
except socket.error as x:
raise ConnectionClosedError("sending: connection lost: " + str(x))
else:
# Socket is in non-blocking mode, use regular send loop.
delays = __retrydelays()
while data:
try:
sent = sock.send(data)
data = data[sent:]
except socket.timeout:
raise TimeoutError("sending: timeout")
except socket.error as x:
err = getattr(x, "errno", x.args[0])
if err not in ERRNO_RETRIES:
raise ConnectionClosedError("sending: connection lost: " + str(x))
time.sleep(next(delays)) # a slight delay to wait before retrying
def create_socket(bind: Union[Tuple, str] = None,
connect: Union[Tuple, str] = None,
reuseaddr: bool = False, keepalive: bool = True,
timeout: Optional[float] = -1, noinherit: bool = False,
ipv6: bool = False, nodelay: bool = True, sslContext: ssl.SSLContext = None) -> socket.socket:
"""
Create a socket. Default socket options are keepalive and IPv4 family, and nodelay (nagle disabled).
If 'bind' or 'connect' is a string, it is assumed a Unix domain socket is requested.
Otherwise, a normal tcp/ip socket tuple (addr, port, ...) is used.
Set ipv6=True to create an IPv6 socket rather than IPv4.
Set ipv6=None to use the PREFER_IP_VERSION config setting.
"""
if bind and connect:
raise ValueError("bind and connect cannot both be specified at the same time")
forceIPv6 = ipv6 or (ipv6 is None and config.PREFER_IP_VERSION == 6)
if isinstance(bind, str) or isinstance(connect, str):
family = socket.AF_UNIX
elif not bind and not connect:
family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
elif isinstance(bind, tuple):
if not bind[0]:
family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
else:
addr = get_ip_address(bind[0])
if addr.version == 4:
if forceIPv6:
raise ValueError("IPv4 address is used bind argument with forceIPv6 argument:" + bind[0] + ".")
family = socket.AF_INET
elif addr.version == 6:
family = socket.AF_INET6
# replace bind addresses by their ipv6 counterparts (4-tuple)
bind = (bind[0], bind[1], 0, 0)
else:
raise ValueError("unknown bind format.")
elif isinstance(connect, tuple):
if not connect[0]:
family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
else:
addr = get_ip_address(connect[0])
if addr.version == 4:
if forceIPv6:
raise ValueError("IPv4 address is used in connect argument with forceIPv6 argument:" + connect[0] + ".")
family = socket.AF_INET
elif addr.version == 6:
family = socket.AF_INET6
# replace connect addresses by their ipv6 counterparts (4-tuple)
connect = (connect[0], connect[1], 0, 0)
else:
raise ValueError("unknown connect format.")
else:
raise ValueError("unknown bind or connect format.")
sock = socket.socket(family, socket.SOCK_STREAM)
if sslContext:
if bind:
sock = sslContext.wrap_socket(sock, server_side=True)
elif connect:
sock = sslContext.wrap_socket(sock, server_side=False, server_hostname=connect[0])
else:
sock = sslContext.wrap_socket(sock, server_side=False)
if nodelay:
set_nodelay(sock)
if reuseaddr:
set_reuseaddr(sock)
if noinherit:
set_noinherit(sock)
if timeout is not None:
if timeout == 0:
timeout = None
elif timeout >= 0:
sock.settimeout(timeout)
if bind:
if type(bind) is tuple and bind[1] == 0:
bind_unused_port(sock, bind[0])
else:
sock.bind(bind)
with contextlib.suppress(OSError, IOError):
sock.listen(100)
if connect:
try:
sock.connect(connect)
except socket.error as xv:
# This can happen when the socket is in non-blocking mode (or has a timeout configured).
# We check if it is a retryable errno (usually EINPROGRESS).
# If so, we use select() to wait until the socket is in writable state,
# essentially rebuilding a blocking connect() call.
errno = getattr(xv, "errno", 0)
if errno in ERRNO_RETRIES:
if timeout:
if timeout < 0:
timeout = None
elif timeout < 0.1:
timeout = 0.1
while True:
try:
sr, sw, se = select.select([], [sock], [sock], timeout)
except InterruptedError:
continue
if sock in sw:
break # yay, writable now, connect() completed
elif sock in se:
sock.close() # close the socket that refused to connect
raise socket.error("connect failed")
else:
sock.close() # close the socket that refused to connect
raise
if keepalive:
set_keepalive(sock)
return sock
def create_bc_socket(bind: Union[Tuple, str] = None, reuseaddr: bool = False,
timeout: Optional[float] = -1, ipv6: bool = False) -> socket.socket:
"""
Create a udp broadcast socket.
Set ipv6=True to create an IPv6 socket rather than IPv4.
Set ipv6=None to use the PREFER_IP_VERSION config setting.
"""
forceIPv6 = ipv6 or (ipv6 is None and config.PREFER_IP_VERSION == 6)
if not bind:
family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
elif isinstance(bind, tuple):
if not bind[0]:
family = socket.AF_INET6 if forceIPv6 else socket.AF_INET
else:
addr = get_ip_address(bind[0])
if addr.version == 4:
if forceIPv6:
raise ValueError("IPv4 address is used with forceIPv6 option:" + bind[0] + ".")
family = socket.AF_INET
elif addr.version == 6:
family = socket.AF_INET6
bind = (bind[0], bind[1], 0, 0)
else:
raise ValueError("unknown bind format: %r" % (bind,))
else:
raise ValueError("unknown bind format: %r" % (bind,))
sock = socket.socket(family, socket.SOCK_DGRAM)
if family == socket.AF_INET:
sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
if reuseaddr:
set_reuseaddr(sock)
if timeout is None:
sock.settimeout(None)
else:
if timeout >= 0:
sock.settimeout(timeout)
if bind:
host = bind[0] or ""
port = bind[1]
if port == 0:
bind_unused_port(sock, host)
else:
if len(bind) == 2:
sock.bind((host, port)) # ipv4
elif len(bind) == 4:
sock.bind((host, port, 0, 0)) # ipv6
else:
raise ValueError("bind must be None, 2-tuple or 4-tuple")
return sock
def set_reuseaddr(sock: socket.socket) -> None:
"""sets the SO_REUSEADDR option on the socket, if possible."""
with contextlib.suppress(Exception):
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
def set_nodelay(sock: socket.socket) -> None:
"""sets the TCP_NODELAY option on the socket (to disable Nagle's algorithm), if possible."""
with contextlib.suppress(Exception):
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
def set_keepalive(sock: socket.socket) -> None:
"""sets the SO_KEEPALIVE option on the socket, if possible."""
with contextlib.suppress(Exception):
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
try:
import fcntl
def set_noinherit(sock: socket.socket) -> None:
"""Mark the given socket fd as non-inheritable to child processes"""
fd = sock.fileno()
flags = fcntl.fcntl(fd, fcntl.F_GETFD)
fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
except ImportError:
# no fcntl available, try the windows version
try:
from ctypes import windll, WinError, wintypes
# help ctypes to set the proper args for this kernel32 call on 64-bit pythons
_SetHandleInformation = windll.kernel32.SetHandleInformation
_SetHandleInformation.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD]
_SetHandleInformation.restype = wintypes.BOOL # don't need this, but might as well
def set_noinherit(sock: socket.socket) -> None:
"""Mark the given socket fd as non-inheritable to child processes"""
if not _SetHandleInformation(sock.fileno(), 1, 0):
raise WinError()
except (ImportError, NotImplementedError):
# nothing available, define a dummy function
def set_noinherit(sock: socket.socket) -> None:
"""Mark the given socket fd as non-inheritable to child processes (dummy)"""
pass
class SocketConnection(object):
"""A wrapper class for plain sockets, containing various methods such as :meth:`send` and :meth:`recv`"""
def __init__(self, sock: socket.socket, objectId: str = None, keep_open: bool = False) -> None:
self.sock = sock
self.objectId = objectId
self.pyroInstances = {} # type: Dict[Type, Any] # pyro objects for instance_mode=session
self.tracked_resources = weakref.WeakSet() # type: weakref.WeakSet[Any] # weakrefs to resources for this connection
self.keep_open = keep_open
def __del__(self):
self.close()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
def send(self, data: bytes) -> None:
send_data(self.sock, data)
def recv(self, size: int) -> bytes:
return receive_data(self.sock, size)
def close(self) -> None:
if self.keep_open:
return
with contextlib.suppress(Exception):
self.sock.shutdown(socket.SHUT_RDWR)
with contextlib.suppress(Exception):
self.sock.close()
self.pyroInstances = {} # release the session instances
for rsc in self.tracked_resources:
with contextlib.suppress(Exception):
rsc.close() # it is assumed a 'resource' has a close method.
self.tracked_resources.clear()
def fileno(self) -> int:
return self.sock.fileno()
def family(self) -> str:
return family_str(self.sock)
def settimeout(self, timeout: Optional[float]) -> None:
self.sock.settimeout(timeout)
def gettimeout(self) -> Optional[float]:
return self.sock.gettimeout()
def getpeercert(self) -> Optional[dict]:
try:
return self.sock.getpeercert() # type: ignore
except AttributeError:
return None
timeout = property(gettimeout, settimeout)
def family_str(sock) -> str:
f = sock.family
if f == socket.AF_INET:
return "IPv4"
if f == socket.AF_INET6:
return "IPv6"
if hasattr(socket, "AF_UNIX") and f == socket.AF_UNIX:
return "Unix"
return "???"
def find_probably_unused_port(family: int = socket.AF_INET, socktype: int = socket.SOCK_STREAM) -> int:
"""Returns an unused port that should be suitable for binding (likely, but not guaranteed).
This code is copied from the stdlib's test.test_support module."""
with socket.socket(family, socktype) as sock:
return bind_unused_port(sock)
def bind_unused_port(sock: socket.socket, host: Union[str, ipaddress.IPv4Address, ipaddress.IPv6Address] = 'localhost') -> int:
"""Bind the socket to a free port and return the port number.
This code is based on the code in the stdlib's test.test_support module."""
if sock.family in (socket.AF_INET, socket.AF_INET6) and sock.type == socket.SOCK_STREAM:
if hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
with contextlib.suppress(socket.error):
sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
if not isinstance(host, str):
host = str(host)
if sock.family == socket.AF_INET:
if host == 'localhost':
sock.bind(('127.0.0.1', 0))
else:
sock.bind((host, 0))
elif sock.family == socket.AF_INET6:
if host == 'localhost':
sock.bind(('::1', 0, 0, 0))
else:
sock.bind((host, 0, 0, 0))
else:
raise CommunicationError("unsupported socket family: " + str(sock.family))
return sock.getsockname()[1]
def interrupt_socket(address: Tuple[str, int]) -> None:
"""bit of a hack to trigger a blocking server to get out of the loop, useful at clean shutdowns"""
with contextlib.suppress(socket.error):
sock = create_socket(connect=address, keepalive=False, timeout=-1)
with contextlib.suppress(socket.error, AttributeError):
sock.sendall(b"!" * 16)
with contextlib.suppress(OSError, socket.error):
sock.shutdown(socket.SHUT_RDWR)
sock.close()
__ssl_server_context = None
__ssl_client_context = None
def get_ssl_context(servercert: str = "", serverkey: str = "", clientcert: str = "", clientkey: str = "",
cacerts: str = "", keypassword: str = "") -> ssl.SSLContext:
"""creates an SSL context and caches it, so you have to set the parameters correctly before doing anything"""
global __ssl_client_context, __ssl_server_context
if servercert:
if clientcert:
raise ValueError("can't have both server cert and client cert")
# server context
if __ssl_server_context:
return __ssl_server_context
if not os.path.isfile(servercert):
raise IOError("server cert file not found")
if serverkey and not os.path.isfile(serverkey):
raise IOError("server key file not found")
__ssl_server_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
__ssl_server_context.load_cert_chain(servercert, serverkey or None, keypassword or None) # type: ignore
if cacerts:
if os.path.isdir(cacerts):
__ssl_server_context.load_verify_locations(capath=cacerts)
else:
__ssl_server_context.load_verify_locations(cafile=cacerts)
if config.SSL_REQUIRECLIENTCERT:
__ssl_server_context.verify_mode = ssl.CERT_REQUIRED # 2-way ssl, server+client certs
else:
__ssl_server_context.verify_mode = ssl.CERT_NONE # 1-way ssl, server cert only
return __ssl_server_context
else:
# client context
if __ssl_client_context:
return __ssl_client_context
__ssl_client_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
if clientcert:
if not os.path.isfile(clientcert):
raise IOError("client cert file not found")
__ssl_client_context.load_cert_chain(clientcert, clientkey or None, keypassword or None) # type: ignore
if cacerts:
if os.path.isdir(cacerts):
__ssl_client_context.load_verify_locations(capath=cacerts)
else:
__ssl_client_context.load_verify_locations(cafile=cacerts)
return __ssl_client_context
You can’t perform that action at this time.