async def test_proxy_subscription(self):
socket_id = uuid.uuid1()
async with AsyncExitStack() as exit_stack:
remote_sock, remote_port = testing.bind_unused_port()
remote_server = testing_utils.EchoServer()
remote_server.add_sockets([remote_sock])
# Register the remote server to stop on exit.
exit_stack.callback(remote_server.stop)
# Make a request to proxy traffic from a local port to the remote
# server.
state = await self.ws_connect()
async with AsyncExitStack() as sub_stack:
local_proxy = BufferedProxyContext(state, socket_id)
await sub_stack.enter_async_context(local_proxy)
sub = json_protocol.setup_subscription(
state, "proxy_socket",
dict(host="localhost",
port=remote_port,
protocol="tcp",
socket_id=socket_id.hex,
buffsize=1))
await sub_stack.enter_async_context(sub)
msg = await sub.next()
self.assertIn('connection_status', msg)
self.assertEqual('connected', msg['connection_status'])
self.assertIsNotNone(msg['socket_id'])
self.assertEqual(0, msg['count'])
async def run_socks_proxy(cxn, port, cxn_id=None):
args = dict(port=port)
if cxn_id:
args['cxn_id'] = cxn_id
await cxn.open()
async with setup_subscription(cxn.state, 'socks5_proxy', args) as sub:
async for msg in sub.result_generator():
click.echo(msg)
async def test_multiple_subscriptions(self):
state = await self.ws_connect()
self.assertIsNotNone(state)
sub1 = json_protocol.setup_subscription(state, 'count',
dict(count=2, timeout=0.001))
sub2 = json_protocol.setup_subscription(state, 'count',
dict(count=2, timeout=0.001))
async with sub1, sub2:
msg1 = await sub1.next()
self.assertEqual(2, msg1)
msg1 = await sub1.next()
self.assertEqual(1, msg1)
msg2 = await sub2.next()
self.assertEqual(2, msg2)
msg2 = await sub2.next()
self.assertEqual(1, msg2)
async def run_socks_proxy(port, state):
"""Run a socks proxy on the given port and connection.
This proxies all traffic from this socks proxy over the given
connection (as defined by the passed in WebsocketState).
"""
try:
async with setup_subscription(state, 'socks5_proxy',
dict(port=port)) as sub:
async for msg in sub.result_generator():
logger.info("%s socks5 handler: %s", state.cxn_id, msg)
except Exception:
logger.exception("Error starting SOCSKv5 proxy!")
async def test_basic_subscription(self):
state = await self.ws_connect()
self.assertIsNotNone(state)
# This handler counts down to 0. The messages should start at 3 and end at 0.
expected = 3
sub = json_protocol.setup_subscription(state, 'count',
dict(count=3, timeout=0.001))
async with sub:
async for msg in sub.result_generator():
self.assertFalse(expected < 0)
self.assertEqual(expected, msg)
expected -= 1
# Try again, with a long timeout, but unsubscribe first.
sub = json_protocol.setup_subscription(state, 'count',
dict(count=3, timeout=10))
async with sub:
await sub.close()
async for msg in sub.result_generator():
self.fail(
"Should not return any messages, because this was unsubscribed."
)
# Try yet again, with a short timeout, and a few calls before unsubscribing.
sub = json_protocol.setup_subscription(state, 'count',
dict(count=5, timeout=0.001))
async with sub:
count = 5
async for msg in sub.result_generator():
self.assertEqual(count, msg)
if count > 3:
count -= 1
else:
await sub.close()
self.assertEqual(3, count)
self.assertEqual(3, msg)
async def test_proxy_socket(self):
async with contextlib.AsyncExitStack() as exit_stack:
# Setup a remote server.
remote_sock, remote_port = testing.bind_unused_port()
remote_server = EchoServer()
remote_server.add_sockets([remote_sock])
# Register the remote server to stop on exit.
exit_stack.callback(remote_server.stop)
# Make a request to proxy traffic from a local port to the remote
# server.
state = await self.ws_connect()
socket_id = uuid.uuid1()
async with json_request.setup_subscription(
state, "proxy_socket",
dict(host="localhost",
port=remote_port,
protocol="tcp",
socket_id=socket_id.hex,
buffsize=1)) as sub:
self.assertIn(sub.msg_id, state.msg_mapping)
msg = await sub.next()
self.assertIn('connection_status', msg)
self.assertEqual('connected', msg['connection_status'])
self.assertIsNotNone(msg['socket_id'])
# The socket_id should be passed, to indicate the ID to used.
socket_id = uuid.UUID(msg['socket_id'])
remote_socket = tunnel_routes.TcpRemoteTunneledSocket(
state, socket_id)
await remote_socket.send(b"Hello World\n")
result = await remote_socket.receive()
self.assertEqual(b"Hello World\n", result)
# We are expecting a message that the socket has closed.
msg = await sub.next()
self.assertIn('connection_status', msg)
self.assertEqual('closed', msg['connection_status'])
# The sub should be closed. Assert that the socket was cleaned up.
self.assertNotIn(sub.msg_id, state.msg_mapping)
async def open(self):
"""Open the remote socket and setup the applicable structures."""
async with AsyncExitStack() as exit_stack:
self._proxy_stream = RawProxyStreamContext(
self.state,
self._local_stream,
self.socket_id,
buffsize=self._buff_size)
self.state.add_proxy_socket(self.socket_id, self._proxy_stream)
exit_stack.callback(self.state.remove_proxy_socket, self.socket_id)
self._monitor_sub = setup_subscription(
self.state, "proxy_socket",
dict(protocol="tcp",
host=self.host,
port=self.port,
socket_id=self.socket_id.hex))
await self._monitor_sub.open()
exit_stack.push_async_callback(self._monitor_sub.close)
# Wait for a response to see if we connected.
#If not, raise an Exception.
msg = await self._monitor_sub.next()
status = msg.get('connection_status')
if status != 'connected':
raise Exception("Error connecting.")
socket_id = uuid.UUID(msg['socket_id'])
self.bind_host = msg['bind_host']
self.bind_port = msg['bind_port']
# Now, schedule the monitor sub to update in the background with
# the current value of the received bytes.
monitor_fut = asyncio.create_task(self._monitor_update())
exit_stack.callback(monitor_fut.cancel)
# Everything is setup correctly, so do not run the exit_stack
# callbacks.
self._close_context = exit_stack.pop_all()
async def proxy_socket_subscription(endpoint, args):
# Parse the socket configuration.
try:
# Assume TCP by default, as this is the most common.
protocol = args.get('protocol', 'tcp')
host = args['host']
port = args['port']
try:
buffsize = int(args.get('buffsize', DEFAULT_BUFFSIZE))
except Exception:
raise KeyError('buffsize')
try:
socket_id = uuid.UUID(args['socket_id'])
except Exception:
raise KeyError('socket_id')
# For now, just handle TCP
assert protocol.lower() == 'tcp'
# Check the proxy request against the current auth manager.
endpoint.state.auth_manager.check_proxy_request(host, port, protocol)
except KeyError as ke:
await endpoint.error("Missing or Invalid field: {}".format(ke))
return
except NotAuthorized:
await endpoint.error("Not Authorized.")
return
except Exception:
logger.exception("Error in proxy_socket subscription.")
await endpoint.error("Internal Error")
return
state = endpoint.state
# Connect to the socket with the parsed information.
try:
resolver = AuthManagerResolver(netutil.DefaultExecutorResolver(),
endpoint.state.auth_manager)
client = tcpclient.TCPClient(resolver)
local_stream = await client.connect(host, port)
res = local_stream.socket.getsockname()
bind_host = res[0]
bind_port = res[1]
except iostream.StreamClosedError as exc:
await endpoint.error(
dict(connection_status="error",
message="Could not connect; stream is not open!"))
return
except Exception:
logger.exception("Error opening connection to %s:%s", host, port)
await endpoint.error(
dict(connection_status="error", message="Could not connect."))
return
logger.info("Setup proxy (ID: %s) to: %s:%d", socket_id.hex, host, port)
try:
async with AsyncExitStack() as exit_stack:
# Next, check whether the proxy socket monitor is valid. This is needed
# for throttling; the caller of this sub should appropriately setup the
# proxy as well, or else this should fail.
monitor_sub = setup_subscription(state, 'proxy_socket_monitor',
dict(socket_id=socket_id.hex))
await monitor_sub.open()
exit_stack.push_async_callback(monitor_sub.close)
# Wait for a response from the monitor sub, for some timeout.
try:
msg = await asyncio.wait_for(monitor_sub.next(), 5)
# This should be defined, but just assume 0 in the case of a
# malformed response.
count = msg.get('count', 0)
except asyncio.TimeoutError:
await endpoint.error(
"Did not receive a response from the monitor!")
return
except SubscriptionError as e:
await endpoint.error("Unexpected error: {}".format(e))
except Exception:
logger.exception(
"Unexpected exception in proxy monitor setup!")
await endpoint.error("Unexpected error!")
return
proxy_stream = RawProxyStreamContext(state,
local_stream,
socket_id,
buffsize=buffsize)
await proxy_stream.open()
exit_stack.push_async_callback(proxy_stream.close)
# Now that everything is setup, send the connection over
# Now that the local proxy is setup, send the connection over.
await endpoint.next(
dict(connection_status="connected",
socket_id=socket_id.hex,
bind_host=bind_host,
bind_port=bind_port,
count=0))
monitor_fut = asyncio.create_task(proxy_stream.run(monitor_sub))
write_fut = asyncio.create_task(
_write_monitor_updates(endpoint, proxy_stream))
exit_stack.callback(monitor_fut.cancel)
exit_stack.callback(write_fut.cancel)
await asyncio.gather(monitor_fut, write_fut)
await endpoint.next(
dict(connection_status="closed", socket_id=socket_id.hex))
except (SubscriptionComplete, iostream.StreamClosedError):
logger.info("Closing proxy socket.")