mirror of
https://github.com/revng/revng
synced 2026-06-21 14:07:57 +00:00
Giacomo Vercesi
828bd9c49f
Add the PRSS (Pypeline Remote Storage Server) and the relay companion server.
122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
#
|
|
# This file is distributed under the MIT License. See LICENSE.md for details.
|
|
#
|
|
|
|
import asyncio
|
|
from asyncio import Queue
|
|
from collections import defaultdict
|
|
from contextlib import asynccontextmanager
|
|
from typing import override
|
|
|
|
from starlette.applications import Starlette
|
|
from starlette.requests import Request
|
|
from starlette.responses import PlainTextResponse
|
|
from starlette.routing import Route, WebSocketRoute
|
|
from starlette.websockets import WebSocket
|
|
|
|
from websockets.frames import CloseCode
|
|
|
|
from revng.pypeline.storage.notification_queue import MultiQueue
|
|
from revng.pypeline.storage.storage_provider import ProjectID
|
|
from revng.pypeline.utils.notification_broker import NotificationBroker, NotificationSubscriber
|
|
from revng.pypeline.utils.notification_broker import Stream
|
|
from revng.pypeline.utils.starlette import NotificationWebsocket, get_middlewares, get_project_id
|
|
|
|
shutdown_begun: asyncio.Event | None = None
|
|
|
|
|
|
class MultiNotificationBroker(NotificationBroker):
|
|
def __init__(self, queues: dict[ProjectID, MultiQueue[bytes]]):
|
|
super().__init__()
|
|
self.queues = queues
|
|
|
|
@override
|
|
async def subscribe(
|
|
self, project_id: ProjectID | None, stream: Stream
|
|
) -> NotificationSubscriber:
|
|
assert project_id is not None, "project_id = None is unsupported"
|
|
return await super().subscribe(project_id, stream)
|
|
|
|
@override
|
|
async def get_queue(self, project_id: ProjectID | None) -> Queue[bytes]:
|
|
assert project_id is not None
|
|
return self.queues[project_id].get_queue()
|
|
|
|
|
|
def _parse_bind(string: str) -> tuple[str, int]:
|
|
# TODO: fix for IPv6
|
|
assert string.count(":") == 1
|
|
parts = string.split(":")
|
|
return (parts[0], int(parts[1]))
|
|
|
|
|
|
def _compare_connection(bind: tuple[str, int], connection: tuple[str, int]):
|
|
"""
|
|
Given a bind address, e.g. `("127.0.0.1", 8000)` check that the connection
|
|
address matches
|
|
"""
|
|
|
|
# TODO: fix for IPv6
|
|
if bind[0] == "0.0.0.0":
|
|
# If the server is bound on all interfaces just check that the port
|
|
# matches
|
|
return bind[1] == connection[1]
|
|
else:
|
|
return bind == connection
|
|
|
|
|
|
def make_starlette(
|
|
production: bool, notifications_bind: str, publish_bind: str, publish_psk: str
|
|
) -> Starlette:
|
|
queues: dict[ProjectID, MultiQueue[bytes]] = defaultdict(MultiQueue)
|
|
notification_broker = MultiNotificationBroker(queues)
|
|
ws_notifications = NotificationWebsocket(notification_broker, lambda: shutdown_begun)
|
|
|
|
notifications_tuple = _parse_bind(notifications_bind)
|
|
publish_tuple = _parse_bind(publish_bind)
|
|
|
|
async def status(request):
|
|
return PlainTextResponse("OK")
|
|
|
|
async def publish(request: Request):
|
|
if not _compare_connection(publish_tuple, request["server"]):
|
|
return PlainTextResponse(status_code=403)
|
|
|
|
project_id = get_project_id(request.headers)
|
|
if project_id is None:
|
|
return PlainTextResponse("Project ID must be supplied", 400)
|
|
|
|
# NOTE: at some point this static PSK will be dropped in favor of
|
|
# mTLS, but it has a high upfront infrastructural cost
|
|
authorization = request.headers.get("authorization")
|
|
if authorization != f"Bearer {publish_psk}":
|
|
return PlainTextResponse("Invalid authorization header", 403)
|
|
|
|
body = await request.body()
|
|
queues[project_id].send(body)
|
|
return PlainTextResponse("Sent")
|
|
|
|
async def notifications(websocket: WebSocket):
|
|
if not _compare_connection(notifications_tuple, websocket["server"]):
|
|
await websocket.close(CloseCode.POLICY_VIOLATION)
|
|
return
|
|
return await ws_notifications.endpoint(websocket)
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app):
|
|
global shutdown_begun
|
|
shutdown_begun = asyncio.Event()
|
|
yield
|
|
|
|
# Create the Starlette application
|
|
return Starlette(
|
|
debug=not production,
|
|
routes=[
|
|
Route("/publish", publish, methods=["POST"]),
|
|
WebSocketRoute("/notifications", notifications),
|
|
Route("/status", status, methods=["GET"]),
|
|
],
|
|
middleware=get_middlewares(production, unauthenticated_paths={"/publish"}),
|
|
lifespan=lifespan,
|
|
)
|