From fe056a2ccce5cdd7dd869dd5bbe8893952c4fe94 Mon Sep 17 00:00:00 2001 From: Kenneth Bingham Date: Tue, 28 Jan 2025 18:33:31 -0500 Subject: [PATCH] migrate the proxy example to ProxyShare class so it's a reusable feature of the Py SDK --- sdk/python/examples/proxy/proxy.py | 159 +++++--------------------- sdk/python/sdk/zrok/zrok/README.md | 47 ++++++++ sdk/python/sdk/zrok/zrok/proxy.py | 178 +++++++++++++++++++++++++++++ 3 files changed, 253 insertions(+), 131 deletions(-) create mode 100644 sdk/python/sdk/zrok/zrok/README.md create mode 100644 sdk/python/sdk/zrok/zrok/proxy.py diff --git a/sdk/python/examples/proxy/proxy.py b/sdk/python/examples/proxy/proxy.py index 97f7d2a4..94f3f010 100644 --- a/sdk/python/examples/proxy/proxy.py +++ b/sdk/python/examples/proxy/proxy.py @@ -1,153 +1,50 @@ -import argparse -import atexit -import logging -import sys -import urllib.parse +#!/usr/bin/env python3 -import requests -from flask import Flask, Response, request -from waitress import serve -from zrok.model import ShareRequest, Share -from zrok.overview import EnvironmentAndResources, Overview +""" +Example of using zrok's proxy facility to create an HTTP proxy server. + +This example demonstrates how to: +1. Create a proxy share (optionally with a unique name for persistence) +2. Handle HTTP requests/responses through the proxy +3. Automatically clean up non-reserved shares on exit +""" + +import argparse +import logging import zrok +from zrok.proxy import ProxyShare # Setup logging logging.basicConfig( level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s' + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) -app = Flask(__name__) -target_url = None -zrok_opts = {} -bindPort = 18081 -# List of hop-by-hop headers that should not be returned to the viewer -HOP_BY_HOP_HEADERS = { - 'connection', - 'keep-alive', - 'proxy-authenticate', - 'proxy-authorization', - 'te', - 'trailers', - 'transfer-encoding', - 'upgrade' -} - - -@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH']) -@app.route('/', methods=['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH']) -def proxy(path): - global target_url - logger.info(f"Incoming {request.method} request to {request.path}") - logger.info(f"Headers: {dict(request.headers)}") - - # Forward the request to target URL - full_url = urllib.parse.urljoin(target_url, request.path) - logger.info(f"Forwarding to: {full_url}") - - # Copy request headers, excluding hop-by-hop headers - headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS and k.lower() != 'host'} - - try: - response = requests.request( - method=request.method, - url=full_url, - headers=headers, - data=request.get_data(), - stream=True - ) - - logger.info(f"Response status: {response.status_code}") - logger.info(f"Response headers: {dict(response.headers)}") - - # Filter out hop-by-hop headers from the response - filtered_headers = {k: v for k, v in response.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS} - - return Response( - response.iter_content(chunk_size=8192), - status=response.status_code, - headers=filtered_headers - ) - - except Exception as e: - logger.error(f"Proxy error: {str(e)}", exc_info=True) - return str(e), 502 - - -@zrok.decor.zrok(opts=zrok_opts) -def run_proxy(): - # the port is only used to integrate zrok with frameworks that expect a "hostname:port" combo - serve(app, port=bindPort) - - -if __name__ == '__main__': +def main(): + """Main entry point.""" parser = argparse.ArgumentParser(description='Start a zrok proxy server') parser.add_argument('target_url', help='Target URL to proxy requests to') parser.add_argument('-n', '--unique-name', help='Unique name for the proxy instance') args = parser.parse_args() - target_url = args.target_url logger.info("=== Starting proxy server ===") - logger.info(f"Target URL: {target_url}") - logger.info(f"Logging level: {logger.getEffectiveLevel()}") + logger.info(f"Target URL: {args.target_url}") + # Load environment and create proxy share root = zrok.environment.root.Load() - my_env = EnvironmentAndResources( - environment=None, - shares=[] + proxy_share = ProxyShare.create( + root=root, + target=args.target_url, + unique_name=args.unique_name ) - overview = Overview.create(root=root) - for env_stuff in overview.environments: - if env_stuff.environment.z_id == root.env.ZitiIdentity: - my_env = EnvironmentAndResources( - environment=env_stuff.environment, - shares=env_stuff.shares - ) - break + + # Log access information and start the proxy + logger.info(f"Access proxy at: {', '.join(proxy_share.endpoints)}") + proxy_share.run() - if my_env: - logger.debug( - f"Found environment in overview with Ziti identity " - f"matching local environment: {my_env.environment.z_id}" - ) - else: - logger.error("No matching environment found") - sys.exit(1) - existing_reserved_share = None - for share in my_env.shares: - if share.token == args.unique_name: - existing_reserved_share = share - break - - if existing_reserved_share: - logger.debug(f"Found existing share with token: {existing_reserved_share.token}") - shr = Share(Token=existing_reserved_share.token, FrontendEndpoints=[existing_reserved_share.frontend_endpoint]) - else: - logger.debug(f"No existing share found with token: {args.unique_name}") - share_request = ShareRequest( - BackendMode=zrok.model.PROXY_BACKEND_MODE, - ShareMode=zrok.model.PUBLIC_SHARE_MODE, - Frontends=['public'], - Target="http-proxy", - Reserved=True - ) - if args.unique_name: - share_request.UniqueName = args.unique_name - - shr = zrok.share.CreateShare(root=root, request=share_request) - - def cleanup(): - zrok.share.ReleaseReservedShare(root=root, shr=shr) - logger.info(f"Share {shr.Token} released") - if not args.unique_name: - atexit.register(cleanup) - - zrok_opts['cfg'] = zrok.decor.Opts(root=root, shrToken=shr.Token, bindPort=bindPort) - - logger.info(f"Access proxy at: {', '.join(shr.FrontendEndpoints)}") - - run_proxy() +if __name__ == '__main__': + main() diff --git a/sdk/python/sdk/zrok/zrok/README.md b/sdk/python/sdk/zrok/zrok/README.md new file mode 100644 index 00000000..76b5fe8f --- /dev/null +++ b/sdk/python/sdk/zrok/zrok/README.md @@ -0,0 +1,47 @@ +# Zrok Python SDK + +## Proxy Facility + +The SDK includes a proxy facility that makes it easy to create and manage proxy shares. This is particularly useful when you need to: + +1. Create an HTTP proxy with zrok +2. Optionally reserve the proxy with a unique name for persistence +3. Automatically handle cleanup of non-reserved shares + +### Basic Usage + +```python +from zrok.proxy import ProxyShare +import zrok + +# Load the environment +root = zrok.environment.root.Load() + +# Create a temporary proxy share (will be cleaned up on exit) +proxy = ProxyShare.create(root=root, target="http://my-target-service") + +# Access the proxy's endpoints and token +print(f"Access proxy at: {proxy.endpoints}") +print(f"Share token: {proxy.token}") +``` + +### Creating a Reserved Proxy Share + +To create a proxy share that persists and can be reused: + +```python +# Create/retrieve a reserved proxy share with a unique name +proxy = ProxyShare.create( + root=root, + target="http://my-target-service", + unique_name="my-persistent-proxy" +) +``` + +When a `unique_name` is provided: + +1. If the zrok environment already has a share with that name, it will be reused +2. If no share exists, a new reserved share will be created +3. The share will be automatically cleaned up on exit if no `unique_name` is provided + +When a `unique_name` is not provided, the randomly generated share will be cleaned up on exit. diff --git a/sdk/python/sdk/zrok/zrok/proxy.py b/sdk/python/sdk/zrok/zrok/proxy.py new file mode 100644 index 00000000..65e2306b --- /dev/null +++ b/sdk/python/sdk/zrok/zrok/proxy.py @@ -0,0 +1,178 @@ +""" +Proxy share management functionality for the zrok SDK. +""" + +import atexit +import logging +import urllib.parse +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import requests +from flask import Flask, Response, request +from waitress import serve +from zrok.environment.root import Root +from zrok.model import (PROXY_BACKEND_MODE, PUBLIC_SHARE_MODE, Share, + ShareRequest) +from zrok.overview import Overview +from zrok.share import CreateShare, ReleaseReservedShare + +import zrok + +logger = logging.getLogger(__name__) + +# List of hop-by-hop headers that should not be returned to the viewer +HOP_BY_HOP_HEADERS = { + 'connection', + 'keep-alive', + 'proxy-authenticate', + 'proxy-authorization', + 'te', + 'trailers', + 'transfer-encoding', + 'upgrade' +} + +# The proxy only listens on the zrok socket, the port used to initialize the Waitress server is not actually bound or +# listening +DUMMY_PORT = 18081 + + +@dataclass +class ProxyShare: + """Represents a proxy share with its configuration and state.""" + root: Root + share: Share + target: str + unique_name: Optional[str] = None + _cleanup_registered: bool = False + _app: Optional[Flask] = None + + @classmethod + def create(cls, root: Root, target: str, unique_name: Optional[str] = None) -> 'ProxyShare': + """ + Create a new proxy share, handling reservation and cleanup logic based on unique_name. + + Args: + root: The zrok root environment + target: Target URL or service to proxy to + unique_name: Optional unique name for a reserved share + + Returns: + ProxyShare instance configured with the created share + """ + # First check if we have an existing reserved share with this name + if unique_name: + existing_share = cls._find_existing_share(root, unique_name) + if existing_share: + logger.debug(f"Found existing share with token: {existing_share.Token}") + return cls( + root=root, + share=existing_share, + target=target, + unique_name=unique_name + ) + + # Create new share request + share_request = ShareRequest( + BackendMode=PROXY_BACKEND_MODE, + ShareMode=PUBLIC_SHARE_MODE, + Target="http-proxy", + Frontends=['public'], + Reserved=bool(unique_name) + ) + if unique_name: + share_request.UniqueName = unique_name + + # Create the share + share = CreateShare(root=root, request=share_request) + logger.info(f"Created new proxy share with endpoints: {', '.join(share.FrontendEndpoints)}") + + # Create instance and setup cleanup if needed + instance = cls( + root=root, + share=share, + target=target, + unique_name=unique_name + ) + if not unique_name: + instance.register_cleanup() + return instance + + @staticmethod + def _find_existing_share(root: Root, unique_name: str) -> Optional[Share]: + """Find an existing share with the given unique name.""" + overview = Overview.create(root=root) + for env in overview.environments: + if env.environment.z_id == root.env.ZitiIdentity: + for share in env.shares: + if share.token == unique_name: + return Share(Token=share.token, FrontendEndpoints=[share.frontend_endpoint]) + return None + + def register_cleanup(self): + """Register cleanup handler to release the share on exit.""" + if not self._cleanup_registered: + def cleanup(): + ReleaseReservedShare(root=self.root, shr=self.share) + logger.info(f"Share {self.share.Token} released") + atexit.register(cleanup) + self._cleanup_registered = True + + def _create_app(self) -> Flask: + """Create and configure the Flask app for proxying.""" + app = Flask(__name__) + + @app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS']) + @app.route('/', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS']) + def proxy(path): + # Construct the target URL + url = urllib.parse.urljoin(self.target, path) + + # Forward the request + resp = requests.request( + method=request.method, + url=url, + headers={key: value for (key, value) in request.headers + if key.lower() not in HOP_BY_HOP_HEADERS}, + data=request.get_data(), + cookies=request.cookies, + allow_redirects=False, + stream=True + ) + + # Create the response + excluded_headers = HOP_BY_HOP_HEADERS.union({'host'}) + headers = [(name, value) for (name, value) in resp.raw.headers.items() + if name.lower() not in excluded_headers] + + return Response( + resp.iter_content(chunk_size=10*1024), + status=resp.status_code, + headers=headers + ) + return app + + def run(self): + """Start the proxy server.""" + if self._app is None: + self._app = self._create_app() + + # Create options dictionary for zrok decorator + zrok_opts: Dict[str, Any] = {} + zrok_opts['cfg'] = zrok.decor.Opts(root=self.root, shrToken=self.token, bindPort=DUMMY_PORT) + + @zrok.decor.zrok(opts=zrok_opts) + def run_server(): + serve(self._app, port=DUMMY_PORT) + run_server() + + @property + def endpoints(self) -> List[str]: + """Get the frontend endpoints for this share.""" + return self.share.FrontendEndpoints + + @property + def token(self) -> str: + """Get the share token.""" + return self.share.Token