From e9e0992dce0e61276a22acba3560f2551ce89e01 Mon Sep 17 00:00:00 2001 From: Jens Langhammer Date: Tue, 27 Apr 2021 16:21:44 +0200 Subject: [PATCH] root: add middleware to properly report websocket connection to sentry Signed-off-by: Jens Langhammer --- authentik/lib/sentry.py | 17 +++++++++++++++++ authentik/root/websocket.py | 7 +++++-- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/authentik/lib/sentry.py b/authentik/lib/sentry.py index df3704ba9..27c5bd020 100644 --- a/authentik/lib/sentry.py +++ b/authentik/lib/sentry.py @@ -5,6 +5,7 @@ from aioredis.errors import ConnectionClosedError, ReplyError from billiard.exceptions import WorkerLostError from botocore.client import ClientError from celery.exceptions import CeleryError +from channels.middleware import BaseMiddleware from channels_redis.core import ChannelFull from django.core.exceptions import SuspiciousOperation, ValidationError from django.db import InternalError, OperationalError, ProgrammingError @@ -14,12 +15,28 @@ from ldap3.core.exceptions import LDAPException from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import RedisError, ResponseError from rest_framework.exceptions import APIException +from sentry_sdk import Hub +from sentry_sdk.tracing import Transaction from structlog.stdlib import get_logger from websockets.exceptions import WebSocketException +from authentik.lib.utils.reflection import class_to_path + LOGGER = get_logger() +class SentryWSMiddleware(BaseMiddleware): + """Sentry Websocket middleweare to set the transaction name based on + consumer class path""" + + async def __call__(self, scope, receive, send): + transaction: Optional[Transaction] = Hub.current.scope.transaction + class_path = class_to_path(self.inner.consumer_class) + if transaction: + transaction.name = class_path + return await self.inner(scope, receive, send) + + class SentryIgnoredException(Exception): """Base Class for all errors that are suppressed, and not sent to sentry.""" diff --git a/authentik/root/websocket.py b/authentik/root/websocket.py index d53b52a12..398b47f62 100644 --- a/authentik/root/websocket.py +++ b/authentik/root/websocket.py @@ -2,10 +2,13 @@ from channels.auth import AuthMiddlewareStack from django.urls import path +from authentik.lib.sentry import SentryWSMiddleware from authentik.outposts.channels import OutpostConsumer from authentik.root.messages.consumer import MessageConsumer websocket_urlpatterns = [ - path("ws/outpost//", OutpostConsumer.as_asgi()), - path("ws/client/", AuthMiddlewareStack(MessageConsumer.as_asgi())), + path("ws/outpost//", SentryWSMiddleware(OutpostConsumer.as_asgi())), + path( + "ws/client/", AuthMiddlewareStack(SentryWSMiddleware(MessageConsumer.as_asgi())) + ), ]