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

This commit is contained in:
Kenneth Bingham 2025-01-28 18:33:31 -05:00
parent 9097643617
commit fe056a2ccc
No known key found for this signature in database
GPG Key ID: 31709281860130B6
3 changed files with 253 additions and 131 deletions

View File

@ -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('/<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__':
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()

View File

@ -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.

View File

@ -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