diff --git a/README.md b/README.md index a4ae9b7d5..8690211d4 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,23 @@ docker-compose up -d DigiScript will be available at `http://localhost:8080` +### Request size metrics + +The server exposes two Prometheus histograms at `/debug/metrics` so that request size limits can be +tuned from real usage (Tornado's defaults currently apply: 100 MB HTTP bodies, 10 MiB WebSocket messages): + +| Metric | Labels | Meaning | +|--------|--------|---------| +| `digiscript_http_request_body_bytes` | `handler`, `method` | Size of HTTP request bodies (requests with a body only) | +| `digiscript_ws_message_bytes` | none | Size of inbound WebSocket messages | + +Buckets are powers of 4 from 1 KiB to 256 MiB. For example, the 99th percentile WebSocket message size +over the last day in Prometheus: + +``` +histogram_quantile(0.99, sum(rate(digiscript_ws_message_bytes_bucket[1d])) by (le)) +``` + ### Manual Installation #### Client @@ -110,6 +127,8 @@ python3 main.py The server will be available at `http://localhost:8080` +The server runs in production mode by default. When developing (for example with the `client-v3` Vite dev server), start it with `python3 main.py --debug` to enable autoreload and accept WebSocket connections from other origins. Never use `--debug` on a deployed server. + ### First Launch Setup On first launch, DigiScript will prompt you to create an administrator user. From there, you can create shows, configure users, and begin building your production. diff --git a/client-v3/README.md b/client-v3/README.md index e929a4803..49e6717a4 100644 --- a/client-v3/README.md +++ b/client-v3/README.md @@ -12,7 +12,7 @@ pnpm install pnpm run dev ``` -The dev server starts at `http://localhost:5173` and proxies API requests to `http://localhost:8080`. +The dev server starts at `http://localhost:5173` and proxies API requests to `http://localhost:8080`. Run the backend with `python3 main.py --debug` while developing. ### Compiles and minifies for production ``` diff --git a/client-v3/src/composables/useWebSocket.closeTooBig.test.ts b/client-v3/src/composables/useWebSocket.closeTooBig.test.ts new file mode 100644 index 000000000..cc2dcd563 --- /dev/null +++ b/client-v3/src/composables/useWebSocket.closeTooBig.test.ts @@ -0,0 +1,147 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; + +// Close code 1009 (message too big) must back off instead of reconnecting at the minimum +// delay forever. + +vi.mock('@/router', () => ({ default: { currentRoute: { value: { path: '/' } }, push: vi.fn() } })); +vi.mock('@/js/toast', () => ({ + toast: { info: vi.fn(), error: vi.fn(), success: vi.fn() }, +})); + +class FakeWebSocket { + static OPEN = 1; + static CONNECTING = 0; + static instances: FakeWebSocket[] = []; + readyState = 3; + onopen: (() => Promise | void) | null = null; + onmessage: ((e: { data: string }) => void) | null = null; + onclose: ((e?: { code: number }) => void) | null = null; + onerror: (() => void) | null = null; + + constructor() { + FakeWebSocket.instances.push(this); + } + + send(): void {} +} + +// The composable keeps its backoff state at module level, so load a fresh copy per test. +let useWebSocket: typeof import('./useWebSocket').useWebSocket; +let toast: typeof import('@/js/toast').toast; + +describe('close code 1009', () => { + beforeEach(async () => { + vi.resetModules(); + FakeWebSocket.instances = []; + vi.useFakeTimers(); + vi.stubGlobal('WebSocket', FakeWebSocket); + const { createPinia, setActivePinia } = await import('pinia'); + setActivePinia(createPinia()); + ({ useWebSocket } = await import('./useWebSocket')); + ({ toast } = await import('@/js/toast')); + }); + + afterEach(() => { + vi.useRealTimers(); + vi.unstubAllGlobals(); + }); + + it('re-syncs the live session after reconnecting from a 1009 close', async () => { + const { useShowStore } = await import('@/stores/show'); + const showStore = useShowStore(); + showStore.currentSession = { id: 1 } as never; + const resync = vi.spyOn(showStore, 'getShowSessionData').mockResolvedValue(); + + useWebSocket().connect(); + FakeWebSocket.instances[0].onclose?.({ code: 1009 }); + vi.advanceTimersByTime(2100); + await FakeWebSocket.instances[1].onopen?.(); + + expect(resync).toHaveBeenCalledTimes(1); + expect(toast.success).toHaveBeenCalledWith('WebSocket reconnected after 1 attempt'); + }); + + it('forgets earlier 1009 closes once a connection has stayed up for 30s', async () => { + useWebSocket().connect(); + FakeWebSocket.instances[0].onclose?.({ code: 1009 }); + vi.advanceTimersByTime(2100); + // Second connection stays up long enough to clear the earlier 1009 count. + await FakeWebSocket.instances[1].onopen?.(); + vi.advanceTimersByTime(30100); + FakeWebSocket.instances[1].onclose?.({ code: 1009 }); + // Count was cleared, so this is a first close again: 2s delay, not 4s. + vi.advanceTimersByTime(2100); + expect(FakeWebSocket.instances).toHaveLength(3); + }); + + it('keeps the 1009 count when the connection was up for less than 30s', async () => { + useWebSocket().connect(); + FakeWebSocket.instances[0].onclose?.({ code: 1009 }); + vi.advanceTimersByTime(2100); + await FakeWebSocket.instances[1].onopen?.(); + vi.advanceTimersByTime(10000); + FakeWebSocket.instances[1].onclose?.({ code: 1009 }); + // Second consecutive close: 4s delay. + vi.advanceTimersByTime(3500); + expect(FakeWebSocket.instances).toHaveLength(2); + vi.advanceTimersByTime(600); + expect(FakeWebSocket.instances).toHaveLength(3); + }); + + it('keeps backing off when authentication succeeds between closes', async () => { + useWebSocket().connect(); + FakeWebSocket.instances[0].onclose?.({ code: 1009 }); + vi.advanceTimersByTime(2100); + expect(FakeWebSocket.instances).toHaveLength(2); + + // Reconnect and re-authenticate (this resets the generic error counter). + FakeWebSocket.instances[1].onmessage?.({ data: JSON.stringify({ OP: 'WS_AUTH_SUCCESS' }) }); + await vi.advanceTimersByTimeAsync(0); + FakeWebSocket.instances[1].onclose?.({ code: 1009 }); + + // Delay must have doubled to 4s rather than starting over at 2s. + vi.advanceTimersByTime(3500); + expect(FakeWebSocket.instances).toHaveLength(2); + vi.advanceTimersByTime(600); + expect(FakeWebSocket.instances).toHaveLength(3); + expect(toast.error).toHaveBeenCalledTimes(1); + }); + + it('reconnects even if showing the toast throws', () => { + vi.mocked(toast.error).mockImplementationOnce(() => { + throw new Error('toast failed'); + }); + useWebSocket().connect(); + FakeWebSocket.instances[0].onclose?.({ code: 1009 }); + vi.advanceTimersByTime(2100); + expect(FakeWebSocket.instances).toHaveLength(2); + }); + + it('does not back off or notify on a normal close', () => { + useWebSocket().connect(); + FakeWebSocket.instances[0].onclose?.({ code: 1000 }); + expect(toast.error).not.toHaveBeenCalled(); + vi.advanceTimersByTime(1100); + expect(FakeWebSocket.instances).toHaveLength(2); + }); + + it('notifies the user once and backs off between reconnects', () => { + useWebSocket().connect(); + FakeWebSocket.instances[0].onclose?.({ code: 1009 }); + expect(toast.error).toHaveBeenCalledTimes(1); + + // First reconnect delay is 2s (1s * 2^1): nothing after 1.5s, a new socket after 2s. + vi.advanceTimersByTime(1500); + expect(FakeWebSocket.instances).toHaveLength(1); + vi.advanceTimersByTime(600); + expect(FakeWebSocket.instances).toHaveLength(2); + + FakeWebSocket.instances[1].onclose?.({ code: 1009 }); + expect(toast.error).toHaveBeenCalledTimes(1); + // Second delay doubles to 4s. + vi.advanceTimersByTime(3500); + expect(FakeWebSocket.instances).toHaveLength(2); + vi.advanceTimersByTime(600); + expect(FakeWebSocket.instances).toHaveLength(3); + }); +}); diff --git a/client-v3/src/composables/useWebSocket.ts b/client-v3/src/composables/useWebSocket.ts index 28f3193a8..38e6d3d4e 100644 --- a/client-v3/src/composables/useWebSocket.ts +++ b/client-v3/src/composables/useWebSocket.ts @@ -12,13 +12,24 @@ import type { WsMessage } from '@/types/api/websocket'; const INITIAL_RECONNECT_DELAY_MS = 1000; const MAX_RECONNECT_DELAY_MS = 30000; +/** RFC 6455 close code sent by the server when a message exceeds its size limit. */ +const WS_CLOSE_MESSAGE_TOO_BIG = 1009; +const STABLE_CONNECTION_MS = 30000; let ws: WebSocket | null = null; let reconnectTimer: ReturnType | null = null; let errorCount = 0; +// Consecutive 1009 closes. Kept apart from errorCount, which a successful authentication resets, +// so a client that keeps tripping the limit still backs off. Cleared once a connection has stayed +// up for STABLE_CONNECTION_MS. +let tooBigCount = 0; +let stableTimer: ReturnType | null = null; function getReconnectDelay(): number { - return Math.min(INITIAL_RECONNECT_DELAY_MS * 2 ** errorCount, MAX_RECONNECT_DELAY_MS); + return Math.min( + INITIAL_RECONNECT_DELAY_MS * 2 ** Math.max(errorCount, tooBigCount), + MAX_RECONNECT_DELAY_MS + ); } /** Returns false (after logging) if the socket isn't open, so callers can react to a dropped frame. */ @@ -177,14 +188,18 @@ function connect(): void { ws.onopen = async () => { try { - const wasErrored = errorCount > 0; + const attempts = Math.max(errorCount, tooBigCount); + const wasErrored = attempts > 0; wsStore.$patch({ isConnected: true }); if (wasErrored) { - toast.success( - `WebSocket reconnected after ${errorCount} attempt${errorCount > 1 ? 's' : ''}` - ); + toast.success(`WebSocket reconnected after ${attempts} attempt${attempts > 1 ? 's' : ''}`); } log.info('WebSocket connected'); + if (stableTimer) clearTimeout(stableTimer); + stableTimer = setTimeout(() => { + stableTimer = null; + tooBigCount = 0; + }, STABLE_CONNECTION_MS); if (wasErrored) { const showStore = useShowStore(); if (showStore.currentSession != null) { @@ -205,10 +220,30 @@ function connect(): void { } }; - ws.onclose = () => { + ws.onclose = (event?: CloseEvent) => { wsStore.$patch({ isConnected: false, authenticated: false }); + if (stableTimer) { + clearTimeout(stableTimer); + stableTimer = null; + } + const tooBig = event?.code === WS_CLOSE_MESSAGE_TOO_BIG; + if (tooBig) { + // The server dropped us for sending an over-sized message. Back off (and tell the user) + // rather than reconnecting at the minimum delay, which would loop if the oversized + // message is re-sent on every connect. + tooBigCount++; + log.error('WebSocket closed by server: message exceeded the maximum size'); + } log.info('WebSocket closed, scheduling reconnect'); + // Schedule before notifying so a failing toast can never stop the reconnect. scheduleReconnect(); + if (tooBig && tooBigCount === 1) { + try { + toast.error('Server closed the connection: message too large'); + } catch (e) { + log.error('Failed to show WebSocket close toast:', e); + } + } }; ws.onerror = () => { diff --git a/server/README.md b/server/README.md index 4c16b6a7f..5e882f882 100644 --- a/server/README.md +++ b/server/README.md @@ -10,6 +10,8 @@ pip install -r requirements.txt python3 ./main.py ``` +Pass `--debug` for development (autoreload, relaxed WebSocket origin checks). It is off by default and must not be used on a deployed server. + This requires that the client has been build first in order to server the static files ## Project Structure diff --git a/server/controllers/api/v1/show/shows.py b/server/controllers/api/v1/show/shows.py index 6c295714e..3e38e236b 100644 --- a/server/controllers/api/v1/show/shows.py +++ b/server/controllers/api/v1/show/shows.py @@ -244,6 +244,8 @@ async def patch(self): self.set_status(404) await self.finish({"message": ERROR_SHOW_NOT_FOUND}) + @api_authenticated + @require_admin @requires_show @no_live_session async def delete(self): diff --git a/server/controllers/ws_controller.py b/server/controllers/ws_controller.py index 5953c27dc..295dcb5f6 100644 --- a/server/controllers/ws_controller.py +++ b/server/controllers/ws_controller.py @@ -32,6 +32,7 @@ from rbac.role import Role from utils.web.base_controller import DatabaseMixin from utils.web.pending_disconnects import disconnect_key +from utils.web.request_size_metrics import WSMessageSizeMetricsMixin from utils.web.route import ApiRoute, ApiVersion from utils.web.ws_session_lifecycle import ( assign_session_user, @@ -63,7 +64,7 @@ @ApiRoute("ws", ApiVersion.V1) -class WebSocketController(DatabaseMixin, WebSocketHandler): +class WebSocketController(WSMessageSizeMetricsMixin, DatabaseMixin, WebSocketHandler): # Live-show ops only the leader may send, and the method that handles each. _LEADER_OP_HANDLERS = { "SCRIPT_SCROLL": "_op_script_scroll", diff --git a/server/digi_server/app_server.py b/server/digi_server/app_server.py index f80b0952a..6eefcd6ba 100644 --- a/server/digi_server/app_server.py +++ b/server/digi_server/app_server.py @@ -44,6 +44,7 @@ from utils.version_checker import VersionChecker from utils.web.jwt_service import JWTService from utils.web.pending_disconnects import PendingDisconnects +from utils.web.request_size_metrics import observe_http_request from utils.web.route import Route @@ -331,6 +332,8 @@ class AlembicVersion(self._db.Model): super().__init__( handlers=handlers, debug=debug, + # Never return stack traces to clients, even in debug mode. + serve_traceback=False, db=self._db, websocket_ping_interval=5, cookie_secret=self._get_cookie_secret(), @@ -340,6 +343,8 @@ class AlembicVersion(self._db.Model): def log_request(self, handler): from tornado.log import access_log # noqa: PLC0415 + observe_http_request(handler) + if handler.request.path in Route.ignored_logging_routes(): return diff --git a/server/main.py b/server/main.py index 80e8994d9..ccbf5b02a 100755 --- a/server/main.py +++ b/server/main.py @@ -38,7 +38,13 @@ def ensure_writable_db_path(path): add_logging_level("TRACE", logging.DEBUG - 5) get_logger().setLevel(logging.DEBUG) -define("debug", type=bool, default=not IS_FROZEN, help="auto reload") +define( + "debug", + type=bool, + default=False, + help="Development mode: autoreload, relaxed WebSocket origin checks. " + "Never enable on a deployed server.", +) define("port", type=int, default=8080, help="port to listen on") define("settings_path", type=str, default=None, help="Path to settings JSON file") define("skip_migrations", type=bool, default=False, help="skip database migrations") diff --git a/server/test/controllers/api/v1/show/test_shows.py b/server/test/controllers/api/v1/show/test_shows.py index be77ddc43..e70df019a 100644 --- a/server/test/controllers/api/v1/show/test_shows.py +++ b/server/test/controllers/api/v1/show/test_shows.py @@ -420,6 +420,40 @@ def test_delete_show_cleans_up_rbac_mapping_rows(self): " the script", ) + def test_delete_show_unauthenticated(self): + """An unauthenticated request cannot delete a show.""" + show_b_id = self._create_show() + response = self.fetch( + f"/api/v1/show?id={show_b_id}", + method="DELETE", + allow_nonstandard_methods=True, + ) + self.assertEqual(401, response.code) + with self._app.get_db().sessionmaker() as session: + self.assertIsNotNone(session.get(Show, show_b_id)) + + def test_delete_show_unauthenticated_nonexistent_show(self): + """Auth is checked before the show lookup, so existence is not leaked.""" + response = self.fetch( + "/api/v1/show?id=99999", method="DELETE", allow_nonstandard_methods=True + ) + self.assertEqual(401, response.code) + + def test_delete_show_non_admin(self): + """A logged-in non-admin cannot delete a show.""" + show_b_id = self._create_show() + with self._app.get_db().sessionmaker() as session: + user = User(username="regular", is_admin=False, password="test") + session.add(user) + session.flush() + user_id = user.id + session.commit() + token = self._app.jwt_service.create_access_token({"user_id": user_id}) + response = self._delete(show_b_id, token=token) + self.assertEqual(401, response.code) + with self._app.get_db().sessionmaker() as session: + self.assertIsNotNone(session.get(Show, show_b_id)) + def test_delete_currently_loaded_show(self): """Attempting to delete the currently loaded show returns 400.""" response = self._delete(self.show_a_id) diff --git a/server/test/controllers/test_ws_controller.py b/server/test/controllers/test_ws_controller.py index 08525d2fc..2d2251f88 100644 --- a/server/test/controllers/test_ws_controller.py +++ b/server/test/controllers/test_ws_controller.py @@ -16,7 +16,7 @@ from sqlalchemy import select from tornado.httpclient import HTTPRequest from tornado.testing import gen_test -from tornado.websocket import websocket_connect +from tornado.websocket import WebSocketClosedError, websocket_connect from controllers.ws_controller import WebSocketController from digi_server.logger import get_logger @@ -36,6 +36,7 @@ disconnect_key, room_close_key, ) +from utils.web.request_size_metrics import ws_message_bytes from utils.web.ws_session_lifecycle import FINALISE_ATTEMPTS @@ -3666,6 +3667,64 @@ async def test_force_logout_of_leader_hands_leadership_on(self): ws_l.close() +class TestWSMessageSizeMetrics(DigiScriptTestCase): + @staticmethod + def _message_count(): + return next( + sample.value + for metric in ws_message_bytes.collect() + for sample in metric.samples + if sample.name.endswith("_count") + ) + + @gen_test + async def test_inbound_message_size_observed(self): + """Every inbound WebSocket message is recorded in the size histogram.""" + before = self._message_count() + conn = await websocket_connect( + f"ws://127.0.0.1:{self.get_http_port()}/api/v1/ws" + ) + await conn.read_message() # SET_UUID greeting + await conn.write_message(json.dumps({"OP": "PING", "DATA": "x" * 100})) + await conn.write_message("not json") + for _ in range(50): + if self._message_count() >= before + 2: + break + await asyncio.sleep(0.05) + conn.close() + self.assertGreaterEqual(self._message_count(), before + 2) + + +class TestWSOversizedMessage(DigiScriptTestCase): + @gen_test + async def test_message_over_limit_closes_connection(self): + """The server ends the connection when a message exceeds the size cap. + + The cap is Tornado's ``websocket_max_message_size`` (lowered here). + Tornado sends a 1009 close frame and then drops the socket while the + client is still uploading, so a client may observe 1009 or an abnormal + close; only the termination is asserted. + """ + limit = 1024 * 1024 + self._app.settings["websocket_max_message_size"] = limit + conn = await websocket_connect( + f"ws://127.0.0.1:{self.get_http_port()}/api/v1/ws" + ) + # Control case: a message under the cap leaves the connection open. + await conn.write_message("x" * (limit - 1024)) + await asyncio.sleep(0.2) + conn.ping(b"") + self.assertIsNone(conn.close_code) + try: + await conn.write_message("x" * (limit + 1024)) + except WebSocketClosedError: + pass # the server closed while the oversized frame was being sent + message = await conn.read_message() + while message is not None: + message = await conn.read_message() + self.assertIsNone(message) + + class TestWSHandlerMaps(DigiScriptTestCase): """Every op in the dispatch maps resolves to a handler method.""" diff --git a/server/test/digi_server/test_digi_server.py b/server/test/digi_server/test_digi_server.py index 0572580ce..e272b0ad2 100644 --- a/server/test/digi_server/test_digi_server.py +++ b/server/test/digi_server/test_digi_server.py @@ -1,3 +1,6 @@ +import ast +from pathlib import Path + import tornado.escape from sqlalchemy import func, select from tornado.testing import gen_test @@ -6,6 +9,7 @@ from models.settings import SystemSettings from models.user import User from test.conftest import DigiScriptTestCase +from utils.web.request_size_metrics import http_request_body_bytes class TestDigiScriptServer(DigiScriptTestCase): @@ -17,6 +21,79 @@ def test_debug(self): self.assertTrue("status" in response_body) self.assertEqual("OK", response_body["status"]) + def test_debug_mode_is_opt_in(self): + """Debug mode must be opt-in via ``--debug``.""" + tree = ast.parse( + (Path(__file__).parents[2] / "main.py").read_text(encoding="utf-8") + ) + defaults = [ + getattr(kw.value, "value", "") + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and getattr(node.func, "id", None) == "define" + and node.args + and getattr(node.args[0], "value", None) == "debug" + for kw in node.keywords + if kw.arg == "default" + ] + self.assertEqual( + [False], + defaults, + "--debug must default to the literal False (update this test if the " + "default is computed)", + ) + + def test_traceback_not_returned_to_client(self): + """Even with debug=True (as in the test app) no traceback is returned.""" + self.assertTrue(self._app.settings["debug"]) + self.assertFalse(self._app.settings["serve_traceback"]) + response = self.fetch( + "/api/v1/auth/login", + method="POST", + body="{not valid json", + ) + self.assertGreaterEqual(response.code, 400) + self.assertNotIn(b"Traceback", response.body) + self.assertNotIn(b".py", response.body) + + def test_request_body_size_observed(self): + """Request bodies are recorded in the size histogram, labelled by handler.""" + labels = ("ApiFallback", "POST") + before = self._body_count(*labels) + self.fetch("/api/v1/nonexistent", method="POST", body="x" * 2000) + self.assertEqual(before + 1, self._body_count(*labels)) + + def test_unknown_http_method_does_not_create_new_series(self): + """Arbitrary method strings are reported as OTHER, keeping labels bounded.""" + before = self._body_count("ApiFallback", "OTHER") + for method in ("FOOA1", "FOOB2", "FOOC3"): + self.fetch( + "/api/v1/nonexistent", + method=method, + body="x" * 10, + allow_nonstandard_methods=True, + ) + self.assertEqual(before + 3, self._body_count("ApiFallback", "OTHER")) + self.assertEqual(0, self._body_count("ApiFallback", "FOOA1")) + + def test_request_without_body_not_observed(self): + """Body-less requests do not add to the size histogram.""" + before = self._body_count("ApiFallback", "GET") + self.fetch("/api/v1/nonexistent") + self.assertEqual(before, self._body_count("ApiFallback", "GET")) + + @staticmethod + def _body_count(handler, method): + for metric in http_request_body_bytes.collect(): + for sample in metric.samples: + if ( + sample.name.endswith("_count") + and sample.labels.get("handler") == handler + and sample.labels.get("method") == method + ): + return sample.value + return 0 + def test_api_debug(self): response = self.fetch("/api/v1/debug") response_body = tornado.escape.json_decode(response.body) diff --git a/server/utils/web/request_size_metrics.py b/server/utils/web/request_size_metrics.py new file mode 100644 index 000000000..21e33ad92 --- /dev/null +++ b/server/utils/web/request_size_metrics.py @@ -0,0 +1,92 @@ +"""Prometheus histograms of inbound HTTP body and WebSocket message sizes. + +The histograms are module-level so they are registered exactly once per +process, however many application instances are created (tests create many). +They are exposed alongside the other metrics at ``/debug/metrics`` and exist so +that request size limits can be tuned from real data. +""" + +from __future__ import annotations + +import functools +import inspect + +from prometheus_client import Histogram + + +# Powers of 4 from 1 KiB to 256 MiB. +SIZE_BUCKETS = tuple(1024 * 4**i for i in range(10)) + +http_request_body_bytes = Histogram( + name="digiscript_http_request_body_bytes", + documentation="Size of HTTP request bodies in bytes (requests with a body only)", + buckets=SIZE_BUCKETS, + labelnames=("handler", "method"), +) + +ws_message_bytes = Histogram( + name="digiscript_ws_message_bytes", + documentation="Size of inbound WebSocket messages in bytes", + buckets=SIZE_BUCKETS, +) + + +_KNOWN_METHODS = frozenset({"GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"}) + + +def observe_http_request(handler) -> None: + """Record the body size of a finished HTTP request, if it had a body. + + Both labels are bounded: the handler class name (never the raw path) and + the HTTP method, with anything outside the standard set reported as + ``OTHER`` so a client cannot create unlimited series. + + :param handler: The finished Tornado ``RequestHandler``. + """ + size = len(handler.request.body or b"") + if size: + method = handler.request.method + http_request_body_bytes.labels( + type(handler).__name__, method if method in _KNOWN_METHODS else "OTHER" + ).observe(size) + + +class WSMessageSizeMetricsMixin: + """Observes the size in bytes of every message passed to ``on_message``. + + ``on_message`` is wrapped when the subclass is created, so the subclass's + own method body needs no changes. The subclass must define an ``async`` + ``on_message`` directly in its class body; anything else fails at import + time rather than silently recording nothing. A subclass that overrides + ``on_message`` and calls ``super().on_message`` is observed once per + message, not twice. + """ + + _observing_ws_size = False + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + original = cls.__dict__.get("on_message") + if original is None or not inspect.iscoroutinefunction(original): + raise TypeError( + f"{cls.__name__} must define an async on_message to use " + "WSMessageSizeMetricsMixin" + ) + + @functools.wraps(original) + async def on_message(self, message): + if self._observing_ws_size: + return await original(self, message) + size = ( + len(message.encode("utf-8")) + if isinstance(message, str) + else len(message) + ) + ws_message_bytes.observe(size) + self._observing_ws_size = True + try: + return await original(self, message) + finally: + self._observing_ws_size = False + + cls.on_message = on_message