From ad89019e8a968b8ee02e323d7e70789ef9cfc461 Mon Sep 17 00:00:00 2001 From: cormick Date: Tue, 27 May 2025 19:38:48 +0800 Subject: [PATCH] feat: support setting proxy for reggie client Signed-off-by: cormick --- CHANGELOG.md | 1 + .../distribution/reggie/__init__.py | 1 + opencontainers/distribution/reggie/client.py | 15 +++ opencontainers/distribution/reggie/request.py | 25 +++++ opencontainers/tests/test_distribution.py | 91 +++++++++++++++++++ opencontainers/version.py | 2 +- 6 files changed, 134 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9250bb9..5ffab71 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ Critical items to know are: Versions here coincide with releases on pypi. ## [master](https://github.com/vsoch/oci-python) + - support setting proxy for reggie client (0.0.15) - do not set basic auth if no username/password provided (0.0.14) - allow for update of a structure attribute, if applicable (0.0.13) - fix to bug with parsing www-Authenticate (0.0.12) diff --git a/opencontainers/distribution/reggie/__init__.py b/opencontainers/distribution/reggie/__init__.py index 0a2f01e..71ef515 100644 --- a/opencontainers/distribution/reggie/__init__.py +++ b/opencontainers/distribution/reggie/__init__.py @@ -7,6 +7,7 @@ WithDefaultName, WithDebug, WithUserAgent, + WithProxy, ) from .request import ( WithName, diff --git a/opencontainers/distribution/reggie/client.py b/opencontainers/distribution/reggie/client.py index 63e17e6..b59791f 100644 --- a/opencontainers/distribution/reggie/client.py +++ b/opencontainers/distribution/reggie/client.py @@ -34,6 +34,7 @@ class ClientConfig(BaseConfig): "WithDebug", "WithDefaultName", "WithAuthScope", + "WithProxy", ] def __init__(self, address, opts=None): @@ -48,6 +49,7 @@ def __init__(self, address, opts=None): self.DefaultName = None self.UserAgent = DEFAULT_USER_AGENT self.required = [self.Address, self.UserAgent] + self.Proxy = None super().__init__() def _validate(self): @@ -118,6 +120,17 @@ def WithUserAgent(config): return WithUserAgent +def WithProxy(proxy): + """ + WithProxy sets the proxy configuration setting. + """ + + def WithProxy(config): + config.Proxy = proxy + + return WithProxy + + # Client @@ -178,6 +191,8 @@ def NewRequest(self, method, path, *opts): requestClient.SetUrl(url) requestClient.SetHeader("User-Agent", self.Config.UserAgent) requestClient.SetRetryCallback(rc.RetryCallback) + if self.Config.Proxy: + requestClient.SetProxy(self.Config.Proxy) # Return the Client, which has Request and retryCallback return requestClient diff --git a/opencontainers/distribution/reggie/request.py b/opencontainers/distribution/reggie/request.py index ce45b44..943a27f 100644 --- a/opencontainers/distribution/reggie/request.py +++ b/opencontainers/distribution/reggie/request.py @@ -35,6 +35,7 @@ class RequestConfig(BaseConfig): "WithDigest", "WithSessionID", "WithRetryCallback", + "WithProxy", ] def __init__(self, opts): @@ -46,6 +47,7 @@ def __init__(self, opts): self.Digest = None self.SessionID = None self.RetryCallback = None + self.Proxy = None self.required = [self.Name] super().__init__(opts or []) @@ -108,6 +110,17 @@ def WithRetryCallback(config): return WithRetryCallback +def WithProxy(proxy): + """ + WithProxy sets the proxy configuration setting for requests. + """ + + def WithProxy(config): + config.Proxy = proxy + + return WithProxy + + class RequestClient(requests.Session): """ A Request Client. @@ -244,6 +257,18 @@ def SetBasicAuth(self, username, password): auth_header = base64.b64encode(auth_str.encode("utf-8")) return self.SetHeader("Authorization", "Basic %s" % auth_header.decode("utf-8")) + def SetProxy(self, proxy): + """ + SetProxy sets the proxy configuration setting for requests. + """ + self.proxies.update( + { + "http": proxy, + "https": proxy, + } + ) + return self + def Execute(self, method=None, url=None): """ Execute validates a Request and executes it. diff --git a/opencontainers/tests/test_distribution.py b/opencontainers/tests/test_distribution.py index 07a29e7..7b6131f 100644 --- a/opencontainers/tests/test_distribution.py +++ b/opencontainers/tests/test_distribution.py @@ -11,6 +11,8 @@ import os import re import pytest +from http.server import BaseHTTPRequestHandler, HTTPServer +from threading import Thread # Use the same port across tests @@ -19,18 +21,59 @@ mock_server_thread = None +# Simple HTTP proxy server for testing proxy functionality +class SimpleProxyHandler(BaseHTTPRequestHandler): + """A simple HTTP proxy handler that logs requests""" + + def do_GET(self): + self.send_response(200) + self.send_header("Content-Type", "text/plain") + self.end_headers() + self.wfile.write(b"Request proxied successfully") + print(f"Proxy handled request: {self.path}") + + def do_PUT(self): + self.send_response(200) + self.send_header("Content-Type", "text/plain") + self.end_headers() + self.wfile.write(b"Request proxied successfully") + print(f"Proxy handled PUT request: {self.path}") + + def log_message(self, format, *args): + # Customize logging to show it's from the proxy + print(f"PROXY LOG: {format % args}") + + +# Global variables for proxy server +proxy_port = get_free_port() +proxy_server = None +proxy_server_thread = None + + def setup_module(module): """setup any state specific to the execution of the given module.""" global mock_server global mock_server_thread + global proxy_server + global proxy_server_thread + + # Start the mock registry server mock_server, mock_server_thread = start_mock_server(port) + # Start the proxy server + proxy_server = HTTPServer(("localhost", proxy_port), SimpleProxyHandler) + proxy_server_thread = Thread(target=proxy_server.serve_forever) + proxy_server_thread.setDaemon(True) + proxy_server_thread.start() + print(f"Proxy server started on port {proxy_port}") + def teardown_module(module): """teardown any state that was previously setup with a setup_module method. """ mock_server.server_close() + proxy_server.server_close() def test_distribution_mock_server(tmp_path): @@ -47,6 +90,24 @@ def test_distribution_mock_server(tmp_path): ) assert not client.Config.Debug + print("Testing creation of client with proxy") + proxy_url = f"http://localhost:{proxy_port}" + proxy_client = NewClient( + mock_url, + WithUsernamePassword("testuser", "testpass"), + WithDefaultName("testname"), + WithUserAgent("reggie-tests"), + WithProxy(proxy_url), + ) + assert proxy_client.Config.Proxy == proxy_url + + # Make a request with the proxy client + req = proxy_client.NewRequest("GET", "/v2//tags/list") + response = proxy_client.Do(req) + assert ( + response.status_code == 200 + ), f"Expected status code 200, got {response.status_code}" + print("Testing setting debug option") clientDebug = NewClient(mock_url, WithDebug(True)) assert clientDebug.Config.Debug @@ -189,6 +250,27 @@ def test_distribution_mock_server(tmp_path): print("Check that the body did not get lost somewhere") assert req.body == "abc" + print("Test proxy request with different configuration") + # Create a client with a different proxy configuration + alt_proxy_url = f"http://localhost:{proxy_port}/alternate" + alt_proxy_client = NewClient( + mock_url, + WithProxy(alt_proxy_url), + ) + assert alt_proxy_client.Config.Proxy == alt_proxy_url + + # Verify that proxy setting is correctly passed to the request + proxy_req = alt_proxy_client.NewRequest("GET", "/v2/test/tags/list") + assert ( + proxy_req.proxies + ), "Request should have non-empty proxies dictionary when proxy is set" + assert ( + proxy_req.proxies.get("http") == alt_proxy_url + ), "HTTP proxy not correctly set" + assert ( + proxy_req.proxies.get("https") == alt_proxy_url + ), "HTTPS proxy not correctly set" + print("Test that the retry callback is invoked, if configured.") newBody = "not the original body" @@ -214,3 +296,12 @@ def errorFunc(r): ) except Exception as exc: assert "ruhroh" in str(exc) + + print("Test proxy setting in request client") + # Directly test the SetProxy method on the request client + req = client.NewRequest("GET", "/test/endpoint") + proxy_addr = f"http://localhost:{proxy_port}/direct-test" + req.SetProxy(proxy_addr) + # Verify that the proxy is set in the underlying request object when it's executed + response = client.Do(req) + assert response.status_code == 200, "Request through proxy should succeed" diff --git a/opencontainers/version.py b/opencontainers/version.py index e7031d0..20a5df5 100644 --- a/opencontainers/version.py +++ b/opencontainers/version.py @@ -4,7 +4,7 @@ # Mozilla Public License, v. 2.0. If a copy of the MPL was not distributed # with this file, You can obtain one at http://mozilla.org/MPL/2.0/. -__version__ = "0.0.14" +__version__ = "0.0.15" AUTHOR = "Vanessa Sochat" AUTHOR_EMAIL = "vsoch@users.noreply.github.com" NAME = "opencontainers"