From 007838fcf21dffa34a2fa552bbac43c022761628 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Tue, 6 Jul 2021 14:48:36 +0200
Subject: [PATCH] root: subclass SessionMiddleware to set Secure and SameSite
 flag depending on context

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/root/middleware.py | 89 ++++++++++++++++++++++++++++++++++++
 authentik/root/settings.py   |  6 ++-
 2 files changed, 93 insertions(+), 2 deletions(-)
 create mode 100644 authentik/root/middleware.py

diff --git a/authentik/root/middleware.py b/authentik/root/middleware.py
new file mode 100644
index 000000000..bb9e05044
--- /dev/null
+++ b/authentik/root/middleware.py
@@ -0,0 +1,89 @@
+"""Dynamically set SameSite depending if the upstream connection is TLS or not"""
+import time
+
+from django.conf import settings
+from django.contrib.sessions.backends.base import UpdateError
+from django.contrib.sessions.exceptions import SessionInterrupted
+from django.contrib.sessions.middleware import (
+    SessionMiddleware as UpstreamSessionMiddleware,
+)
+from django.http.request import HttpRequest
+from django.http.response import HttpResponse
+from django.utils.cache import patch_vary_headers
+from django.utils.http import http_date
+
+
+class SessionMiddleware(UpstreamSessionMiddleware):
+    """Dynamically set SameSite depending if the upstream connection is TLS or not"""
+
+    @staticmethod
+    def is_secure(request: HttpRequest) -> bool:
+        """Check if request is TLS'd or localhost"""
+        if request.is_secure():
+            return True
+        host, _, _ = request.get_host().partition(":")
+        if host == "localhost" and settings.DEBUG:
+            return True
+        return False
+
+    def process_response(
+        self, request: HttpRequest, response: HttpResponse
+    ) -> HttpResponse:
+        """
+        If request.session was modified, or if the configuration is to save the
+        session every time, save the changes and set a session cookie or delete
+        the session cookie if the session has been emptied.
+        """
+        try:
+            accessed = request.session.accessed
+            modified = request.session.modified
+            empty = request.session.is_empty()
+        except AttributeError:
+            return response
+        # Set SameSite based on whether or not the request is secure
+        secure = SessionMiddleware.is_secure(request)
+        same_site = "None" if secure else "Lax"
+        # First check if we need to delete this cookie.
+        # The session should be deleted only if the session is entirely empty.
+        if settings.SESSION_COOKIE_NAME in request.COOKIES and empty:
+            response.delete_cookie(
+                settings.SESSION_COOKIE_NAME,
+                path=settings.SESSION_COOKIE_PATH,
+                domain=settings.SESSION_COOKIE_DOMAIN,
+                samesite=same_site,
+            )
+            patch_vary_headers(response, ("Cookie",))
+        else:
+            if accessed:
+                patch_vary_headers(response, ("Cookie",))
+            if (modified or settings.SESSION_SAVE_EVERY_REQUEST) and not empty:
+                if request.session.get_expire_at_browser_close():
+                    max_age = None
+                    expires = None
+                else:
+                    max_age = request.session.get_expiry_age()
+                    expires_time = time.time() + max_age
+                    expires = http_date(expires_time)
+                # Save the session data and refresh the client cookie.
+                # Skip session save for 500 responses, refs #3881.
+                if response.status_code != 500:
+                    try:
+                        request.session.save()
+                    except UpdateError:
+                        raise SessionInterrupted(
+                            "The request's session was deleted before the "
+                            "request completed. The user may have logged "
+                            "out in a concurrent request, for example."
+                        )
+                    response.set_cookie(
+                        settings.SESSION_COOKIE_NAME,
+                        request.session.session_key,
+                        max_age=max_age,
+                        expires=expires,
+                        domain=settings.SESSION_COOKIE_DOMAIN,
+                        path=settings.SESSION_COOKIE_PATH,
+                        secure=secure,
+                        httponly=settings.SESSION_COOKIE_HTTPONLY or None,
+                        samesite=same_site,
+                    )
+        return response
diff --git a/authentik/root/settings.py b/authentik/root/settings.py
index 638bcc169..3b12469a0 100644
--- a/authentik/root/settings.py
+++ b/authentik/root/settings.py
@@ -203,14 +203,16 @@ DJANGO_REDIS_IGNORE_EXCEPTIONS = True
 DJANGO_REDIS_LOG_IGNORED_EXCEPTIONS = True
 SESSION_ENGINE = "django.contrib.sessions.backends.cache"
 SESSION_CACHE_ALIAS = "default"
-SESSION_COOKIE_SAMESITE = "None"
+# Configured via custom SessionMiddleware
+# SESSION_COOKIE_SAMESITE = "None"
+# SESSION_COOKIE_SECURE = True
 SESSION_EXPIRE_AT_BROWSER_CLOSE = True
 
 MESSAGE_STORAGE = "authentik.root.messages.storage.ChannelsStorage"
 
 MIDDLEWARE = [
     "django_prometheus.middleware.PrometheusBeforeMiddleware",
-    "django.contrib.sessions.middleware.SessionMiddleware",
+    "authentik.root.middleware.SessionMiddleware",
     "django.contrib.auth.middleware.AuthenticationMiddleware",
     "authentik.core.middleware.RequestIDMiddleware",
     "authentik.tenants.middleware.TenantMiddleware",