Coverage for /usr/local/lib/python3.10/site-packages/opal_server-0.0.0-py3.10.egg/opal_server/pubsub.py: 46%
153 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 11:54 +0000
1import asyncio
2import random
3import time
4from contextlib import contextmanager
5from contextvars import ContextVar
6from threading import Lock
7from typing import Dict, Generator, List, Optional, Set, Tuple, Union, cast
8from uuid import UUID, uuid4
10from fastapi import APIRouter, Depends, WebSocket
11from fastapi_websocket_pubsub import (
12 ALL_TOPICS,
13 EventBroadcaster,
14 PubSubEndpoint,
15 TopicList,
16)
17from fastapi_websocket_pubsub.event_notifier import (
18 EventCallback,
19 SubscriberId,
20 Subscription,
21)
22from fastapi_websocket_pubsub.websocket_rpc_event_notifier import (
23 WebSocketRpcEventNotifier,
24)
25from fastapi_websocket_rpc import RpcChannel
26from opal_common.authentication.deps import WebsocketJWTAuthenticator
27from opal_common.authentication.signer import JWTSigner
28from opal_common.authentication.types import JWTClaims
29from opal_common.authentication.verifier import Unauthorized
30from opal_common.confi.confi import load_conf_if_none
31from opal_common.config import opal_common_config
32from opal_common.logger import logger
33from opal_server.config import opal_server_config
34from opal_server.pubsub_resilience import (
35 FreezablePubSubEndpoint,
36 ReconnectingBroadcaster,
37 SafeConnectionManager,
38)
39from pydantic import BaseModel
40from starlette.datastructures import QueryParams
42OPAL_CLIENT_INFO_PARAM_PREFIX = "__opal_"
43OPAL_CLIENT_INFO_CLIENT_ID = f"{OPAL_CLIENT_INFO_PARAM_PREFIX}client_id"
46class ClientInfo(BaseModel):
47 client_id: str
48 source_host: Optional[str]
49 source_port: Optional[int]
50 connect_time: float
51 subscribed_topics: Set[str] = set()
52 refcount: int = 0 # Only change this while locking ClientTracker._client_lock
53 query_params: Dict[str, str]
56current_client: ContextVar[ClientInfo] = ContextVar("current_client")
59class ClientTracker:
60 def __init__(self):
61 self._clients_by_ids: Dict[str, ClientInfo] = {}
62 self._client_lock = Lock()
64 def clients(self) -> Dict[str, ClientInfo]:
65 return dict(self._clients_by_ids)
67 @contextmanager
68 def new_client(
69 self,
70 source_host: Optional[str],
71 source_port: Optional[int],
72 query_params: QueryParams,
73 ) -> Generator[ClientInfo, None, None]:
74 client_id = f"opal:{uuid4().hex}"
75 if OPAL_CLIENT_INFO_CLIENT_ID in query_params:
76 client_id = query_params.get(OPAL_CLIENT_INFO_CLIENT_ID)
77 elif source_host is not None and source_port is not None:
78 client_id = f"host:{source_host}:{source_port}"
79 with self._client_lock:
80 client_info = self._clients_by_ids.pop(client_id, None)
81 if client_info is None:
82 client_info = ClientInfo(
83 client_id=client_id,
84 source_host=source_host,
85 source_port=source_port,
86 connect_time=time.time(),
87 query_params=query_params,
88 )
89 client_info.refcount += 1
90 self._clients_by_ids[client_id] = client_info
91 yield client_info
92 with self._client_lock:
93 client_info = self._clients_by_ids.pop(client_id)
94 client_info.refcount -= 1
95 if client_info.refcount >= 1:
96 self._clients_by_ids[client_id] = client_info
98 async def on_subscribe(
99 self,
100 subscriber_id: SubscriberId,
101 topics: Union[TopicList, ALL_TOPICS],
102 ):
103 if not isinstance(topics, list): 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true
104 topics = [topics]
106 client_info = current_client.get(None)
108 # on_subscribe is sometimes called for the broadcaster, when there is no "current client"
109 if client_info is not None: 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true
110 client_info.subscribed_topics.update(topics)
112 async def on_unsubscribe(
113 self,
114 subscriber_id: SubscriberId,
115 topics: Union[TopicList, ALL_TOPICS],
116 ):
117 if not isinstance(topics, list): 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 topics = [topics]
120 client_info = current_client.get(None)
122 # on_subscribe is sometimes called for the broadcaster, when there is no "current client"
123 if client_info is not None: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true
124 client_info.subscribed_topics.difference_update(topics)
127class PubSub:
128 """Wrapper for the Pub/Sub channel used for both policy and data
129 updates."""
131 def __init__(self, signer: JWTSigner, broadcaster_uri: str = None):
132 """
133 Args:
134 broadcaster_uri (str, optional): Which server/medium should the PubSub use for broadcasting. Defaults to BROADCAST_URI.
135 None means no broadcasting.
136 """
137 broadcaster_uri = load_conf_if_none(
138 broadcaster_uri, opal_server_config.BROADCAST_URI
139 )
140 self.pubsub_router = APIRouter()
141 self.api_router = APIRouter()
142 # Pub/Sub Internals
143 self.notifier = WebSocketRpcEventNotifier()
144 self.notifier.add_channel_restriction(type(self)._verify_permitted_topics)
145 self.notifier.add_channel_restriction(type(self)._reject_external_purge_channel)
146 self.client_tracker = ClientTracker()
147 self.notifier.register_subscribe_event(self.client_tracker.on_subscribe)
148 self.notifier.register_unsubscribe_event(self.client_tracker.on_unsubscribe)
150 self.broadcaster = None
151 if broadcaster_uri is not None: 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true
152 if opal_server_config.BROADCAST_RECONNECT_ENABLED:
153 logger.info(
154 "Initializing reconnecting broadcaster for server<->server communication"
155 )
156 self.broadcaster = ReconnectingBroadcaster(
157 broadcaster_uri,
158 notifier=self.notifier,
159 channel=opal_server_config.BROADCAST_CHANNEL_NAME,
160 reconnect_max_retries=opal_server_config.BROADCAST_RECONNECT_MAX_RETRIES,
161 reconnect_backoff_min=opal_server_config.BROADCAST_RECONNECT_BACKOFF_MIN_SECONDS,
162 reconnect_backoff_max=opal_server_config.BROADCAST_RECONNECT_BACKOFF_MAX_SECONDS,
163 replay_buffer_size=opal_server_config.BROADCAST_REPLAY_BUFFER_SIZE,
164 resync_settle_seconds=opal_server_config.BROADCAST_RESYNC_SETTLE_SECONDS,
165 )
166 else:
167 logger.info(
168 "Initializing broadcaster for server<->server communication"
169 )
170 self.broadcaster = EventBroadcaster(
171 broadcaster_uri,
172 notifier=self.notifier,
173 channel=opal_server_config.BROADCAST_CHANNEL_NAME,
174 )
175 else:
176 logger.info("Pub/Sub broadcaster is off")
178 # The server endpoint.
179 #
180 # ignore_broadcaster_disconnected=False races each client's main_loop against the
181 # shared broadcaster reader task, so a completed reader surfaces the disconnect and
182 # the client reconnects. That is only safe with ReconnectingBroadcaster, whose
183 # reader stays *pending* across transient drops and completes only after
184 # BROADCAST_RECONNECT_MAX_RETRIES is exhausted (the intended last resort). For the
185 # stock EventBroadcaster (reconnect disabled) the reader dies on the first drop, so
186 # we keep the library-safe default (True) to degrade to "stale but connected"
187 # rather than the fleet-wide drop storm. (Replaces an earlier experimental
188 # broadcast-connection-loss flag.)
189 # The reconnect resync is the freeze's ONLY recovery path (frozen publishes are
190 # dropped, not buffered) — with the resync disabled the combination would silently
191 # lose every update published during every gap, so refuse it rather than honor it.
192 # Both guardrails apply only where a freeze can actually engage, i.e. when the
193 # reconnecting broadcaster (the gap signal) was built above — warning a single-
194 # worker or reconnect-disabled deployment about the resync would be misdirection.
195 freeze_on_disconnect = opal_server_config.BROADCAST_FREEZE_ON_DISCONNECT
196 if freeze_on_disconnect and isinstance( 196 ↛ 199line 196 didn't jump to line 199 because the condition on line 196 was never true
197 self.broadcaster, ReconnectingBroadcaster
198 ):
199 if not opal_server_config.BROADCAST_RESYNC_ON_RECONNECT:
200 logger.warning(
201 "BROADCAST_FREEZE_ON_DISCONNECT is enabled but BROADCAST_RESYNC_ON_RECONNECT "
202 "is disabled — the resync is the freeze's only recovery path, so freezing is "
203 "DISABLED to avoid silently losing updates published during a backbone gap. "
204 "Re-enable BROADCAST_RESYNC_ON_RECONNECT to get the fleet-consistency freeze."
205 )
206 freeze_on_disconnect = False
207 elif freeze_on_disconnect and self.broadcaster is not None: 207 ↛ 208line 207 didn't jump to line 208 because the condition on line 207 was never true
208 logger.info(
209 "BROADCAST_FREEZE_ON_DISCONNECT is enabled but BROADCAST_RECONNECT_ENABLED "
210 "is disabled — the freeze engages only on the reconnecting broadcaster's "
211 "backbone-gap signal, so it is a no-op with the stock broadcaster."
212 )
213 self.endpoint = FreezablePubSubEndpoint(
214 broadcaster=self.broadcaster,
215 notifier=self.notifier,
216 rpc_channel_get_remote_id=opal_common_config.STATISTICS_ENABLED,
217 ignore_broadcaster_disconnected=not isinstance(
218 self.broadcaster, ReconnectingBroadcaster
219 ),
220 # Freeze client-facing publishes during a backbone gap so a write that cannot
221 # reach the whole fleet is not applied on a single worker (see
222 # FreezablePubSubEndpoint). No-op unless the reconnecting broadcaster is in use.
223 freeze_on_disconnect=freeze_on_disconnect,
224 # Exempt: the git-webhook trigger (targets the server-side policy watcher, not
225 # clients — freezing it would drop repo-pull triggers with nothing to replay
226 # them) and the server-to-server coordination channels (statistics + broadcaster
227 # keepalive — dropping those corrupts state no resync rebuilds). The coordination
228 # channels are exempted by their CONFIGURED names: the endpoint's own "__" prefix
229 # rule covers only the defaults, and every one of these is operator-overridable.
230 # Worker-to-worker cache purge on scope delete/repoint.
231 # Freezing it during a backbone gap would leave stale GitPolicyFetcher
232 # caches fleet-wide with nothing to replay the purge. Exempt by its
233 # CONFIGURED name because it is operator-overridable to a non-"__" value.
234 freeze_exempt_topics=[
235 opal_server_config.POLICY_REPO_WEBHOOK_TOPIC,
236 opal_server_config.BROADCAST_KEEPALIVE_TOPIC,
237 opal_server_config.STATISTICS_WAKEUP_CHANNEL,
238 opal_server_config.STATISTICS_STATE_SYNC_CHANNEL,
239 opal_server_config.STATISTICS_SERVER_KEEPALIVE_CHANNEL,
240 opal_common_config.STATISTICS_ADD_CLIENT_CHANNEL,
241 opal_common_config.STATISTICS_REMOVE_CLIENT_CHANNEL,
242 opal_server_config.SCOPES_PURGE_CHANNEL,
243 ],
244 )
245 # fastapi_websocket_rpc's ConnectionManager.disconnect is not idempotent: the RPC
246 # endpoint can call it twice for one socket (handle_disconnect plus the outer
247 # except in WebsocketRPCEndpoint.main_loop), raising
248 # ValueError('list.remove(x): x not in list'). Swap in an idempotent manager
249 # before any connection is served. Reaching into the wrapped endpoint is the
250 # only injection point the library offers (PubSubEndpoint takes no manager),
251 # so fail loudly if its internals ever move (a bare assert would be stripped
252 # under python -O, silently restoring the storm).
253 if not ( 253 ↛ 257line 253 didn't jump to line 257 because the condition on line 253 was never true
254 hasattr(self.endpoint, "endpoint")
255 and hasattr(self.endpoint.endpoint, "manager")
256 ):
257 raise RuntimeError(
258 "Unexpected fastapi_websocket_pubsub internals: cannot install "
259 "SafeConnectionManager (endpoint.endpoint.manager not found)"
260 )
261 self.endpoint.endpoint.manager = SafeConnectionManager()
263 if isinstance(self.broadcaster, ReconnectingBroadcaster): 263 ↛ 264line 263 didn't jump to line 264 because the condition on line 263 was never true
264 self._wire_broadcaster_resync()
266 authenticator = WebsocketJWTAuthenticator(signer)
268 @self.api_router.get(
269 "/pubsub_client_info", response_model=Dict[str, ClientInfo]
270 )
271 async def client_info():
272 return self.client_tracker.clients()
274 @self.pubsub_router.websocket("/ws")
275 async def websocket_rpc_endpoint(
276 websocket: WebSocket, claims: Optional[JWTClaims] = Depends(authenticator)
277 ):
278 """This is the main websocket endpoint the sidecar uses to register
279 on policy updates.
281 as you can see, this endpoint is protected by an HTTP
282 Authorization Bearer token.
283 """
284 try:
285 if claims is None:
286 logger.info(
287 "Closing connection, remote address: {remote_address}",
288 remote_address=websocket.client,
289 reason="Authentication failed",
290 )
291 return
293 source_host = None
294 source_port = None
295 if websocket.client is not None:
296 source_host = websocket.client.host
297 source_port = websocket.client.port
298 with self.client_tracker.new_client(
299 source_host, source_port, websocket.query_params
300 ) as client_info:
301 token = current_client.set(client_info)
302 try:
303 await self.endpoint.main_loop(websocket, claims=claims)
304 finally:
305 current_client.reset(token)
306 finally:
307 await websocket.close()
309 def _wire_broadcaster_resync(self):
310 """Register the post-gap resync.
312 After any backbone gap, force this worker's clients to reconnect
313 so they re-run their full (scope-aware) policy + data
314 reconciliation. Every worker hit the same gap, so each
315 reconciles its own clients and the fleet converges — this is the
316 consistency guarantee; the broadcaster's replay buffer only
317 narrows the staleness window.
318 """
319 manager = self.endpoint.endpoint.manager
320 broadcaster = self.broadcaster
321 resync_enabled = opal_server_config.BROADCAST_RESYNC_ON_RECONNECT
322 settle = opal_server_config.BROADCAST_RESYNC_SETTLE_SECONDS
324 async def _on_broadcaster_reconnect():
325 if not resync_enabled:
326 logger.info("Broadcaster recovered after a gap; client resync disabled")
327 return
328 # Every worker hit the same gap; add a per-worker random delay so the
329 # fleet does not recycle its clients in lockstep.
330 if settle > 0:
331 await asyncio.sleep(random.uniform(0, settle))
332 logger.warning(
333 "Broadcaster recovered after a gap; resyncing this worker's clients "
334 "so they re-fetch current policy + data state"
335 )
336 # The reader is kept alive across this recycle by the listening context the
337 # broadcaster pins around the whole recovery (see _recover_after_gap), so
338 # closing every client here cannot cancel it.
339 await manager.close_all_staggered()
340 if settle > 0:
341 await asyncio.sleep(settle)
343 broadcaster.set_reconnect_callback(_on_broadcaster_reconnect)
345 @staticmethod
346 async def _verify_permitted_topics(
347 topics: Union[TopicList, ALL_TOPICS], channel: RpcChannel
348 ):
349 if "permitted_topics" not in channel.context.get("claims", {}):
350 return
351 unauthorized_topics = set(topics).difference(
352 channel.context["claims"]["permitted_topics"]
353 )
354 if unauthorized_topics:
355 raise Unauthorized(
356 description=f"Invalid 'topics' to subscribe {unauthorized_topics}"
357 )
359 @staticmethod
360 async def _reject_external_purge_channel(
361 topics: Union[TopicList, ALL_TOPICS], channel: RpcChannel
362 ):
363 """Forbid external RPC peers from touching the scope-purge channel.
365 ``SCOPES_PURGE_CHANNEL`` is a server-internal control channel: a purge
366 command evicts every worker's ``GitPolicyFetcher`` caches fleet-wide.
367 The only legitimate publishers
368 are opal-server itself (delete / repoint) and the
369 cross-server broadcaster relay — both call ``notify()`` with
370 ``channel=None``, and channel restrictions run **only when a channel is
371 present** (see ``EventNotifier.notify``/``subscribe``: ``if channel:``).
372 So this restriction never fires for legitimate server traffic; it only
373 sees an external websocket peer (a client/PDP).
375 Without this gate any connected peer could forge a purge and churn the
376 whole fleet's caches, because ``_verify_permitted_topics`` above
377 default-allows tokens that carry no ``permitted_topics`` claim (the
378 common case). No legitimate client ever names this channel — clients
379 only publish to ``STATISTICS_ADD_CLIENT_CHANNEL`` — so rejecting it here
380 blocks the forgery without affecting any real client (publish or
381 subscribe); no client change/redeploy is required.
383 ``ALL_TOPICS`` is rejected too. It is exempt for *publish* (a publish
384 never fans out to a specific-topic subscriber through it), but the same
385 callback also guards *subscribe*, and ``EventNotifier.notify`` fans
386 every published topic — the purge channel included — to the
387 ``ALL_TOPICS`` subscriber bucket. So a peer subscribing to
388 ``ALL_TOPICS`` would still receive purge traffic. No opal-client ever
389 subscribes to ``ALL_TOPICS`` (only the broadcaster does, and that is
390 ``channel=None`` so it never reaches this callback), so rejecting it
391 from external peers is safe.
392 """
393 # Normalize: the notifier may hand us a single topic (str), a list, or
394 # the ALL_TOPICS sentinel (also a str). An external peer may name
395 # neither the purge channel nor ALL_TOPICS.
396 topic_list = [topics] if isinstance(topics, str) else list(topics)
397 if (
398 ALL_TOPICS in topic_list
399 or opal_server_config.SCOPES_PURGE_CHANNEL in topic_list
400 ):
401 raise Unauthorized(
402 description=(
403 f"Topic '{opal_server_config.SCOPES_PURGE_CHANNEL}' (and "
404 "ALL_TOPICS, which would receive it) is server-internal and "
405 "may not be published or subscribed by external peers"
406 )
407 )