|
| 1 | +import asyncio |
| 2 | +import errno |
| 3 | +import json |
| 4 | +import sys |
| 5 | + |
| 6 | +import rich_click as click |
| 7 | + |
| 8 | +import flyte.remote as remote |
| 9 | + |
| 10 | +from . import _common as common |
| 11 | + |
| 12 | +# Headers that must not be forwarded verbatim across a proxy hop. |
| 13 | +_HOP_BY_HOP = { |
| 14 | + "connection", |
| 15 | + "keep-alive", |
| 16 | + "proxy-authenticate", |
| 17 | + "proxy-authorization", |
| 18 | + "te", |
| 19 | + "trailer", |
| 20 | + "transfer-encoding", |
| 21 | + "upgrade", |
| 22 | + "host", |
| 23 | + "content-length", |
| 24 | +} |
| 25 | + |
| 26 | + |
| 27 | +@click.group(name="proxy") |
| 28 | +def proxy(): |
| 29 | + """Proxy a local port into a Flyte App through the authenticated app edge.""" |
| 30 | + |
| 31 | + |
| 32 | +@proxy.command(cls=common.CommandBase) |
| 33 | +@click.argument("name", type=str, required=False) |
| 34 | +@click.option("--url", "url", type=str, default=None, help="Proxy this app URL directly; skip name resolution.") |
| 35 | +@click.option("--port", type=int, default=8600, help="Local port to listen on (0 = pick a free port).") |
| 36 | +@click.option("--address", type=str, default="127.0.0.1", help="Local bind address; non-loopback triggers a warning.") |
| 37 | +@click.option( |
| 38 | + "--emit-mcp-config", |
| 39 | + is_flag=True, |
| 40 | + default=False, |
| 41 | + help="Print a generic HTTP-MCP config block for the local endpoint.", |
| 42 | +) |
| 43 | +@click.option("-v", "--verbose", is_flag=True, default=False, help="Log each proxied request (never the token).") |
| 44 | +@click.pass_obj |
| 45 | +def app( |
| 46 | + cfg: common.CLIConfig, |
| 47 | + name: str | None = None, |
| 48 | + project: str | None = None, |
| 49 | + domain: str | None = None, |
| 50 | + url: str | None = None, |
| 51 | + port: int = 8600, |
| 52 | + address: str = "127.0.0.1", |
| 53 | + emit_mcp_config: bool = False, |
| 54 | + verbose: bool = False, |
| 55 | +): |
| 56 | + """ |
| 57 | + Open an authenticated localhost proxy into a no-auth Flyte App. |
| 58 | +
|
| 59 | + Reuses the same Union auth (with auto-refresh) the CLI uses, injecting a fresh bearer on every |
| 60 | + request, so a local HTTP client — a Grafana/Prometheus MCP, curl, a browser — reaches an |
| 61 | + edge-gated app with no token handling. Think: kubectl port-forward for Flyte Apps. |
| 62 | +
|
| 63 | + Foreground; Ctrl-C to stop. |
| 64 | + """ |
| 65 | + cfg.init(project=project, domain=domain) |
| 66 | + |
| 67 | + if url: |
| 68 | + target = url |
| 69 | + elif name: |
| 70 | + target = remote.App.get(name=name).endpoint |
| 71 | + else: |
| 72 | + raise click.UsageError("Provide an app NAME or --url.") |
| 73 | + target = target.rstrip("/") |
| 74 | + |
| 75 | + label = name or target |
| 76 | + try: |
| 77 | + asyncio.run(_serve(cfg, target, label, address, port, emit_mcp_config, verbose)) |
| 78 | + except KeyboardInterrupt: |
| 79 | + pass |
| 80 | + |
| 81 | + |
| 82 | +def _build_authenticator(cfg: common.CLIConfig): |
| 83 | + import typing |
| 84 | + |
| 85 | + from flyte.remote._client.auth._authenticators.factory import get_async_authenticator |
| 86 | + from flyte.remote._client.auth._client_config import AuthType, RemoteClientConfigStore |
| 87 | + from flyte.remote._client.auth._session import normalize_rpc_endpoint |
| 88 | + |
| 89 | + plat = cfg.config.platform |
| 90 | + if not plat.endpoint: |
| 91 | + raise click.UsageError("No endpoint configured; set one via config or FLYTECTL_CONFIG.") |
| 92 | + insecure = getattr(plat, "insecure", False) |
| 93 | + # Config endpoint is gRPC-style (bare host); the OIDC-metadata client needs an http(s) base. |
| 94 | + endpoint = normalize_rpc_endpoint(plat.endpoint, insecure=insecure) |
| 95 | + auth_type = typing.cast(AuthType, getattr(plat, "auth_mode", None) or "Pkce") |
| 96 | + return get_async_authenticator( |
| 97 | + endpoint=endpoint, |
| 98 | + cfg_store=RemoteClientConfigStore(endpoint), |
| 99 | + auth_type=auth_type, |
| 100 | + insecure_skip_verify=getattr(plat, "insecure_skip_verify", False), |
| 101 | + ca_cert_file_path=getattr(plat, "ca_cert_file_path", None), |
| 102 | + ) |
| 103 | + |
| 104 | + |
| 105 | +def _filter_request_headers(headers) -> dict: |
| 106 | + return {k: v for k, v in headers.items() if k.lower() not in _HOP_BY_HOP and k.lower() != "authorization"} |
| 107 | + |
| 108 | + |
| 109 | +def _filter_response_headers(headers) -> dict: |
| 110 | + return {k: v for k, v in headers.items() if k.lower() not in _HOP_BY_HOP} |
| 111 | + |
| 112 | + |
| 113 | +async def _serve(cfg, target, label, address, port, emit_mcp_config, verbose): |
| 114 | + from aiohttp import ClientSession, ClientTimeout, web |
| 115 | + |
| 116 | + authenticator = _build_authenticator(cfg) |
| 117 | + # Prime credentials (load from keyring / run the auth flow) before we start serving. |
| 118 | + if authenticator.get_credentials() is None: |
| 119 | + await authenticator.refresh_credentials() |
| 120 | + |
| 121 | + # auto_decompress=False -> stream upstream bytes verbatim (Content-Encoding preserved). |
| 122 | + # total=None -> allow long-lived SSE streams (MCP streamable-HTTP, Grafana Live). |
| 123 | + upstream = ClientSession(timeout=ClientTimeout(total=None), auto_decompress=False) |
| 124 | + |
| 125 | + async def handle(request: "web.Request"): |
| 126 | + body = await request.read() |
| 127 | + fwd = _filter_request_headers(request.headers) |
| 128 | + up_url = target + request.raw_path |
| 129 | + |
| 130 | + async def send(refresh: bool): |
| 131 | + if refresh: |
| 132 | + await authenticator.refresh_credentials() |
| 133 | + ah = await authenticator.get_auth_headers() |
| 134 | + hdrs = dict(fwd) |
| 135 | + if ah: |
| 136 | + # Normalize the injected bearer onto the standard `authorization` header. |
| 137 | + # The SDK auth flow emits its bearer under the CP's configured metadata key |
| 138 | + # (typically the gRPC-style `flyte-authorization`), but a Flyte App's HTTP |
| 139 | + # auth edge (union-apps) reads only `authorization` and never sees |
| 140 | + # `flyte-authorization`, so the request is rejected. Scope this to the |
| 141 | + # authenticator's own headers — never rewrite forwarded client headers. |
| 142 | + for k, v in ah.headers.items(): |
| 143 | + hdrs.pop(k, None) |
| 144 | + if isinstance(v, str) and v.startswith("Bearer "): |
| 145 | + hdrs["authorization"] = v |
| 146 | + else: |
| 147 | + hdrs[k] = v |
| 148 | + return await upstream.request( |
| 149 | + request.method, up_url, headers=hdrs, data=body or None, allow_redirects=False |
| 150 | + ) |
| 151 | + |
| 152 | + resp = await send(refresh=False) |
| 153 | + # Stale token: the edge 401s or bounces to /login. Refresh once and retry. |
| 154 | + bounced = resp.status in (302, 307) and "/login" in resp.headers.get("Location", "") |
| 155 | + if resp.status in (401, 403) or bounced: |
| 156 | + resp.release() |
| 157 | + resp = await send(refresh=True) |
| 158 | + |
| 159 | + out = web.StreamResponse(status=resp.status, headers=_filter_response_headers(resp.headers)) |
| 160 | + await out.prepare(request) |
| 161 | + async for chunk in resp.content.iter_any(): |
| 162 | + await out.write(chunk) |
| 163 | + await out.write_eof() |
| 164 | + resp.release() |
| 165 | + if verbose: |
| 166 | + click.echo(f"{request.method} {request.path} -> {resp.status}", err=True) |
| 167 | + return out |
| 168 | + |
| 169 | + server = web.Application() |
| 170 | + server.router.add_route("*", "/{tail:.*}", handle) |
| 171 | + runner = web.AppRunner(server) |
| 172 | + await runner.setup() |
| 173 | + # Bind loopback on BOTH IPv4 and IPv6 so the proxy owns `localhost` fully. |
| 174 | + # macOS resolves `localhost` to ::1 (IPv6) first, and a wildcard listener in |
| 175 | + # another process (e.g. OrbStack/Docker on *:PORT) otherwise shadows an |
| 176 | + # IPv4-only 127.0.0.1 bind: bind() still succeeds, so the proxy looks healthy |
| 177 | + # yet never receives the request (it lands on the other listener, which |
| 178 | + # resets it). Binding ::1 too turns that silent half-bind into a loud |
| 179 | + # EADDRINUSE. A non-loopback --address is bound as-is (single family). |
| 180 | + loopback = address in ("127.0.0.1", "localhost", "::1", "loopback") |
| 181 | + hosts = ["127.0.0.1", "::1"] if loopback else [address] |
| 182 | + |
| 183 | + actual = port |
| 184 | + bound: list[str] = [] |
| 185 | + for host in hosts: |
| 186 | + try: |
| 187 | + await web.TCPSite(runner, host, actual).start() |
| 188 | + except OSError as e: |
| 189 | + if e.errno == errno.EADDRINUSE: |
| 190 | + await runner.cleanup() |
| 191 | + raise click.ClickException( |
| 192 | + f"Port {actual} is already in use (binding {host} failed). Another " |
| 193 | + f"process — often a wildcard binder like OrbStack or Docker — holds it " |
| 194 | + f"and would silently shadow this proxy. Re-run with --port <free-port>." |
| 195 | + ) |
| 196 | + if host == "::1" and e.errno in (errno.EADDRNOTAVAIL, errno.EAFNOSUPPORT): |
| 197 | + click.secho( |
| 198 | + f"note: IPv6 loopback unavailable, binding IPv4 only ({e.strerror}).", |
| 199 | + fg="yellow", |
| 200 | + err=True, |
| 201 | + ) |
| 202 | + continue |
| 203 | + await runner.cleanup() |
| 204 | + raise |
| 205 | + bound.append(host) |
| 206 | + if port == 0 and actual == 0: |
| 207 | + # First site picked a free port; pin the other family to the same one. |
| 208 | + actual = runner.addresses[0][1] |
| 209 | + if not bound: |
| 210 | + await runner.cleanup() |
| 211 | + raise click.ClickException("Failed to bind any loopback address.") |
| 212 | + |
| 213 | + display_host = "127.0.0.1" if loopback else address |
| 214 | + local = f"http://{display_host}:{actual}" |
| 215 | + |
| 216 | + if not loopback: |
| 217 | + click.secho( |
| 218 | + f"WARNING: binding {address} exposes your Union identity to anything that can reach it; prefer 127.0.0.1.", |
| 219 | + fg="yellow", |
| 220 | + err=True, |
| 221 | + ) |
| 222 | + identity = _identity(authenticator) |
| 223 | + click.echo(f"Proxying {target} -> {local}", err=True) |
| 224 | + click.echo(f" authenticating as: {identity}", err=True) |
| 225 | + click.echo(" Ctrl-C to stop.", err=True) |
| 226 | + if emit_mcp_config: |
| 227 | + _emit_mcp_config(label, local) |
| 228 | + |
| 229 | + try: |
| 230 | + await asyncio.Event().wait() # serve until interrupted (Ctrl-C) |
| 231 | + finally: |
| 232 | + await upstream.close() |
| 233 | + await runner.cleanup() |
| 234 | + |
| 235 | + |
| 236 | +def _identity(authenticator) -> str: |
| 237 | + creds = authenticator.get_credentials() |
| 238 | + if not creds or not creds.access_token: |
| 239 | + return "<unknown>" |
| 240 | + # Best-effort: decode the JWT payload's sub/email without verifying (display only). |
| 241 | + try: |
| 242 | + import base64 |
| 243 | + |
| 244 | + payload = creds.access_token.split(".")[1] |
| 245 | + payload += "=" * (-len(payload) % 4) |
| 246 | + claims = json.loads(base64.urlsafe_b64decode(payload)) |
| 247 | + return claims.get("email") or claims.get("sub") or "<token>" |
| 248 | + except Exception: |
| 249 | + return "<token>" |
| 250 | + |
| 251 | + |
| 252 | +def _emit_mcp_config(name: str, local: str): |
| 253 | + block = {"mcpServers": {name: {"type": "http", "url": local}}} |
| 254 | + click.echo("", err=True) |
| 255 | + click.echo("# MCP client config (generic HTTP transport) — point your client at the local proxy:", err=True) |
| 256 | + sys.stdout.write(json.dumps(block, indent=2) + "\n") |
| 257 | + sys.stdout.flush() |
0 commit comments