stdlib: add socks proxy to atlas client (#864)

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',

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:

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",

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