core: cleanup channels code, fix error when server side close
This commit is contained in:
parent
3e49acf7ae
commit
3b5e1c7b34
|
@ -1,4 +1,5 @@
|
|||
"""Channels base classes"""
|
||||
from channels.exceptions import DenyConnection
|
||||
from channels.generic.websocket import JsonWebsocketConsumer
|
||||
from structlog import get_logger
|
||||
|
||||
|
@ -17,16 +18,13 @@ class AuthJsonConsumer(JsonWebsocketConsumer):
|
|||
headers = dict(self.scope["headers"])
|
||||
if b"authorization" not in headers:
|
||||
LOGGER.warning("WS Request without authorization header")
|
||||
self.close()
|
||||
return False
|
||||
raise DenyConnection()
|
||||
|
||||
raw_header = headers[b"authorization"]
|
||||
|
||||
token = token_from_header(raw_header)
|
||||
if not token:
|
||||
LOGGER.warning("Failed to authenticate")
|
||||
self.close()
|
||||
return False
|
||||
raise DenyConnection()
|
||||
|
||||
self.user = token.user
|
||||
return True
|
||||
|
|
|
@ -2,8 +2,9 @@
|
|||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import IntEnum
|
||||
from typing import Any, Dict
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from channels.exceptions import DenyConnection
|
||||
from dacite import from_dict
|
||||
from dacite.data import Data
|
||||
from guardian.shortcuts import get_objects_for_user
|
||||
|
@ -39,18 +40,16 @@ class WebsocketMessage:
|
|||
class OutpostConsumer(AuthJsonConsumer):
|
||||
"""Handler for Outposts that connect over websockets for health checks and live updates"""
|
||||
|
||||
outpost: Outpost
|
||||
outpost: Optional[Outpost] = None
|
||||
|
||||
def connect(self):
|
||||
if not super().connect():
|
||||
return
|
||||
super().connect()
|
||||
uuid = self.scope["url_route"]["kwargs"]["pk"]
|
||||
outpost = get_objects_for_user(
|
||||
self.user, "authentik_outposts.view_outpost"
|
||||
).filter(pk=uuid)
|
||||
if not outpost.exists():
|
||||
self.close()
|
||||
return
|
||||
raise DenyConnection()
|
||||
self.accept()
|
||||
self.outpost = outpost.first()
|
||||
OutpostState(
|
||||
|
@ -60,7 +59,8 @@ class OutpostConsumer(AuthJsonConsumer):
|
|||
|
||||
# pylint: disable=unused-argument
|
||||
def disconnect(self, close_code):
|
||||
OutpostState.for_channel(self.outpost, self.channel_name).delete()
|
||||
if self.outpost:
|
||||
OutpostState.for_channel(self.outpost, self.channel_name).delete()
|
||||
LOGGER.debug("removed channel from cache", channel_name=self.channel_name)
|
||||
|
||||
def receive_json(self, content: Data):
|
||||
|
|
|
@ -97,7 +97,12 @@ def outpost_token_ensurer(self: MonitoredTask):
|
|||
all_outposts = Outpost.objects.all()
|
||||
for outpost in all_outposts:
|
||||
_ = outpost.token
|
||||
self.set_status(TaskResult(TaskResultStatus.SUCCESSFUL, f"Successfully checked {len(all_outposts)} Outposts."))
|
||||
self.set_status(
|
||||
TaskResult(
|
||||
TaskResultStatus.SUCCESSFUL,
|
||||
[f"Successfully checked {len(all_outposts)} Outposts."],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@CELERY_APP.task()
|
||||
|
|
|
@ -105,7 +105,16 @@ class ASGILogger:
|
|||
# https://code.djangoproject.com/ticket/31508
|
||||
# https://github.com/encode/uvicorn/issues/266
|
||||
return
|
||||
await self.app(scope, receive, send_hooked)
|
||||
try:
|
||||
await self.app(scope, receive, send_hooked)
|
||||
except TypeError as exc:
|
||||
# https://github.com/encode/uvicorn/issues/244
|
||||
if exc.args == (
|
||||
"An asyncio.Future, a coroutine or an awaitable is required",
|
||||
):
|
||||
pass
|
||||
else:
|
||||
raise exc
|
||||
|
||||
def _get_ip(self) -> str:
|
||||
client_ip = None
|
||||
|
|
Reference in a new issue