stroblme
f569f325a7
The connector now serves /mcp, its OAuth documents and /hooks/ besides /api/v1/, refuses any dot segment (httpx resolves them, which carried /api/v1/../../mcp onto the MCP mount), addresses calls to the hub's host, and bridges served sockets by calling the app with a websocket scope. A hook without a secret answers the tunnel with the 404 (or 1008) of a wrong secret. Requests through the hub get the hub's issuer, resource and consent page in the OAuth metadata and the 401, and the MCP host check admits the hub while enrolled.
851 lines
35 KiB
Python
851 lines
35 KiB
Python
"""The link this instance opens to its portal.
|
|
|
|
Same shape as the remote-worker agent, with the direction of the favour
|
|
reversed: there, a GPU box dials the engine so the engine can give it work;
|
|
here, the engine dials the portal so the portal can pass on what a browser
|
|
asked for. Both exist because the interesting machine is the one nobody can
|
|
route to.
|
|
|
|
A proxied request is dispatched straight into this process's own ASGI app
|
|
rather than over localhost HTTP. That matters beyond saving a hop: HTTP
|
|
trigger nodes install themselves into ``app.routes`` while flows run, and an
|
|
in-process transport sees that table live, with no port, TLS or self-address
|
|
to discover.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import contextlib
|
|
import logging
|
|
import random
|
|
import time
|
|
import uuid
|
|
from collections.abc import MutableMapping
|
|
from dataclasses import replace
|
|
from datetime import timedelta
|
|
from typing import Any
|
|
|
|
import httpx
|
|
import jwt
|
|
from fastapi import FastAPI
|
|
|
|
import fluksio
|
|
from fluksio.cloud import config as cloud_config
|
|
from fluksio.cloud.enroll import is_loopback
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
PROTOCOL = 1
|
|
#: The hub closes a credential it will never accept again with this code (see
|
|
#: the hub's ``tunnel/protocol.py``). The link then stands down until enrolled
|
|
#: again or restarted, rather than dialling a deleted instance forever.
|
|
REVOKED_CLOSE = 4410
|
|
HEARTBEAT_S = 20.0
|
|
STATUS_S = 60.0
|
|
BASE_BACKOFF_S = 1.0
|
|
MAX_BACKOFF_S = 30.0
|
|
#: How long a silent socket is given before it is treated as gone. A laptop
|
|
#: that suspended, or a wifi that changed underneath us, leaves a connection
|
|
#: the operating system still believes in — nothing arrives and nothing fails.
|
|
#: Only an unanswered ping tells us, so this is the worst case for noticing.
|
|
PING_INTERVAL_S = 20.0
|
|
PING_TIMEOUT_S = 20.0
|
|
CHUNK_BYTES = 256 * 1024
|
|
MAX_WS_MESSAGE = 1024 * 1024
|
|
#: What the tunnel serves: the versioned API, which the hub opens only to a
|
|
#: portal session; the MCP endpoint and the OAuth documents leading an agent
|
|
#: to it, which take the agent's own token; and the flows' hooks, of which
|
|
#: only those with a secret answer a request from the tunnel. Everything else,
|
|
#: the API docs included, stays on this machine.
|
|
ALLOWED_PREFIXES = ("/api/v1/", "/hooks/")
|
|
ALLOWED_PATHS = frozenset(
|
|
{
|
|
"/mcp",
|
|
"/.well-known/oauth-protected-resource/mcp",
|
|
"/.well-known/oauth-authorization-server",
|
|
"/.well-known/jwks.json",
|
|
}
|
|
)
|
|
#: Proxied calls, and bridged streams, this instance will run at once.
|
|
#: Neither dict was bounded, and an id the hub reused silently dropped the
|
|
#: reference to a task that was still running.
|
|
MAX_IN_FLIGHT = 256
|
|
#: How often to look for a config that appeared while the engine was up.
|
|
#: Enrolment is something a person does and then waits on, so this is the
|
|
#: delay they sit through — short enough not to be worth a restart.
|
|
ENROL_POLL_S = 3.0
|
|
|
|
|
|
def _identity(config: cloud_config.CloudConfig | None) -> tuple[str, str] | None:
|
|
"""Which enrolment a link is running on, so a replaced one is noticed."""
|
|
return (config.instance_id, config.token) if config is not None else None
|
|
|
|
|
|
def start(app: FastAPI) -> None:
|
|
"""Dial the portal, replacing any link already up."""
|
|
existing = getattr(app.state, "cloud_task", None)
|
|
if existing is not None:
|
|
existing.cancel()
|
|
connector = CloudConnector(app)
|
|
app.state.cloud_connector = connector
|
|
app.state.cloud_identity = _identity(cloud_config.load())
|
|
app.state.cloud_task = asyncio.create_task(
|
|
connector.serve_forever(), name="cloud-connector"
|
|
)
|
|
|
|
|
|
async def watch_enrolment(app: FastAPI) -> None:
|
|
"""Notice an enrolment that happened outside this process.
|
|
|
|
``fluksio enroll`` writes the config against the database directly, with no
|
|
idea whether an engine is running — so without this, pairing an instance
|
|
that is already serving would take a restart to come into effect. Enrolling
|
|
through the API starts the link itself and this sees a task already there.
|
|
"""
|
|
while True:
|
|
await asyncio.sleep(ENROL_POLL_S)
|
|
task = getattr(app.state, "cloud_task", None)
|
|
# `load()`, not `exists()`: a file that cannot be read is not an
|
|
# enrolment, and the connector would give up on it the instant it
|
|
# started — which, started from here, is a restart every few seconds.
|
|
# A config that is fine but unreachable keeps its task, and the
|
|
# retrying belongs to the connector rather than to this.
|
|
config = cloud_config.load()
|
|
if task is not None and task.done():
|
|
connector = getattr(app.state, "cloud_connector", None)
|
|
if (
|
|
connector is not None
|
|
and connector.revoked
|
|
and _identity(config) == getattr(app.state, "cloud_identity", None)
|
|
):
|
|
# Stood down on the portal's word, and kept for its status.
|
|
# Only a new enrolment brings the link back.
|
|
continue
|
|
# It returns of its own accord when the config goes away, which is
|
|
# what Disconnect, here or on the portal, does.
|
|
app.state.cloud_task = None
|
|
app.state.cloud_connector = None
|
|
task = None
|
|
if task is not None and _identity(config) != getattr(
|
|
app.state, "cloud_identity", None
|
|
):
|
|
# Enrolled again, at this portal or another one. `fluksio enroll`
|
|
# is its own process and cannot cancel this task, so the link would
|
|
# otherwise stay up on the credential that was replaced.
|
|
logger.info("Enrolment replaced while running; redialling")
|
|
task.cancel()
|
|
task = None
|
|
if task is None and config is not None:
|
|
logger.info("Enrolled while running; dialling the portal")
|
|
start(app)
|
|
|
|
|
|
class CloudConnector:
|
|
"""Holds one websocket to the portal open, and answers over it."""
|
|
|
|
def __init__(self, app: FastAPI) -> None:
|
|
self._app = app
|
|
self._calls: dict[str, asyncio.Task[None]] = {}
|
|
self._streams: dict[str, asyncio.Task[None]] = {}
|
|
self._connected = False
|
|
self._connected_since: float | None = None
|
|
self._last_error: str | None = None
|
|
self._revoked = False
|
|
#: The portal account that owns this instance, as the last welcome said.
|
|
self._owner: str | None = None
|
|
#: Where the portal opens this instance, as the last welcome said. An
|
|
#: older hub does not say.
|
|
self.app_url: str | None = None
|
|
#: What the hub sent each bridged hook socket, waiting to be received.
|
|
self._inboxes: dict[str, asyncio.Queue[dict[str, Any]]] = {}
|
|
|
|
@property
|
|
def revoked(self) -> bool:
|
|
return self._revoked
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Status, for the local settings screen
|
|
# -------------------------------------------------------------------------
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
config = cloud_config.load()
|
|
return {
|
|
"enrolled": config is not None,
|
|
"connected": self._connected,
|
|
"portal_url": config.portal_url if config else None,
|
|
# Where a browser reaches the portal, which is not always where
|
|
# this process does: enrolment may have named a container.
|
|
"issuer": config.issuer if config else None,
|
|
"portal_account": config.portal_account if config else None,
|
|
"instance_id": config.instance_id if config else None,
|
|
"last_error": self._last_error,
|
|
"connected_since": self._connected_since,
|
|
"revoked": self._revoked,
|
|
}
|
|
|
|
# -------------------------------------------------------------------------
|
|
# The connection
|
|
# -------------------------------------------------------------------------
|
|
|
|
async def serve_forever(self) -> None:
|
|
"""Hold the link up, and put it back up when it goes down.
|
|
|
|
Every ending is the same ending here: a portal that closed the socket,
|
|
a network that changed under it, a laptop that woke up somewhere else.
|
|
The link is dialled again in all of them — the only question is how
|
|
soon, and that turns on whether this attempt got as far as being
|
|
attached. One that did and then dropped is a network event, so it
|
|
retries at once; one that never stood up is met with a longer wait
|
|
each time, because whatever is refusing is unlikely to stop within a
|
|
second.
|
|
"""
|
|
delay = BASE_BACKOFF_S
|
|
while True:
|
|
config = cloud_config.load()
|
|
if config is None:
|
|
# Disconnected locally while we were running.
|
|
return
|
|
attached = False
|
|
try:
|
|
attached = await self._session(config)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
self._last_error = str(exc)
|
|
if _is_revoked(exc):
|
|
# Dialling again would be refused the same way forever.
|
|
# The config stays, so enrolling again replaces it.
|
|
self._revoked = True
|
|
logger.warning("The portal revoked this instance; standing down")
|
|
return
|
|
logger.warning("Portal link down: %s", exc)
|
|
finally:
|
|
self._connected = False
|
|
self._connected_since = None
|
|
|
|
delay = BASE_BACKOFF_S if attached else min(MAX_BACKOFF_S, delay * 2)
|
|
# Jittered, so a portal coming back up is not met by every
|
|
# instance it serves in the same instant.
|
|
wait = delay * (0.75 + random.random() * 0.5)
|
|
logger.info("Reconnecting to the portal in %.0fs", wait)
|
|
await asyncio.sleep(wait)
|
|
|
|
async def _session(self, config: cloud_config.CloudConfig) -> bool:
|
|
"""One connection, from dial to close. True if it ever attached."""
|
|
import websockets
|
|
|
|
attached = False
|
|
# A header rather than `?token=`, since a URL ends up in access logs.
|
|
# A hub from before the header reads only the query: upgrade the hub
|
|
# before its instances.
|
|
async with websockets.connect(
|
|
config.ws_url,
|
|
additional_headers={"Authorization": f"Bearer {config.token}"},
|
|
max_size=MAX_WS_MESSAGE,
|
|
ping_interval=PING_INTERVAL_S,
|
|
ping_timeout=PING_TIMEOUT_S,
|
|
) as socket:
|
|
await socket.send(
|
|
_dump({"op": "hello", "protocol": PROTOCOL, "app_version": _version()})
|
|
)
|
|
welcome = _load(await socket.recv())
|
|
if welcome.get("op") != "welcome":
|
|
raise RuntimeError(
|
|
str(welcome.get("reason") or "refused by the portal")
|
|
)
|
|
|
|
self._owner = welcome.get("owner")
|
|
self.app_url = welcome.get("app_url")
|
|
self._adopt_owner(config, self._owner)
|
|
self._adopt_jwks(welcome.get("jwks"))
|
|
|
|
attached = True
|
|
self._connected = True
|
|
self._connected_since = time.time()
|
|
self._last_error = None
|
|
logger.info("Attached to the portal as instance %s", config.instance_id)
|
|
|
|
keepalive = asyncio.create_task(self._keepalive(socket, config))
|
|
try:
|
|
async for raw in socket:
|
|
frame = _load(raw)
|
|
op = frame.get("op")
|
|
if op == "req":
|
|
self._start_call(socket, frame)
|
|
elif op == "cancel":
|
|
self._cancel(str(frame.get("id") or ""))
|
|
elif op == "ws_open":
|
|
self._start_stream(socket, frame)
|
|
elif op == "ws_msg":
|
|
inbox = self._inboxes.get(str(frame.get("id") or ""))
|
|
if inbox is not None:
|
|
inbox.put_nowait(
|
|
{
|
|
"type": "websocket.receive",
|
|
"text": str(frame.get("text") or ""),
|
|
}
|
|
)
|
|
elif op == "ws_close":
|
|
self._cancel_stream(str(frame.get("id") or ""))
|
|
elif op == "member_removed":
|
|
await asyncio.to_thread(
|
|
self._forget_member,
|
|
str(frame.get("user") or ""),
|
|
self._owner,
|
|
)
|
|
finally:
|
|
keepalive.cancel()
|
|
# Awaited, not just cancelled: a keepalive that died of its own
|
|
# accord holds the reason the link went, and an un-awaited task
|
|
# takes it to the garbage collector instead of the log.
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
await keepalive
|
|
self._connected = False
|
|
self._connected_since = None
|
|
for task in list(self._calls.values()) + list(self._streams.values()):
|
|
task.cancel()
|
|
self._calls.clear()
|
|
self._streams.clear()
|
|
return attached
|
|
|
|
@staticmethod
|
|
def _adopt_owner(config: cloud_config.CloudConfig, owner: Any) -> None:
|
|
"""Map the enrolling account to the portal account that owns us.
|
|
|
|
Enrolment does this itself. This is for the enrolments that predate
|
|
per-user mapping: without it their owner would be refused here until
|
|
somebody enrolled the machine again — which, for a machine reached only
|
|
through the portal, means standing in front of it. Runs on every attach
|
|
because it is a no-op once the mapping is there.
|
|
|
|
Never moves a mapping somebody else holds: that would be this code
|
|
guessing at something enrolment was told.
|
|
"""
|
|
from sqlmodel import Session, select
|
|
|
|
from fluksio.core.db import engine
|
|
from fluksio.models import User
|
|
|
|
if not owner:
|
|
return
|
|
owner_id = str(owner)
|
|
try:
|
|
with Session(engine) as session:
|
|
if session.exec(
|
|
select(User).where(User.portal_sub == owner_id)
|
|
).first():
|
|
return
|
|
user = CloudConnector._local_account(session, config)
|
|
if user is None or user.portal_sub:
|
|
return
|
|
user.portal_sub = owner_id
|
|
session.add(user)
|
|
session.commit()
|
|
logger.info("Mapped %s to the portal account that owns us", user.email)
|
|
except Exception:
|
|
# A mapping that could not be written is not a reason to drop the
|
|
# link: everything else this connection does still works.
|
|
logger.exception("Could not map the enrolling account to the portal owner")
|
|
|
|
@staticmethod
|
|
def _adopt_jwks(jwks: Any) -> None:
|
|
"""Replace the pinned signing keys with the set the hub sends on attach.
|
|
|
|
Taken only from a link that is encrypted or on this machine, the rule
|
|
enrolment applies to the portal address. Over cleartext anything on the
|
|
path could hand this instance a key of its own.
|
|
"""
|
|
config = cloud_config.load()
|
|
if config is None or not isinstance(jwks, dict) or jwks == config.jwks:
|
|
return
|
|
if not config.ws_url.startswith("wss://") and not is_loopback(config.ws_url):
|
|
return
|
|
try:
|
|
jwt.PyJWKSet.from_dict(jwks)
|
|
except jwt.PyJWTError as exc:
|
|
logger.warning("Ignoring unusable portal keys: %s", exc)
|
|
return
|
|
cloud_config.save(replace(config, jwks=jwks))
|
|
logger.info("Took the portal's current signing keys")
|
|
|
|
@staticmethod
|
|
def _forget_member(portal_sub: str, owner: str | None) -> None:
|
|
"""Delete the local user of somebody the portal no longer admits.
|
|
|
|
Never the owner, and never a superuser: the portal only ever creates
|
|
ordinary accounts here, so those are not its to take away.
|
|
"""
|
|
from sqlmodel import Session, select
|
|
|
|
from fluksio.core.db import engine
|
|
from fluksio.models import User
|
|
|
|
if not portal_sub or portal_sub == owner:
|
|
return
|
|
try:
|
|
with Session(engine) as session:
|
|
user = session.exec(
|
|
select(User).where(User.portal_sub == portal_sub)
|
|
).first()
|
|
if user is None or user.is_superuser:
|
|
return
|
|
session.delete(user)
|
|
session.commit()
|
|
logger.info(
|
|
"Deleted %s, removed from this instance at the portal", user.email
|
|
)
|
|
except Exception:
|
|
logger.exception("Could not delete the user the portal removed")
|
|
|
|
@staticmethod
|
|
def _local_account(session: Any, config: cloud_config.CloudConfig) -> Any:
|
|
"""The account this instance acts as, or the one that now stands for it.
|
|
|
|
``local_user_id`` names whoever redeemed the claim code. A database
|
|
restored from a backup, or rebuilt under an enrolment that outlived it,
|
|
has the same operator behind a different row — and then this resolves to
|
|
nobody, which would refuse the owner on a machine they may only be able
|
|
to reach through the portal. That is the lockout the enrolment handshake
|
|
exists to prevent, so it is closed here too: one superuser is not a guess,
|
|
it is the account enrolment would have used. Several is a guess, and this
|
|
code does not make it.
|
|
|
|
The new id is written back, because the same field is what a screen
|
|
paired through the portal borrows — a panel would otherwise be refused
|
|
for the same reason and with no way to say so.
|
|
"""
|
|
from fluksio.models import User
|
|
|
|
try:
|
|
user = session.get(User, uuid.UUID(config.local_user_id))
|
|
except ValueError:
|
|
user = None
|
|
if user is not None:
|
|
return user
|
|
|
|
from sqlmodel import select
|
|
|
|
superusers = list(
|
|
session.exec(select(User).where(User.is_superuser == True)) # noqa: E712
|
|
)
|
|
if len(superusers) != 1:
|
|
logger.warning(
|
|
"The account this instance enrolled as is gone, and there "
|
|
"%s to adopt unambiguously — enrol again from Settings.",
|
|
"is no superuser" if not superusers else "are several superusers",
|
|
)
|
|
return None
|
|
adopted = superusers[0]
|
|
cloud_config.save(replace(config, local_user_id=str(adopted.id)))
|
|
logger.warning(
|
|
"The account this instance enrolled as is gone; acting as %s instead",
|
|
adopted.email,
|
|
)
|
|
return adopted
|
|
|
|
async def _keepalive(self, socket: Any, config: cloud_config.CloudConfig) -> None:
|
|
"""Heartbeats, plus a health snapshot the portal can show while offline."""
|
|
last_status = 0.0
|
|
while True:
|
|
now = time.time()
|
|
if now - last_status >= STATUS_S or last_status == 0.0:
|
|
summary = await self._health_summary(config)
|
|
await socket.send(
|
|
_dump(
|
|
{"op": "status", "summary": summary, "app_version": _version()}
|
|
)
|
|
)
|
|
last_status = now
|
|
else:
|
|
await socket.send(_dump({"op": "heartbeat"}))
|
|
await asyncio.sleep(HEARTBEAT_S)
|
|
|
|
async def _health_summary(
|
|
self, config: cloud_config.CloudConfig
|
|
) -> dict[str, Any] | None:
|
|
"""This instance's own health, as the dashboard reads it.
|
|
|
|
Fetched through the same in-process transport as everything else, with
|
|
a short-lived local token: the observability API is authenticated, and
|
|
this process is entitled to mint one for the enrolling user.
|
|
|
|
``failures_24h`` is added here rather than by ``/summary``: it is the
|
|
portal's tile and nothing on this instance reads it, and summing
|
|
the same rollups over the same window the app's own Home tile sums
|
|
keeps the two from drifting apart the way they already did once.
|
|
"""
|
|
from fluksio.core import security
|
|
|
|
token = security.create_access_token(config.local_user_id, timedelta(minutes=5))
|
|
headers = {"Authorization": f"Bearer {token}"}
|
|
try:
|
|
async with self._client() as client:
|
|
response = await client.get(
|
|
"/api/v1/observability/summary", headers=headers
|
|
)
|
|
if response.status_code != 200:
|
|
return None
|
|
summary: dict[str, Any] = response.json()
|
|
rollups = await client.get(
|
|
"/api/v1/observability/flows",
|
|
params={"hours": 24},
|
|
headers=headers,
|
|
)
|
|
if rollups.status_code == 200:
|
|
summary["failures_24h"] = sum(row["errors"] for row in rollups.json())
|
|
return summary
|
|
except Exception:
|
|
logger.debug("Could not collect a health summary for the portal")
|
|
return None
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Proxied requests
|
|
# -------------------------------------------------------------------------
|
|
|
|
def _client(self) -> httpx.AsyncClient:
|
|
return httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=self._app),
|
|
base_url="http://instance.local",
|
|
timeout=httpx.Timeout(300.0, connect=5.0),
|
|
)
|
|
|
|
def _start_call(self, socket: Any, frame: dict[str, Any]) -> None:
|
|
call_id = str(frame.get("id") or "")
|
|
if not call_id:
|
|
return
|
|
if len(self._calls) >= MAX_IN_FLIGHT:
|
|
logger.warning(
|
|
"Refusing proxied call %s: %d already in flight",
|
|
call_id,
|
|
len(self._calls),
|
|
)
|
|
return
|
|
# An id the hub reuses would otherwise drop the reference to a task
|
|
# still running, leaving it with nothing able to cancel it.
|
|
self._cancel(call_id)
|
|
task = asyncio.create_task(self._serve(socket, frame))
|
|
self._calls[call_id] = task
|
|
task.add_done_callback(lambda _t: self._calls.pop(call_id, None))
|
|
|
|
async def _serve(self, socket: Any, frame: dict[str, Any]) -> None:
|
|
call_id = str(frame["id"])
|
|
path = str(frame.get("path") or "")
|
|
try:
|
|
if not _allowed(path):
|
|
# Defence in depth: the hub forwards nothing else, and an
|
|
# instance that trusted it to would be an SSRF gadget aimed at
|
|
# its own process.
|
|
await socket.send(
|
|
_dump({"op": "err", "id": call_id, "message": "path not allowed"})
|
|
)
|
|
return
|
|
body = frame.get("body")
|
|
content = base64.b64decode(body) if body else None
|
|
query = str(frame.get("query") or "")
|
|
|
|
headers = dict(frame.get("headers") or {})
|
|
# Set here rather than trusted from the frame: a browser can send
|
|
# any header it likes through the proxy, and this one decides where
|
|
# a pairing device's credential is minted. Arriving on this socket
|
|
# is the only thing that makes it true, and the value is this
|
|
# process's own so nothing off the network can imitate it.
|
|
headers[cloud_config.VIA_HEADER] = cloud_config.VIA_PORTAL
|
|
# Addressed as the caller addressed it, which is the hub. The MCP
|
|
# endpoint checks the Host against the names it answers to.
|
|
headers["host"] = _hub_host()
|
|
|
|
async with self._client() as client:
|
|
request = client.build_request(
|
|
str(frame.get("method") or "GET"),
|
|
f"{path}?{query}" if query else path,
|
|
headers=headers,
|
|
content=content,
|
|
)
|
|
response = await client.send(request, stream=True)
|
|
try:
|
|
await socket.send(
|
|
_dump(
|
|
{
|
|
"op": "res",
|
|
"id": call_id,
|
|
"status": response.status_code,
|
|
"headers": dict(response.headers),
|
|
}
|
|
)
|
|
)
|
|
async for chunk in response.aiter_bytes(CHUNK_BYTES):
|
|
await socket.send(
|
|
_dump(
|
|
{
|
|
"op": "res_body",
|
|
"id": call_id,
|
|
"body": base64.b64encode(chunk).decode("ascii"),
|
|
"done": False,
|
|
}
|
|
)
|
|
)
|
|
await socket.send(
|
|
_dump({"op": "res_body", "id": call_id, "done": True})
|
|
)
|
|
finally:
|
|
await response.aclose()
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.warning("Proxied %s failed: %s", path, exc)
|
|
with contextlib.suppress(Exception):
|
|
await socket.send(
|
|
_dump({"op": "err", "id": call_id, "message": str(exc)})
|
|
)
|
|
|
|
def _cancel(self, call_id: str) -> None:
|
|
task = self._calls.pop(call_id, None)
|
|
if task is not None:
|
|
task.cancel()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Bridged live stream
|
|
# -------------------------------------------------------------------------
|
|
|
|
def _start_stream(self, socket: Any, frame: dict[str, Any]) -> None:
|
|
stream_id = str(frame.get("id") or "")
|
|
if not stream_id:
|
|
return
|
|
if len(self._streams) >= MAX_IN_FLIGHT:
|
|
logger.warning(
|
|
"Refusing stream %s: %d already open", stream_id, len(self._streams)
|
|
)
|
|
return
|
|
existing = self._streams.pop(stream_id, None)
|
|
if existing is not None:
|
|
existing.cancel()
|
|
task = asyncio.create_task(self._stream(socket, frame))
|
|
self._streams[stream_id] = task
|
|
task.add_done_callback(lambda _t: self._streams.pop(stream_id, None))
|
|
|
|
async def _stream(self, socket: Any, frame: dict[str, Any]) -> None:
|
|
"""Feed the engine's live events down the tunnel.
|
|
|
|
The flows websocket is subscribed to directly rather than dialled:
|
|
httpx cannot speak websocket, and everything that endpoint does — check
|
|
the token, send a snapshot, forward the bus — is a few lines against
|
|
objects this process already holds. The snapshot is built by the route
|
|
itself, so the two cannot drift apart behind the portal's back.
|
|
"""
|
|
from sqlmodel import Session
|
|
|
|
from fluksio.api.deps import user_from_token
|
|
from fluksio.api.routes.flows import (
|
|
event_for_panel,
|
|
panel_scope,
|
|
snapshot_payload,
|
|
)
|
|
from fluksio.api.socket_auth import redeem_ticket
|
|
from fluksio.core.db import engine
|
|
from fluksio.flow.events import event_bus
|
|
|
|
stream_id = str(frame["id"])
|
|
path = str(frame.get("path") or "")
|
|
query = str(frame.get("query") or "")
|
|
if _allowed(path) and path.startswith("/hooks/"):
|
|
await self._hook_stream(socket, stream_id, path, query)
|
|
return
|
|
# The browser's ticket when it brought one; the hub restates the
|
|
# session it checked as `token` either way, which is the tunnel's own
|
|
# channel and stays when the public `?token=` goes.
|
|
ticket = _query_value(query, "ticket")
|
|
token = redeem_ticket(ticket) if ticket else _query_value(query, "token")
|
|
|
|
if not path.startswith("/api/v1/") or not path.endswith("/flows/ws"):
|
|
await socket.send(_dump({"op": "ws_close", "id": stream_id, "code": 1008}))
|
|
return
|
|
with Session(engine) as session:
|
|
if user_from_token(session, token) is None:
|
|
await socket.send(
|
|
_dump({"op": "ws_close", "id": stream_id, "code": 1008})
|
|
)
|
|
return
|
|
only = panel_scope(token, self._app)
|
|
|
|
await socket.send(_dump({"op": "ws_open_ok", "id": stream_id}))
|
|
|
|
controller = getattr(self._app.state, "flow_controller", None)
|
|
if controller is not None:
|
|
await socket.send(
|
|
_dump(
|
|
{
|
|
"op": "ws_msg",
|
|
"id": stream_id,
|
|
"text": _dump(snapshot_payload(controller, only)),
|
|
}
|
|
)
|
|
)
|
|
|
|
try:
|
|
async with event_bus.subscribe() as queue:
|
|
while True:
|
|
event = await queue.get()
|
|
if only is not None and event.get("type") == "dashboard_changed":
|
|
# What this screen shows may have just changed under
|
|
# it. Rescope, then resend the snapshot so a dashboard
|
|
# it has only now been assigned draws values instead of
|
|
# blanks. ``or set()`` is the point: a panel that was
|
|
# deleted scopes to nothing, and None would widen this
|
|
# socket to everything on the bus.
|
|
only = panel_scope(token, self._app) or set()
|
|
if controller is not None:
|
|
catch_up = snapshot_payload(controller, only)
|
|
await socket.send(
|
|
_dump(
|
|
{
|
|
"op": "ws_msg",
|
|
"id": stream_id,
|
|
"text": _dump(catch_up),
|
|
}
|
|
)
|
|
)
|
|
elif only is not None and not event_for_panel(event, only):
|
|
continue
|
|
await socket.send(
|
|
_dump({"op": "ws_msg", "id": stream_id, "text": _dump(event)})
|
|
)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.debug("Live stream %s ended: %s", stream_id, exc)
|
|
with contextlib.suppress(Exception):
|
|
await socket.send(
|
|
_dump({"op": "ws_close", "id": stream_id, "code": 1011})
|
|
)
|
|
|
|
async def _hook_stream(
|
|
self, socket: Any, stream_id: str, path: str, query: str
|
|
) -> None:
|
|
"""Serve a flow's socket over the tunnel, as uvicorn would serve it here.
|
|
|
|
The app is called with a websocket scope of its own, so the route a hook
|
|
installed answers exactly as it does a local client, secret included.
|
|
"""
|
|
inbox: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
|
|
inbox.put_nowait({"type": "websocket.connect"})
|
|
self._inboxes[stream_id] = inbox
|
|
state = {"accepted": False, "closed": False}
|
|
|
|
async def send(message: MutableMapping[str, Any]) -> None:
|
|
kind = message["type"]
|
|
if kind == "websocket.accept":
|
|
state["accepted"] = True
|
|
await socket.send(_dump({"op": "ws_open_ok", "id": stream_id}))
|
|
elif kind == "websocket.send":
|
|
text = message.get("text")
|
|
if text is None:
|
|
text = (message.get("bytes") or b"").decode(errors="replace")
|
|
await socket.send(
|
|
_dump({"op": "ws_msg", "id": stream_id, "text": text})
|
|
)
|
|
elif kind == "websocket.close":
|
|
state["closed"] = True
|
|
# Refused before it was accepted is a refusal whatever the code,
|
|
# so no route, a wrong secret and no secret all look alike.
|
|
code = message.get("code", 1000) if state["accepted"] else 1008
|
|
await socket.send(
|
|
_dump({"op": "ws_close", "id": stream_id, "code": code})
|
|
)
|
|
|
|
scope = {
|
|
"type": "websocket",
|
|
"asgi": {"version": "3.0"},
|
|
"http_version": "1.1",
|
|
"scheme": "ws",
|
|
"path": path,
|
|
"raw_path": path.encode(),
|
|
"root_path": "",
|
|
"query_string": query.encode(),
|
|
"headers": [
|
|
(b"host", _hub_host().encode()),
|
|
(cloud_config.VIA_HEADER.encode(), cloud_config.VIA_PORTAL.encode()),
|
|
],
|
|
"client": ("127.0.0.1", 0),
|
|
"server": None,
|
|
"subprotocols": [],
|
|
}
|
|
code = 1000
|
|
try:
|
|
await self._app(scope, inbox.get, send)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.debug("Hook stream %s ended: %s", stream_id, exc)
|
|
code = 1011
|
|
finally:
|
|
self._inboxes.pop(stream_id, None)
|
|
if not state["closed"]:
|
|
with contextlib.suppress(Exception):
|
|
await socket.send(
|
|
_dump({"op": "ws_close", "id": stream_id, "code": code})
|
|
)
|
|
|
|
def _cancel_stream(self, stream_id: str) -> None:
|
|
task = self._streams.pop(stream_id, None)
|
|
if task is not None:
|
|
task.cancel()
|
|
|
|
|
|
def _allowed(path: str) -> bool:
|
|
"""Whether the tunnel may serve this path.
|
|
|
|
A dot segment is refused, literal or encoded: httpx resolves a literal one
|
|
before dispatching, so ``/hooks/../api/v1/…`` would otherwise reach the API
|
|
as a path the hub forwarded without a session.
|
|
"""
|
|
from urllib.parse import unquote
|
|
|
|
if {".", ".."} & set(unquote(path).split("/")):
|
|
return False
|
|
return path in ALLOWED_PATHS or path.startswith(ALLOWED_PREFIXES)
|
|
|
|
|
|
def _hub_host() -> str:
|
|
"""The host a caller addressed through the tunnel, the hub's own."""
|
|
from urllib.parse import urlsplit
|
|
|
|
config = cloud_config.load()
|
|
return urlsplit(config.issuer).netloc if config is not None else "instance.local"
|
|
|
|
|
|
def _is_revoked(exc: Exception) -> bool:
|
|
"""Whether the hub closed the link as revoked, at the handshake or later."""
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
return (
|
|
isinstance(exc, ConnectionClosed)
|
|
and exc.rcvd is not None
|
|
and exc.rcvd.code == REVOKED_CLOSE
|
|
)
|
|
|
|
|
|
def _dump(payload: Any) -> str:
|
|
import json
|
|
|
|
return json.dumps(payload, default=str)
|
|
|
|
|
|
def _load(raw: Any) -> dict[str, Any]:
|
|
import json
|
|
|
|
data: dict[str, Any] = json.loads(raw)
|
|
return data
|
|
|
|
|
|
def _query_value(query: str, key: str) -> str:
|
|
from urllib.parse import parse_qs
|
|
|
|
values = parse_qs(query).get(key) or [""]
|
|
return values[0]
|
|
|
|
|
|
def _version() -> str:
|
|
return fluksio.__version__
|