From aa7540045b00d072a1fe5a469c9b68caca9e0905 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Wed, 18 May 2022 14:42:54 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20session=20guard?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/authentication/middleware.py      | 20 +++++++++++++++++---
 apps/authentication/signal_handlers.py |  3 +++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/apps/authentication/middleware.py b/apps/authentication/middleware.py
index 96d0017e9..f1f60bbc5 100644
--- a/apps/authentication/middleware.py
+++ b/apps/authentication/middleware.py
@@ -60,14 +60,28 @@ class SessionCookieMiddleware(MiddlewareMixin):
         response.set_cookie(pub_key_name, public_key_decode)
 
     @staticmethod
-    def set_session_cooke_prefix(request, response):
+    def set_cookie_session_prefix(request, response):
         key = settings.SESSION_COOKIE_NAME_PREFIX_KEY
         value = settings.SESSION_COOKIE_NAME_PREFIX
         if request.COOKIES.get(key) == value:
             return response
         response.set_cookie(key, value)
 
+    @staticmethod
+    def set_cookie_session_expire(request, response):
+        if not request.session.get('auth_session_expiration_required'):
+            return
+        value = 'age'
+        if settings.SESSION_EXPIRE_AT_BROWSER_CLOSE_FORCE or \
+                not request.session.get('auto_login', False):
+            value = 'close'
+
+        age = request.session.get_expiry_age()
+        response.set_cookie('jms_session_expire', value, max_age=age)
+        request.session.pop('auth_session_expiration_required', None)
+
     def process_response(self, request, response: HttpResponse):
-        self.set_session_cooke_prefix(request, response)
-        self.set_session_cooke_prefix(request, response)
+        self.set_cookie_session_prefix(request, response)
+        self.set_cookie_public_key(request, response)
+        self.set_cookie_session_expire(request, response)
         return response
diff --git a/apps/authentication/signal_handlers.py b/apps/authentication/signal_handlers.py
index 0d2a617f9..ac155dcf0 100644
--- a/apps/authentication/signal_handlers.py
+++ b/apps/authentication/signal_handlers.py
@@ -35,6 +35,9 @@ def on_user_auth_login_success(sender, user, request, **kwargs):
             session.delete()
         cache.set(lock_key, request.session.session_key, None)
 
+    # 标记登录，设置 cookie，前端可以控制刷新, Middleware 会拦截这个生成 cookie
+    request.session['auth_session_expiration_required'] = 1
+
 
 @receiver(openid_user_login_success)
 def on_oidc_user_login_success(sender, request, user, create=False, **kwargs):