tests: add unit test for portage.util.socks5 (FEATURES=network-sandbox-proxy)

Bug: https://bugs.gentoo.org/604474
Signed-off-by: Zac Medico <zmedico@gentoo.org>

2 files changed, 256 insertions, 3 deletions

diff --git a/lib/portage/tests/util/test_socks5.py b/lib/portage/tests/util/test_socks5.py
new file mode 100644
index 000000000..5db85b0a6
--- /dev/null
+++ b/lib/portage/tests/util/test_socks5.py
@@ -0,0 +1,211 @@
+# Copyright 2019 Gentoo Authors
+# Distributed under the terms of the GNU General Public License v2
+
+import functools
+import platform
+import shutil
+import socket
+import struct
+import sys
+import tempfile
+import time
+
+import portage
+from portage.tests import TestCase
+from portage.util._eventloop.global_event_loop import global_event_loop
+from portage.util import socks5
+from portage.const import PORTAGE_BIN_PATH
+
+try:
+	from http.server import BaseHTTPRequestHandler, HTTPServer
+except ImportError:
+	from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
+
+try:
+	from urllib.request import urlopen
+except ImportError:
+	from urllib import urlopen
+
+
+class _Handler(BaseHTTPRequestHandler):
+
+	def __init__(self, content, *args, **kwargs):
+		self.content = content
+		BaseHTTPRequestHandler.__init__(self, *args, **kwargs)
+
+	def do_GET(self):
+		doc = self.send_head()
+		if doc is not None:
+			self.wfile.write(doc)
+
+	def do_HEAD(self):
+		self.send_head()
+
+	def send_head(self):
+		doc = self.content.get(self.path)
+		if doc is None:
+			self.send_error(404, "File not found")
+			return None
+
+		self.send_response(200)
+		self.send_header("Content-type", "text/plain")
+		self.send_header("Content-Length", len(doc))
+		self.send_header("Last-Modified", self.date_time_string(time.time()))
+		self.end_headers()
+		return doc
+
+	def log_message(self, fmt, *args):
+		pass
+
+
+class AsyncHTTPServer(object):
+	def __init__(self, host, content, loop):
+		self._host = host
+		self._content = content
+		self._loop = loop
+		self.server_port = None
+		self._httpd = None
+
+	def __enter__(self):
+		httpd = self._httpd = HTTPServer((self._host, 0), functools.partial(_Handler, self._content))
+		self.server_port = httpd.server_port
+		self._loop.add_reader(httpd.socket.fileno(), self._httpd._handle_request_noblock)
+		return self
+
+	def __exit__(self, exc_type, exc_value, exc_traceback):
+		if self._httpd is not None:
+			self._loop.remove_reader(self._httpd.socket.fileno())
+			self._httpd.socket.close()
+			self._httpd = None
+
+
+class AsyncHTTPServerTestCase(TestCase):
+
+	@staticmethod
+	def _fetch_directly(host, port, path):
+		# NOTE: python2.7 does not have context manager support here
+		try:
+			f = urlopen('http://{host}:{port}{path}'.format( # nosec
+				host=host, port=port, path=path))
+			return f.read()
+		finally:
+			if f is not None:
+				f.close()
+
+	def test_http_server(self):
+		host = '127.0.0.1'
+		content = b'Hello World!\n'
+		path = '/index.html'
+		loop = global_event_loop()
+		for i in range(2):
+			with AsyncHTTPServer(host, {path: content}, loop) as server:
+				for j in range(2):
+					result = loop.run_until_complete(loop.run_in_executor(None,
+						self._fetch_directly, host, server.server_port, path))
+					self.assertEqual(result, content)
+
+
+class _socket_file_wrapper(portage.proxy.objectproxy.ObjectProxy):
+	"""
+	A file-like object that wraps a socket and closes the socket when
+	closed. Since python2.7 does not support socket.detach(), this is a
+	convenient way to have a file attached to a socket that closes
+	automatically (without resource warnings about unclosed sockets).
+	"""
+
+	__slots__ = ('_file', '_socket')
+
+	def __init__(self, socket, f):
+		object.__setattr__(self, '_socket', socket)
+		object.__setattr__(self, '_file', f)
+
+	def _get_target(self):
+		return object.__getattribute__(self, '_file')
+
+	def __getattribute__(self, attr):
+		if attr == 'close':
+			return object.__getattribute__(self, 'close')
+		return super(_socket_file_wrapper, self).__getattribute__(attr)
+
+	def __enter__(self):
+		return self
+
+	def close(self):
+		object.__getattribute__(self, '_file').close()
+		object.__getattribute__(self, '_socket').close()
+
+	def __exit__(self, exc_type, exc_value, traceback):
+		self.close()
+
+
+def socks5_http_get_ipv4(proxy, host, port, path):
+	"""
+	Open http GET request via socks5 proxy listening on a unix socket,
+	and return a file to read the response body from.
+	"""
+	s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+	f = _socket_file_wrapper(s, s.makefile('rb', 1024))
+	try:
+		s.connect(proxy)
+		s.send(struct.pack('!BBB', 0x05, 0x01, 0x00))
+		vers, method = struct.unpack('!BB', s.recv(2))
+		s.send(struct.pack('!BBBB', 0x05, 0x01, 0x00, 0x01))
+		s.send(socket.inet_pton(socket.AF_INET, host))
+		s.send(struct.pack('!H', port))
+		reply = struct.unpack('!BBB', s.recv(3))
+		if reply != (0x05, 0x00, 0x00):
+			raise AssertionError(repr(reply))
+		struct.unpack('!B4sH', s.recv(7)) # contains proxied address info
+		s.send("GET {} HTTP/1.1\r\nHost: {}:{}\r\nAccept: */*\r\nConnection: close\r\n\r\n".format(
+			path, host, port).encode())
+		headers = []
+		while True:
+			headers.append(f.readline())
+			if headers[-1] == b'\r\n':
+				return f
+	except Exception:
+		f.close()
+		raise
+
+
+class Socks5ServerTestCase(TestCase):
+
+	@staticmethod
+	def _fetch_via_proxy(proxy, host, port, path):
+		with socks5_http_get_ipv4(proxy, host, port, path) as f:
+			return f.read()
+
+	def test_socks5_proxy(self):
+
+		loop = global_event_loop()
+
+		host = '127.0.0.1'
+		content = b'Hello World!'
+		path = '/index.html'
+		proxy = None
+		tempdir = tempfile.mkdtemp()
+
+		try:
+			with AsyncHTTPServer(host, {path: content}, loop) as server:
+
+				settings = {
+					'PORTAGE_TMPDIR': tempdir,
+					'PORTAGE_BIN_PATH': PORTAGE_BIN_PATH,
+				}
+
+				try:
+					proxy = socks5.get_socks5_proxy(settings)
+				except NotImplementedError:
+					# bug 658172 for python2.7
+					self.skipTest('get_socks5_proxy not implemented for {} {}.{}'.format(
+						platform.python_implementation(), *sys.version_info[:2]))
+				else:
+					loop.run_until_complete(socks5.proxy.ready())
+
+					result = loop.run_until_complete(loop.run_in_executor(None,
+						self._fetch_via_proxy, proxy, host, server.server_port, path))
+
+					self.assertEqual(result, content)
+		finally:
+			socks5.proxy.stop()
+			shutil.rmtree(tempdir)
diff --git a/lib/portage/util/socks5.py b/lib/portage/util/socks5.py
index 74b0714eb..59e6699ec 100644
--- a/lib/portage/util/socks5.py
+++ b/lib/portage/util/socks5.py
@@ -1,13 +1,18 @@
 # SOCKSv5 proxy manager for network-sandbox
-# Copyright 2015 Gentoo Foundation
+# Copyright 2015-2019 Gentoo Authors
 # Distributed under the terms of the GNU General Public License v2
 
+import errno
 import os
 import signal
+import socket
 
+import portage.data
 from portage import _python_interpreter
 from portage.data import portage_gid, portage_uid, userpriv_groups
 from portage.process import atexit_register, spawn
+from portage.util.futures.compat_coroutine import coroutine
+from portage.util.futures import asyncio
 
 
 class ProxyManager(object):
@@ -36,9 +41,16 @@ class ProxyManager(object):
 		self.socket_path = os.path.join(settings['PORTAGE_TMPDIR'],
 				'.portage.%d.net.sock' % os.getpid())
 		server_bin = os.path.join(settings['PORTAGE_BIN_PATH'], 'socks5-server.py')
+		spawn_kwargs = {}
+		# The portage_uid check solves EPERM failures in Travis CI.
+		if portage.data.secpass > 1 and os.geteuid() != portage_uid:
+			spawn_kwargs.update(
+				uid=portage_uid,
+				gid=portage_gid,
+				groups=userpriv_groups,
+				umask=0o077)
 		self._pids = spawn([_python_interpreter, server_bin, self.socket_path],
-				returnpid=True, uid=portage_uid, gid=portage_gid,
-				groups=userpriv_groups, umask=0o077)
+				returnpid=True, **spawn_kwargs)
 
 	def stop(self):
 		"""
@@ -60,6 +72,36 @@ class ProxyManager(object):
 		return self.socket_path is not None
 
 
+	@coroutine
+	def ready(self):
+		"""
+		Wait for the proxy socket to become ready. This method is a coroutine.
+		"""
+
+		while True:
+			try:
+				wait_retval = os.waitpid(self._pids[0], os.WNOHANG)
+			except OSError as e:
+				if e.errno == errno.EINTR:
+					continue
+				raise
+
+			if wait_retval is not None and wait_retval != (0, 0):
+				raise OSError(3, 'No such process')
+
+			try:
+				s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+				s.connect(self.socket_path)
+			except EnvironmentError as e:
+				if e.errno != errno.ENOENT:
+					raise
+				yield asyncio.sleep(0.2)
+			else:
+				break
+			finally:
+				s.close()
+
+
 proxy = ProxyManager()