add share mode

This commit is contained in:
Kenneth Bingham 2025-01-29 04:06:37 -05:00
parent 41ce0809ab
commit 5c64a73aad
No known key found for this signature in database
GPG Key ID: 31709281860130B6
2 changed files with 23 additions and 11 deletions

View File

@ -31,23 +31,30 @@ def main():
parser.add_argument('-f', '--frontends', nargs='+', help='One or more space-separated frontends to use')
parser.add_argument('-k', '--insecure', action='store_false', dest='verify_ssl', default=True,
help='Skip SSL verification')
parser.add_argument('-s', '--share-mode', default='public', choices=['public', 'private'],
help='Share mode (default: public)')
args = parser.parse_args()
logger.info("=== Starting proxy server ===")
logger.info(f"Target URL: {args.target_url}")
logger.info(f"Share mode: {args.share_mode}")
# Load environment and create proxy share
root = zrok.environment.root.Load()
proxy_share = ProxyShare.create(
root=root,
target=args.target_url,
share_mode=args.share_mode,
unique_name=args.unique_name,
frontends=args.frontends,
verify_ssl=args.verify_ssl
)
# Log access information and start the proxy
logger.info(f"Access proxy at: {', '.join(proxy_share.endpoints)}")
if args.share_mode == "public":
logger.info(f"Access proxy at: {', '.join(proxy_share.endpoints)}")
elif args.share_mode == "private":
logger.info(f"Run a private access frontend: 'zrok access private {proxy_share.token}'")
proxy_share.run()

View File

@ -12,7 +12,7 @@ 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,
from zrok.model import (PROXY_BACKEND_MODE, PUBLIC_SHARE_MODE, PRIVATE_SHARE_MODE, Share,
ShareRequest)
from zrok.overview import Overview
from zrok.share import CreateShare, ReleaseReservedShare
@ -50,7 +50,7 @@ class ProxyShare:
verify_ssl: bool = True
@classmethod
def create(cls, root: Root, target: str, unique_name: Optional[str] = None,
def create(cls, root: Root, target: str, share_mode: str = PUBLIC_SHARE_MODE, unique_name: Optional[str] = None,
frontends: Optional[List[str]] = None, verify_ssl: bool = True) -> 'ProxyShare':
"""
Create a new proxy share, handling reservation and cleanup logic based on unique_name.
@ -79,16 +79,18 @@ class ProxyShare:
)
# Compose the share request
if frontends:
share_frontends = frontends
elif root.cfg and root.cfg.DefaultFrontend:
share_frontends = [root.cfg.DefaultFrontend]
else:
share_frontends = ['public']
share_frontends = []
if share_mode == PUBLIC_SHARE_MODE:
if frontends:
share_frontends = frontends
elif root.cfg and root.cfg.DefaultFrontend:
share_frontends = [root.cfg.DefaultFrontend]
else:
share_frontends = ['public']
share_request = ShareRequest(
BackendMode=PROXY_BACKEND_MODE,
ShareMode=PUBLIC_SHARE_MODE,
ShareMode=share_mode,
Target=target,
Frontends=share_frontends,
Reserved=True
@ -98,7 +100,10 @@ class ProxyShare:
# Create the share
share = CreateShare(root=root, request=share_request)
logger.debug(f"Created new proxy share with endpoints: {', '.join(share.FrontendEndpoints)}")
if share_mode == PUBLIC_SHARE_MODE:
logger.debug(f"Created new proxy share with endpoints: {', '.join(share.FrontendEndpoints)}")
elif share_mode == PRIVATE_SHARE_MODE:
logger.debug(f"Created new private share with token: {share.Token}")
# Create class instance and setup cleanup-at-exit if we reserved a random share token
instance = cls(