From 9207458fd771e402cc98f67806dc83f5a0f79a0c Mon Sep 17 00:00:00 2001 From: Harshil Patel Date: Thu, 28 Mar 2024 14:30:02 -0700 Subject: [PATCH] stdlib: add socks proxy to atlas client (#864) --- src/python/SConscript | 2 + .../gem5/resources/client_api/atlasclient.py | 3 +- src/python/gem5/resources/downloader.py | 28 ++------- src/python/gem5/utils/socks_ssl_context.py | 58 +++++++++++++++++++ 4 files changed, 68 insertions(+), 23 deletions(-) create mode 100644 src/python/gem5/utils/socks_ssl_context.py diff --git a/src/python/SConscript b/src/python/SConscript index 49bc616ae2..1016fee5c6 100644 --- a/src/python/SConscript +++ b/src/python/SConscript @@ -306,6 +306,8 @@ PySource('gem5.utils', 'gem5/utils/filelock.py') PySource('gem5.utils', 'gem5/utils/override.py') PySource('gem5.utils', 'gem5/utils/progress_bar.py') PySource('gem5.utils', 'gem5/utils/requires.py') +PySource('gem5.utils', + 'gem5/utils/socks_ssl_context.py') PySource('gem5.utils.multiprocessing', 'gem5/utils/multiprocessing/__init__.py') PySource('gem5.utils.multiprocessing', diff --git a/src/python/gem5/resources/client_api/atlasclient.py b/src/python/gem5/resources/client_api/atlasclient.py index 1a4f9737c9..83933f5497 100644 --- a/src/python/gem5/resources/client_api/atlasclient.py +++ b/src/python/gem5/resources/client_api/atlasclient.py @@ -43,6 +43,7 @@ from urllib import ( from m5.util import warn +from ...utils.socks_ssl_context import get_proxy_context from .abstract_client import AbstractClient @@ -133,7 +134,7 @@ class AtlasClient(AbstractClient): for attempt in itertools.count(start=1): try: - response = request.urlopen(req) + response = request.urlopen(req, context=get_proxy_context()) break except Exception as e: if attempt >= max_failed_attempts: diff --git a/src/python/gem5/resources/downloader.py b/src/python/gem5/resources/downloader.py index 1d495509ee..bb6d7109cd 100644 --- a/src/python/gem5/resources/downloader.py +++ b/src/python/gem5/resources/downloader.py @@ -48,6 +48,7 @@ from ..utils.progress_bar import ( progress_hook, tqdm, ) +from ..utils.socks_ssl_context import get_proxy_context from .client import get_resource_json_obj from .client import list_resources as client_list_resources from .md5_utils import ( @@ -87,30 +88,13 @@ def _download(url: str, download_to: str, max_attempts: int = 6) -> None: # number of download attempts has been reached or if a HTTP status code # other than 408, 429, or 5xx is received. try: - # check to see if user requests a proxy connection - use_proxy = os.getenv("GEM5_USE_PROXY") - if use_proxy: - # If the "use_proxy" variable is specified we setup a socks5 - # connection. - - import socket - import ssl - - import socks - - IP_ADDR, host_port = use_proxy.split(":") - PORT = int(host_port) - socks.set_default_proxy(socks.SOCKS5, IP_ADDR, PORT) - socket.socket = socks.socksocket - - # base SSL context for https connection - ctx = ssl.create_default_context() - ctx.check_hostname = False - ctx.verify_mode = ssl.CERT_NONE - + proxy_context = get_proxy_context() + if proxy_context: # get the file as a bytes blob request = urllib.request.Request(url) - with urllib.request.urlopen(request, context=ctx) as fr: + with urllib.request.urlopen( + request, context=proxy_context + ) as fr: with tqdm.wrapattr( open(download_to, "wb"), "write", diff --git a/src/python/gem5/utils/socks_ssl_context.py b/src/python/gem5/utils/socks_ssl_context.py new file mode 100644 index 0000000000..3881adb3e1 --- /dev/null +++ b/src/python/gem5/utils/socks_ssl_context.py @@ -0,0 +1,58 @@ +# Copyright (c) 2024 The Regents of the University of California +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are +# met: redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer; +# redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution; +# neither the name of the copyright holders nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import os +import ssl +from typing import Optional + +_gem5_ssl_context = None + + +def get_proxy_context() -> Optional[ssl.SSLContext]: + """ + This function returns a SSL context for https connection with SOCKS proxy + support. It uses the environment variable GEM5_USE_PROXY to determine + the proxy server to use. If the environment variable is not set, it + returns None. + """ + + global _gem5_ssl_context + use_proxy = os.getenv("GEM5_USE_PROXY") + if use_proxy and not _gem5_ssl_context: + import socket + + import socks + + ip_addr, host_port = use_proxy.split(":") + port = int(host_port) + socks.set_default_proxy(socks.SOCKS5, ip_addr, port) + socket.socket = socks.socksocket + + # base SSL context for https connection + _gem5_ssl_context = ssl.create_default_context() + _gem5_ssl_context.check_hostname = False + _gem5_ssl_context.verify_mode = ssl.CERT_NONE + return _gem5_ssl_context