core: cleanup channels code, fix error when server side close

This commit is contained in:
Jens Langhammer 2020-12-13 17:46:34 +01:00
parent 3e49acf7ae
commit 3b5e1c7b34
4 changed files with 26 additions and 14 deletions

View File

@ -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

View File

@ -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):

View File

@ -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()

View File

@ -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