Innovenergy_trunk/NodeRed/NodeRedFiles/pika-0.13.1/tests/acceptance/forward_server.py

527 lines
20 KiB
Python
Raw Normal View History

2024-06-05 10:33:44 +00:00
"""TCP/IP forwarding/echo service for testing."""
from __future__ import print_function
import array
from datetime import datetime
import errno
from functools import partial
import logging
import multiprocessing
import os
import socket
import struct
import sys
import threading
import traceback
import pika.compat
if pika.compat.PY3:
def buffer(object, offset, size): # pylint: disable=W0622
"""array etc. have the buffer protocol"""
return object[offset:offset + size]
try:
import SocketServer
except ImportError:
import socketserver as SocketServer # pylint: disable=F0401
def _trace(fmt, *args):
"""Format and output the text to stderr"""
print((fmt % args) + "\n", end="", file=sys.stderr)
class ForwardServer(object): # pylint: disable=R0902
""" Implement a TCP/IP forwarding/echo service for testing. Listens for
an incoming TCP/IP connection, accepts it, then connects to the given
remote address and forwards data back and forth between the two
endpoints.
This is similar to the subset of `netcat` functionality, but without
dependency on any specific flavor of netcat
Connection forwarding example; forward local connection to default
rabbitmq addr, connect to rabbit via forwarder, then disconnect
forwarder, then attempt another pika operation to see what happens
with ForwardServer(("localhost", 5672)) as fwd:
params = pika.ConnectionParameters(
host=fwd.server_address[0],
port=fwd.server_address[1])
conn = pika.BlockingConnection(params)
# Once outside the context, the forwarder is disconnected
# Let's see what happens in pika with a disconnected server
channel = conn.channel()
Echo server example
def produce(sock):
sock.sendall("12345")
sock.shutdown(socket.SHUT_WR)
with ForwardServer(None) as fwd:
sock = socket.socket()
sock.connect(fwd.server_address)
worker = threading.Thread(target=produce,
args=[sock])
worker.start()
data = sock.makefile().read()
assert data == "12345", data
worker.join()
"""
# Amount of time, in seconds, we're willing to wait for the subprocess
_SUBPROC_TIMEOUT = 10
def __init__(
self, # pylint: disable=R0913
remote_addr,
remote_addr_family=socket.AF_INET,
remote_socket_type=socket.SOCK_STREAM,
server_addr=("127.0.0.1", 0),
server_addr_family=socket.AF_INET,
server_socket_type=socket.SOCK_STREAM,
local_linger_args=None):
"""
:param tuple remote_addr: remote server's IP address, whose structure
depends on remote_addr_family; pair (host-or-ip-addr, port-number).
Pass None to have ForwardServer behave as echo server.
:param remote_addr_family: socket.AF_INET (the default), socket.AF_INET6
or socket.AF_UNIX.
:param remote_socket_type: only socket.SOCK_STREAM is supported at this
time
:param server_addr: optional address for binding this server's listening
socket; the format depends on server_addr_family; defaults to
("127.0.0.1", 0)
:param server_addr_family: Address family for this server's listening
socket; socket.AF_INET (the default), socket.AF_INET6 or
socket.AF_UNIX; defaults to socket.AF_INET
:param server_socket_type: only socket.SOCK_STREAM is supported at this
time
:param tuple local_linger_args: SO_LINGER sockoverride for the local
connection sockets, to be configured after connection is accepted.
None for default, which is to not change the SO_LINGER option.
Otherwise, its a two-tuple, where the first element is the `l_onoff`
switch, and the second element is the `l_linger` value, in seconds
"""
self._logger = logging.getLogger(__name__)
self._remote_addr = remote_addr
self._remote_addr_family = remote_addr_family
assert remote_socket_type == socket.SOCK_STREAM, remote_socket_type
self._remote_socket_type = remote_socket_type
assert server_addr is not None
self._server_addr = server_addr
assert server_addr_family is not None
self._server_addr_family = server_addr_family
assert server_socket_type == socket.SOCK_STREAM, server_socket_type
self._server_socket_type = server_socket_type
self._local_linger_args = local_linger_args
self._subproc = None
@property
def running(self):
"""Property: True if ForwardServer is active"""
return self._subproc is not None
@property
def server_address_family(self):
"""Property: Get listening socket's address family
NOTE: undefined before server starts and after it shuts down
"""
assert self._server_addr_family is not None, "Not in context"
return self._server_addr_family
@property
def server_address(self):
""" Property: Get listening socket's address; the returned value
depends on the listening socket's address family
NOTE: undefined before server starts and after it shuts down
"""
assert self._server_addr is not None, "Not in context"
return self._server_addr
def __enter__(self):
""" Context manager entry. Starts the forwarding server
:returns: self
"""
return self.start()
def __exit__(self, *args):
""" Context manager exit; stops the forwarding server
"""
self.stop()
def start(self):
""" Start the server
NOTE: The context manager is the recommended way to use
ForwardServer. start()/stop() are alternatives to the context manager
use case and are mutually exclusive with it.
:returns: self
"""
queue = multiprocessing.Queue()
self._subproc = multiprocessing.Process(
target=_run_server,
kwargs=dict(
local_addr=self._server_addr,
local_addr_family=self._server_addr_family,
local_socket_type=self._server_socket_type,
local_linger_args=self._local_linger_args,
remote_addr=self._remote_addr,
remote_addr_family=self._remote_addr_family,
remote_socket_type=self._remote_socket_type,
queue=queue))
self._subproc.daemon = True
self._subproc.start()
try:
# Get server socket info from subprocess
self._server_addr_family, self._server_addr = queue.get(
block=True, timeout=self._SUBPROC_TIMEOUT)
queue.close()
except Exception: # pylint: disable=W0703
try:
self._logger.exception(
"Failed while waiting for local socket info")
# Preserve primary exception and traceback
raise
finally:
# Clean up
try:
self.stop()
except Exception: # pylint: disable=W0703
# Suppress secondary exception in favor of the primary
self._logger.exception(
"Emergency subprocess shutdown failed")
return self
def stop(self):
"""Stop the server
NOTE: The context manager is the recommended way to use
ForwardServer. start()/stop() are alternatives to the context manager
use case and are mutually exclusive with it.
"""
self._logger.info("ForwardServer STOPPING")
try:
self._subproc.terminate()
self._subproc.join(timeout=self._SUBPROC_TIMEOUT)
if self._subproc.is_alive():
self._logger.error(
"ForwardServer failed to terminate, killing it")
os.kill(self._subproc.pid)
self._subproc.join(timeout=self._SUBPROC_TIMEOUT)
assert not self._subproc.is_alive(), self._subproc
# Log subprocess's exit code; NOTE: negative signal.SIGTERM (usually
# -15) is normal on POSIX systems - it corresponds to SIGTERM
exit_code = self._subproc.exitcode
self._logger.info("ForwardServer terminated with exitcode=%s",
exit_code)
finally:
self._subproc = None
def _run_server(
local_addr,
local_addr_family,
local_socket_type, # pylint: disable=R0913
local_linger_args,
remote_addr,
remote_addr_family,
remote_socket_type,
queue):
""" Run the server; executed in the subprocess
:param local_addr: listening address
:param local_addr_family: listening address family; one of socket.AF_*
:param local_socket_type: listening socket type; typically
socket.SOCK_STREAM
:param tuple local_linger_args: SO_LINGER sockoverride for the local
connection sockets, to be configured after connection is accepted.
Pass None to not change SO_LINGER. Otherwise, its a two-tuple, where the
first element is the `l_onoff` switch, and the second element is the
`l_linger` value in seconds
:param remote_addr: address of the target server. Pass None to have
ForwardServer behave as echo server
:param remote_addr_family: address family for connecting to target server;
one of socket.AF_*
:param remote_socket_type: socket type for connecting to target server;
typically socket.SOCK_STREAM
:param multiprocessing.Queue queue: queue for depositing the forwarding
server's actual listening socket address family and bound address. The
parent process waits for this.
"""
# NOTE: We define _ThreadedTCPServer class as a closure in order to
# override some of its class members dynamically
# NOTE: we add `object` to the base classes because `_ThreadedTCPServer`
# isn't derived from `object`, which prevents `super` from working properly
class _ThreadedTCPServer(SocketServer.ThreadingMixIn,
SocketServer.TCPServer, object):
"""Threaded streaming server for forwarding"""
# Override TCPServer's class members
address_family = local_addr_family
socket_type = local_socket_type
allow_reuse_address = True
def __init__(self):
handler_class_factory = partial(
_TCPHandler,
local_linger_args=local_linger_args,
remote_addr=remote_addr,
remote_addr_family=remote_addr_family,
remote_socket_type=remote_socket_type)
super(_ThreadedTCPServer, self).__init__(
local_addr, handler_class_factory, bind_and_activate=True)
server = _ThreadedTCPServer()
# Send server socket info back to parent process
queue.put([server.socket.family, server.server_address])
queue.close()
server.serve_forever()
# NOTE: we add `object` to the base classes because `StreamRequestHandler` isn't
# derived from `object`, which prevents `super` from working properly
class _TCPHandler(SocketServer.StreamRequestHandler, object):
"""TCP/IP session handler instantiated by TCPServer upon incoming
connection. Implements forwarding/echo of the incoming connection.
"""
_SOCK_RX_BUF_SIZE = 16 * 1024
def __init__(
self, # pylint: disable=R0913
request,
client_address,
server,
local_linger_args,
remote_addr,
remote_addr_family,
remote_socket_type):
"""
:param request: for super
:param client_address: for super
"paarm server: for super
:param tuple local_linger_args: SO_LINGER sockoverride for the local
connection sockets, to be configured after connection is accepted.
Pass None to not change SO_LINGER. Otherwise, its a two-tuple, where
the first element is the `l_onoff` switch, and the second element is
the `l_linger` value in seconds
:param remote_addr: address of the target server. Pass None to have
ForwardServer behave as echo server.
:param remote_addr_family: address family for connecting to target
server; one of socket.AF_*
:param remote_socket_type: socket type for connecting to target server;
typically socket.SOCK_STREAM
:param **kwargs: kwargs for super class
"""
self._local_linger_args = local_linger_args
self._remote_addr = remote_addr
self._remote_addr_family = remote_addr_family
self._remote_socket_type = remote_socket_type
super(_TCPHandler, self).__init__(
request=request, client_address=client_address, server=server)
def handle(self): # pylint: disable=R0912
"""Connect to remote and forward data between local and remote"""
local_sock = self.connection
if self._local_linger_args is not None:
# Set SO_LINGER socket options on local socket
l_onoff, l_linger = self._local_linger_args
local_sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER,
struct.pack('ii', l_onoff, l_linger))
if self._remote_addr is not None:
# Forwarding set-up
remote_dest_sock = remote_src_sock = socket.socket(
family=self._remote_addr_family,
type=self._remote_socket_type,
proto=socket.IPPROTO_IP)
remote_dest_sock.connect(self._remote_addr)
_trace("%s _TCPHandler connected to remote %s", datetime.utcnow(),
remote_dest_sock.getpeername())
else:
# Echo set-up
remote_dest_sock, remote_src_sock = \
pika.compat._nonblocking_socketpair()
try:
local_forwarder = threading.Thread(
target=self._forward, args=(
local_sock,
remote_dest_sock,
))
local_forwarder.setDaemon(True)
local_forwarder.start()
try:
self._forward(remote_src_sock, local_sock)
finally:
# Wait for local forwarder thread to exit
local_forwarder.join()
finally:
try:
try:
_safe_shutdown_socket(remote_dest_sock, socket.SHUT_RDWR)
finally:
if remote_src_sock is not remote_dest_sock:
_safe_shutdown_socket(remote_src_sock, socket.SHUT_RDWR)
finally:
remote_dest_sock.close()
if remote_src_sock is not remote_dest_sock:
remote_src_sock.close()
def _forward(self, src_sock, dest_sock): # pylint: disable=R0912
"""Forward from src_sock to dest_sock"""
src_peername = src_sock.getpeername()
_trace("%s forwarding from %s to %s", datetime.utcnow(), src_peername,
dest_sock.getpeername())
try:
# NOTE: python 2.6 doesn't support bytearray with recv_into, so
# we use array.array instead; this is only okay as long as the
# array instance isn't shared across threads. See
# http://bugs.python.org/issue7827 and
# groups.google.com/forum/#!topic/comp.lang.python/M6Pqr-KUjQw
rx_buf = array.array("B", [0] * self._SOCK_RX_BUF_SIZE)
while True:
try:
nbytes = src_sock.recv_into(rx_buf)
except pika.compat.SOCKET_ERROR as exc:
if exc.errno == errno.EINTR:
continue
elif exc.errno == errno.ECONNRESET:
# Source peer forcibly closed connection
_trace("%s errno.ECONNRESET from %s", datetime.utcnow(),
src_peername)
break
else:
_trace("%s Unexpected errno=%s from %s\n%s",
datetime.utcnow(), exc.errno, src_peername,
"".join(traceback.format_stack()))
raise
if not nbytes:
# Source input EOF
_trace("%s EOF on %s", datetime.utcnow(), src_peername)
break
try:
dest_sock.sendall(buffer(rx_buf, 0, nbytes))
except pika.compat.SOCKET_ERROR as exc:
if exc.errno == errno.EPIPE:
# Destination peer closed its end of the connection
_trace("%s Destination peer %s closed its end of "
"the connection: errno.EPIPE", datetime.utcnow(),
dest_sock.getpeername())
break
elif exc.errno == errno.ECONNRESET:
# Destination peer forcibly closed connection
_trace("%s Destination peer %s forcibly closed "
"connection: errno.ECONNRESET",
datetime.utcnow(), dest_sock.getpeername())
break
else:
_trace("%s Unexpected errno=%s in sendall to %s\n%s",
datetime.utcnow(), exc.errno,
dest_sock.getpeername(), "".join(
traceback.format_stack()))
raise
except:
_trace("forward failed\n%s", "".join(traceback.format_exc()))
raise
finally:
_trace("%s done forwarding from %s", datetime.utcnow(),
src_peername)
try:
# Let source peer know we're done receiving
_safe_shutdown_socket(src_sock, socket.SHUT_RD)
finally:
# Let destination peer know we're done sending
_safe_shutdown_socket(dest_sock, socket.SHUT_WR)
def echo(port=0):
""" This function implements a simple echo server for testing the
Forwarder class.
:param int port: port number on which to listen
We run this function and it prints out the listening socket binding.
Then, we run Forwarder and point it at this echo "server".
Then, we run telnet and point it at forwarder and see if whatever we
type gets echoed back to us.
This function exits when the remote end connects, then closes connection
"""
lsock = socket.socket()
lsock.bind(("", port))
lsock.listen(1)
_trace("Listening on sockname=%s", lsock.getsockname())
sock, remote_addr = lsock.accept()
try:
_trace("Connection from peer=%s", remote_addr)
while True:
try:
data = sock.recv(4 * 1024) # pylint: disable=E1101
except pika.compat.SOCKET_ERROR as exc:
if exc.errno == errno.EINTR:
continue
else:
raise
if not data:
break
sock.sendall(data) # pylint: disable=E1101
finally:
try:
_safe_shutdown_socket(sock, socket.SHUT_RDWR)
finally:
sock.close()
def _safe_shutdown_socket(sock, how=socket.SHUT_RDWR):
""" Shutdown a socket, suppressing ENOTCONN
"""
try:
sock.shutdown(how)
except pika.compat.SOCKET_ERROR as exc:
if exc.errno != errno.ENOTCONN:
raise