migrate the proxy example to ProxyShare class so it's a reusable feature of the Py SDK

sdk/python/examples/proxy/proxy.py

@@ -1,153 +1,50 @@
-import argparse
+#!/usr/bin/env python3
-import atexit
-import logging
-import sys
-import urllib.parse
 
-import requests
+"""
-from flask import Flask, Response, request
+Example of using zrok's proxy facility to create an HTTP proxy server.
-from waitress import serve
+
-from zrok.model import ShareRequest, Share
+This example demonstrates how to:
-from zrok.overview import EnvironmentAndResources, Overview
+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
+def main():
-HOP_BY_HOP_HEADERS = {
+    """Main entry point."""
-    '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('/<path:path>', 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__':
     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"Target URL: {args.target_url}")
-    logger.info(f"Logging level: {logger.getEffectiveLevel()}")
 
+    # Load environment and create proxy share
     root = zrok.environment.root.Load()
-    my_env = EnvironmentAndResources(
+    proxy_share = ProxyShare.create(
-        environment=None,
+        root=root,
-        shares=[]
+        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
     
-    if my_env:
+    # Log access information and start the proxy
-        logger.debug(
+    logger.info(f"Access proxy at: {', '.join(proxy_share.endpoints)}")
-            f"Found environment in overview with Ziti identity "
+    proxy_share.run()
-            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:
+if __name__ == '__main__':
-        logger.debug(f"Found existing share with token: {existing_reserved_share.token}")
+    main()
-        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()

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.

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('/<path:path>', 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