From de5c4f4b130f4d52747c131ed9fad601d2a478d3 Mon Sep 17 00:00:00 2001 From: nibrag Date: Mon, 16 May 2016 17:16:44 +0300 Subject: [PATCH] Use contextmanager for fake_socks_srv --- tests/helpers.py | 39 ++-- tests/test_functional.py | 414 ++++++++++++++++----------------------- 2 files changed, 195 insertions(+), 258 deletions(-) diff --git a/tests/helpers.py b/tests/helpers.py index 7d81462..a379c56 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,6 +1,8 @@ import asyncio +import contextlib import socket from unittest import mock +import gc try: from asyncio import ensure_future except ImportError: @@ -24,28 +26,33 @@ def find_unused_port(): return port -class SocksPrimitiveProtocol(asyncio.Protocol): - def __init__(self, write_buff): - self._write_buff = write_buff - self._transport = None - - def connection_made(self, transport): - self._transport = transport +@contextlib.contextmanager +def fake_socks_srv(loop, write_buff): + transports = [] - def data_received(self, data): - self._transport.write(self._write_buff) + class SocksPrimitiveProtocol(asyncio.Protocol): + _transport = None - def connection_lost(self, exc): - self._transport.close() + def connection_made(self, transport): + self._transport = transport + transports.append(transport) + def data_received(self, data): + self._transport.write(write_buff) -@asyncio.coroutine -def fake_socks_srv(loop, write_buff): port = find_unused_port() def factory(): - return SocksPrimitiveProtocol(write_buff) + return SocksPrimitiveProtocol() + + srv = loop.run_until_complete( + loop.create_server(factory, '127.0.0.1', port)) + + yield port - server = yield from loop.create_server(factory, '127.0.0.1', port) - return server, port + for tr in transports: + tr.close() + srv.close() + loop.run_until_complete(srv.wait_closed()) + gc.collect() diff --git a/tests/test_functional.py b/tests/test_functional.py index 0ce165f..2ec0d0c 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -1,16 +1,13 @@ import unittest -import aiohttp import aiosocks import asyncio -from aiohttp import web -from aiosocks.connector import SocksConnector try: from asyncio import ensure_future except ImportError: ensure_future = asyncio.async -from .helpers import fake_socks_srv, find_unused_port +from .helpers import fake_socks_srv class TestCreateSocks4Connection(unittest.TestCase): @@ -22,59 +19,49 @@ class TestCreateSocks4Connection(unittest.TestCase): self.loop.close() def test_connect_success(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x00\x5a\x04W\x01\x01\x01\x01test') - ) - addr = aiosocks.Socks4Addr('127.0.0.1', port) - auth = aiosocks.Socks4Auth('usr') - dst = ('python.org', 80) + with fake_socks_srv(self.loop, + b'\x00\x5a\x04W\x01\x01\x01\x01test') as port: + addr = aiosocks.Socks4Addr('127.0.0.1', port) + auth = aiosocks.Socks4Auth('usr') + dst = ('python.org', 80) - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - transport, protocol = self.loop.run_until_complete(coro) + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + transport, protocol = self.loop.run_until_complete(coro) - self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) + self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) - data = self.loop.run_until_complete(protocol._stream_reader.read(4)) - self.assertEqual(data, b'test') + data = self.loop.run_until_complete( + protocol._stream_reader.read(4)) + self.assertEqual(data, b'test') - transport.close() - server.close() - self.loop.run_until_complete(server.wait_closed()) + transport.close() def test_invalid_data(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x01\x5a\x04W\x01\x01\x01\x01') - ) - addr = aiosocks.Socks4Addr('127.0.0.1', port) - auth = aiosocks.Socks4Auth('usr') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('invalid data', str(ct.exception)) - - server.close() - self.loop.run_until_complete(server.wait_closed()) + with fake_socks_srv(self.loop, + b'\x01\x5a\x04W\x01\x01\x01\x01') as port: + addr = aiosocks.Socks4Addr('127.0.0.1', port) + auth = aiosocks.Socks4Auth('usr') + dst = ('python.org', 80) + + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('invalid data', str(ct.exception)) def test_socks_srv_error(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x00\x5b\x04W\x01\x01\x01\x01') - ) - addr = aiosocks.Socks4Addr('127.0.0.1', port) - auth = aiosocks.Socks4Auth('usr') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('0x5b', str(ct.exception)) + with fake_socks_srv(self.loop, + b'\x00\x5b\x04W\x01\x01\x01\x01') as port: + addr = aiosocks.Socks4Addr('127.0.0.1', port) + auth = aiosocks.Socks4Auth('usr') + dst = ('python.org', 80) - server.close() - self.loop.run_until_complete(server.wait_closed()) + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('0x5b', str(ct.exception)) class TestCreateSocks5Connect(unittest.TestCase): @@ -86,250 +73,193 @@ class TestCreateSocks5Connect(unittest.TestCase): self.loop.close() def test_connect_success_anonymous(self): - server, port = self.loop.run_until_complete( - fake_socks_srv( + with fake_socks_srv( self.loop, - b'\x05\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest' - ) - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) + b'\x05\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - transport, protocol = self.loop.run_until_complete(coro) + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + transport, protocol = self.loop.run_until_complete(coro) - self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) + self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) - data = self.loop.run_until_complete(protocol._stream_reader.read(4)) - self.assertEqual(data, b'test') + data = self.loop.run_until_complete( + protocol._stream_reader.read(4)) + self.assertEqual(data, b'test') - transport.close() - server.close() - self.loop.run_until_complete(server.wait_closed()) + transport.close() def test_connect_success_usr_pwd(self): - server, port = self.loop.run_until_complete( - fake_socks_srv( + with fake_socks_srv( self.loop, b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest' - ) - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - transport, protocol = self.loop.run_until_complete(coro) + ) as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + transport, protocol = self.loop.run_until_complete(coro) - data = self.loop.run_until_complete(protocol._stream_reader.read(4)) - self.assertEqual(data, b'test') + self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) - transport.close() - server.close() - self.loop.run_until_complete(server.wait_closed()) + data = self.loop.run_until_complete( + protocol._stream_reader.read(4)) + self.assertEqual(data, b'test') + transport.close() def test_auth_ver_err(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x04\x02') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('invalid version', str(ct.exception)) + with fake_socks_srv(self.loop, b'\x04\x02') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - server.close() - self.loop.run_until_complete(server.wait_closed()) + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('invalid version', str(ct.exception)) def test_auth_method_rejected(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x05\xFF') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('authentication methods were rejected', - str(ct.exception)) - - server.close() - self.loop.run_until_complete(server.wait_closed()) + with fake_socks_srv(self.loop, b'\x05\xFF') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) + + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('authentication methods were rejected', + str(ct.exception)) def test_auth_status_invalid(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x05\xF0') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('invalid data', str(ct.exception)) + with fake_socks_srv(self.loop, b'\x05\xF0') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - server.close() - self.loop.run_until_complete(server.wait_closed()) + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('invalid data', str(ct.exception)) def test_auth_status_invalid2(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x05\x02\x02\x00') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('invalid data', str(ct.exception)) + with fake_socks_srv(self.loop, b'\x05\x02\x02\x00') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - server.close() - self.loop.run_until_complete(server.wait_closed()) + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('invalid data', str(ct.exception)) def test_auth_failed(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x05\x02\x01\x01') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('authentication failed', str(ct.exception)) + with fake_socks_srv(self.loop, b'\x05\x02\x01\x01') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - server.close() - self.loop.run_until_complete(server.wait_closed()) + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('authentication failed', str(ct.exception)) def test_cmd_ver_err(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x05\x02\x01\x00\x04\x00\x00') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('invalid version', str(ct.exception)) - - server.close() - self.loop.run_until_complete(server.wait_closed()) + with fake_socks_srv(self.loop, + b'\x05\x02\x01\x00\x04\x00\x00') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) + + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('invalid version', str(ct.exception)) def test_cmd_not_granted(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x05\x02\x01\x00\x05\x01\x00') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('General SOCKS server failure', str(ct.exception)) - - server.close() - self.loop.run_until_complete(server.wait_closed()) + with fake_socks_srv(self.loop, + b'\x05\x02\x01\x00\x05\x01\x00') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) + + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('General SOCKS server failure', str(ct.exception)) def test_invalid_address_type(self): - server, port = self.loop.run_until_complete( - fake_socks_srv(self.loop, b'\x05\x02\x01\x00\x05\x00\x00\xFF') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - with self.assertRaises(aiosocks.SocksError) as ct: - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - self.loop.run_until_complete(coro) - self.assertIn('invalid data', str(ct.exception)) - - server.close() - self.loop.run_until_complete(server.wait_closed()) + with fake_socks_srv(self.loop, + b'\x05\x02\x01\x00\x05\x00\x00\xFF') as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) + + with self.assertRaises(aiosocks.SocksError) as ct: + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + self.loop.run_until_complete(coro) + self.assertIn('invalid data', str(ct.exception)) def test_atype_ipv4(self): - server, port = self.loop.run_until_complete( - fake_socks_srv( + with fake_socks_srv( self.loop, b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04W' - ) - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) + ) as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - transport, protocol = self.loop.run_until_complete(coro) + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + transport, protocol = self.loop.run_until_complete(coro) - self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) + self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) - transport.close() - server.close() - self.loop.run_until_complete(server.wait_closed()) + transport.close() def test_atype_ipv6(self): - server, port = self.loop.run_until_complete( - fake_socks_srv( + with fake_socks_srv( self.loop, b'\x05\x02\x01\x00\x05\x00\x00\x04\x00\x00\x00\x00' - b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x11\x04W') - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x11\x04W' + ) as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - transport, protocol = self.loop.run_until_complete(coro) + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + transport, protocol = self.loop.run_until_complete(coro) - self.assertEqual(protocol.proxy_sockname, ('::111', 1111)) + self.assertEqual(protocol.proxy_sockname, ('::111', 1111)) - transport.close() - server.close() - self.loop.run_until_complete(server.wait_closed()) + transport.close() def test_atype_domain(self): - server, port = self.loop.run_until_complete( - fake_socks_srv( + with fake_socks_srv( self.loop, b'\x05\x02\x01\x00\x05\x00\x00\x03\x0apython.org\x04W' - ) - ) - addr = aiosocks.Socks5Addr('127.0.0.1', port) - auth = aiosocks.Socks5Auth('usr', 'pwd') - dst = ('python.org', 80) - - coro = aiosocks.create_connection( - None, addr, auth, dst, loop=self.loop) - transport, protocol = self.loop.run_until_complete(coro) + ) as port: + addr = aiosocks.Socks5Addr('127.0.0.1', port) + auth = aiosocks.Socks5Auth('usr', 'pwd') + dst = ('python.org', 80) - self.assertEqual(protocol.proxy_sockname, (b'python.org', 1111)) + coro = aiosocks.create_connection( + None, addr, auth, dst, loop=self.loop) + transport, protocol = self.loop.run_until_complete(coro) - transport.close() - server.close() - self.loop.run_until_complete(server.wait_closed()) + self.assertEqual(protocol.proxy_sockname, (b'python.org', 1111)) + transport.close()