perf: 优化用户session 会话过期

apps/audits/api.py

@@ -20,6 +20,7 @@ from common.const.http import GET, POST
 from common.drf.filters import DatetimeRangeFilterBackend
 from common.permissions import IsServiceAccount
 from common.plugins.es import QuerySet as ESQuerySet
+from common.sessions.cache import user_session_manager
 from common.storage.ftp_file import FTPFileStorageHandler
 from common.utils import is_uuid, get_logger, lazyproperty
 from orgs.mixins.api import OrgReadonlyModelViewSet, OrgModelViewSet
@@ -289,8 +290,7 @@ class UserSessionViewSet(CommonApiMixin, viewsets.ModelViewSet):
             return Response(status=status.HTTP_200_OK)
 
         keys = queryset.values_list('key', flat=True)
-        session_store_cls = import_module(settings.SESSION_ENGINE).SessionStore
         for key in keys:
-            session_store_cls(key).delete()
+            user_session_manager.decrement_or_remove(key)
         queryset.delete()
         return Response(status=status.HTTP_200_OK)

apps/audits/filters.py

@@ -1,10 +1,9 @@
-from django.core.cache import cache
 from django_filters import rest_framework as drf_filters
 from rest_framework import filters
 from rest_framework.compat import coreapi, coreschema
 
 from common.drf.filters import BaseFilterSet
-from notifications.ws import WS_SESSION_KEY
+from common.sessions.cache import user_session_manager
 from orgs.utils import current_org
 from .models import UserSession
 
@@ -41,13 +40,11 @@ class UserSessionFilterSet(BaseFilterSet):
 
     @staticmethod
     def filter_is_active(queryset, name, is_active):
-        redis_client = cache.client.get_client()
-        members = redis_client.smembers(WS_SESSION_KEY)
-        members = [member.decode('utf-8') for member in members]
+        keys = user_session_manager.get_active_keys()
         if is_active:
-            queryset = queryset.filter(key__in=members)
+            queryset = queryset.filter(key__in=keys)
         else:
-            queryset = queryset.exclude(key__in=members)
+            queryset = queryset.exclude(key__in=keys)
         return queryset
 
     class Meta:

apps/audits/models.py

@@ -4,15 +4,15 @@ from datetime import timedelta
 from importlib import import_module
 
 from django.conf import settings
-from django.core.cache import caches, cache
+from django.core.cache import caches
 from django.db import models
 from django.db.models import Q
 from django.utils import timezone
 from django.utils.translation import gettext, gettext_lazy as _
 
 from common.db.encoder import ModelJSONFieldEncoder
+from common.sessions.cache import user_session_manager
 from common.utils import lazyproperty, i18n_trans
-from notifications.ws import WS_SESSION_KEY
 from ops.models import JobExecution
 from orgs.mixins.models import OrgModelMixin, Organization
 from orgs.utils import current_org
@@ -278,8 +278,7 @@ class UserSession(models.Model):
 
     @property
     def is_active(self):
-        redis_client = cache.client.get_client()
-        return redis_client.sismember(WS_SESSION_KEY, self.key)
+        return user_session_manager.check_active(self.key)
 
     @property
     def date_expired(self):

apps/authentication/forms.py

@@ -18,7 +18,7 @@ class EncryptedField(forms.CharField):
 
 class UserLoginForm(forms.Form):
     days_auto_login = int(settings.SESSION_COOKIE_AGE / 3600 / 24)
-    disable_days_auto_login = settings.SESSION_EXPIRE_AT_BROWSER_CLOSE_FORCE \
+    disable_days_auto_login = settings.SESSION_EXPIRE_AT_BROWSER_CLOSE \
                               or days_auto_login < 1
 
     username = forms.CharField(

apps/authentication/middleware.py

@@ -142,23 +142,7 @@ class SessionCookieMiddleware(MiddlewareMixin):
             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()
-        expire_timestamp = request.session.get_expiry_date().timestamp()
-        response.set_cookie('jms_session_expire_timestamp', expire_timestamp)
-        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_cookie_session_prefix(request, response)
         self.set_cookie_public_key(request, response)
-        self.set_cookie_session_expire(request, response)
         return response

apps/authentication/signal_handlers.py

@@ -37,9 +37,6 @@ def on_user_auth_login_success(sender, user, request, **kwargs):
             UserSession.objects.filter(key=session_key).delete()
         cache.set(lock_key, request.session.session_key, None)
 
-    # 标记登录，设置 cookie，前端可以控制刷新, Middleware 会拦截这个生成 cookie
-    request.session['auth_session_expiration_required'] = 1
-
 
 @receiver(cas_user_authenticated)
 def on_cas_user_login_success(sender, request, user, **kwargs):

apps/common/sessions/cache.py

@@ -0,0 +1,56 @@
+import re
+
+from django.contrib.sessions.backends.cache import (
+    SessionStore as DjangoSessionStore
+)
+from django.core.cache import cache
+
+from jumpserver.utils import get_current_request
+
+
+class SessionStore(DjangoSessionStore):
+    ignore_urls = [
+        r'^/api/v1/users/profile/'
+    ]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ignore_pattern = re.compile('|'.join(self.ignore_urls))
+
+    def save(self, *args, **kwargs):
+        request = get_current_request()
+        if request is None or not self.ignore_pattern.match(request.path):
+            super().save(*args, **kwargs)
+
+
+class RedisUserSessionManager:
+    JMS_SESSION_KEY = 'jms_session_key'
+
+    def __init__(self):
+        self.client = cache.client.get_client()
+
+    def add_or_increment(self, session_key):
+        self.client.hincrby(self.JMS_SESSION_KEY, session_key, 1)
+
+    def decrement_or_remove(self, session_key):
+        new_count = self.client.hincrby(self.JMS_SESSION_KEY, session_key, -1)
+        if new_count <= 0:
+            self.client.hdel(self.JMS_SESSION_KEY, session_key)
+
+    def check_active(self, session_key):
+        count = self.client.hget(self.JMS_SESSION_KEY, session_key)
+        count = 0 if count is None else int(count.decode('utf-8'))
+        return count > 0
+
+    def get_active_keys(self):
+        session_keys = []
+        for k, v in self.client.hgetall(self.JMS_SESSION_KEY).items():
+            count = int(v.decode('utf-8'))
+            if count <= 0:
+                continue
+            key = k.decode('utf-8')
+            session_keys.append(key)
+        return session_keys
+
+
+user_session_manager = RedisUserSessionManager()

apps/jumpserver/conf.py

@@ -547,7 +547,6 @@ class Config(dict):
         'REFERER_CHECK_ENABLED': False,
         'SESSION_ENGINE': 'cache',
         'SESSION_SAVE_EVERY_REQUEST': True,
-        'SESSION_EXPIRE_AT_BROWSER_CLOSE_FORCE': False,
         'SERVER_REPLAY_STORAGE': {},
         'SECURITY_DATA_CRYPTO_ALGO': None,
         'GMSSL_ENABLED': False,

apps/jumpserver/middleware.py

@@ -66,11 +66,6 @@ class RequestMiddleware:
     def __call__(self, request):
         set_current_request(request)
         response = self.get_response(request)
-        is_request_api = request.path.startswith('/api')
-        if not settings.SESSION_EXPIRE_AT_BROWSER_CLOSE and \
-                not is_request_api:
-            age = request.session.get_expiry_age()
-            request.session.set_expiry(age)
         return response
 
 

apps/jumpserver/settings/base.py

@@ -234,11 +234,9 @@ CSRF_COOKIE_NAME = '{}csrftoken'.format(SESSION_COOKIE_NAME_PREFIX)
 SESSION_COOKIE_NAME = '{}sessionid'.format(SESSION_COOKIE_NAME_PREFIX)
 
 SESSION_COOKIE_AGE = CONFIG.SESSION_COOKIE_AGE
-SESSION_EXPIRE_AT_BROWSER_CLOSE = True
-# 自定义的配置，SESSION_EXPIRE_AT_BROWSER_CLOSE 始终为 True, 下面这个来控制是否强制关闭后过期 cookie
-SESSION_EXPIRE_AT_BROWSER_CLOSE_FORCE = CONFIG.SESSION_EXPIRE_AT_BROWSER_CLOSE_FORCE
 SESSION_SAVE_EVERY_REQUEST = CONFIG.SESSION_SAVE_EVERY_REQUEST
-SESSION_ENGINE = "django.contrib.sessions.backends.{}".format(CONFIG.SESSION_ENGINE)
+SESSION_EXPIRE_AT_BROWSER_CLOSE = CONFIG.SESSION_EXPIRE_AT_BROWSER_CLOSE
+SESSION_ENGINE = "common.sessions.{}".format(CONFIG.SESSION_ENGINE)
 
 MESSAGE_STORAGE = 'django.contrib.messages.storage.cookie.CookieStorage'
 # Database

apps/notifications/ws.py

@@ -1,28 +1,32 @@
 import json
+import time
+from threading import Thread
 
 from channels.generic.websocket import JsonWebsocketConsumer
-from django.core.cache import cache
+from django.conf import settings
 
 from common.db.utils import safe_db_connection
+from common.sessions.cache import user_session_manager
 from common.utils import get_logger
 from .signal_handlers import new_site_msg_chan
 from .site_msg import SiteMessageUtil
 
 logger = get_logger(__name__)
-WS_SESSION_KEY = 'ws_session_key'
 
 
 class SiteMsgWebsocket(JsonWebsocketConsumer):
     sub = None
     refresh_every_seconds = 10
 
+    @property
+    def session(self):
+        return self.scope['session']
+
     def connect(self):
         user = self.scope["user"]
         if user.is_authenticated:
             self.accept()
-            session = self.scope['session']
-            redis_client = cache.client.get_client()
-            redis_client.sadd(WS_SESSION_KEY, session.session_key)
+            user_session_manager.add_or_increment(self.session.session_key)
             self.sub = self.watch_recv_new_site_msg()
         else:
             self.close()
@@ -66,6 +70,32 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         if not self.sub:
             return
         self.sub.unsubscribe()
-        session = self.scope['session']
-        redis_client = cache.client.get_client()
-        redis_client.srem(WS_SESSION_KEY, session.session_key)
+
+        user_session_manager.decrement_or_remove(self.session.session_key)
+        if self.should_delete_session():
+            thread = Thread(target=self.delay_delete_session)
+            thread.start()
+
+    def should_delete_session(self):
+        return (self.session.modified or settings.SESSION_SAVE_EVERY_REQUEST) and \
+            not self.session.is_empty() and \
+            self.session.get_expire_at_browser_close() and \
+            not user_session_manager.check_active(self.session.session_key)
+
+    def delay_delete_session(self):
+        timeout = 3
+        check_interval = 0.5
+
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            time.sleep(check_interval)
+            if user_session_manager.check_active(self.session.session_key):
+                return
+
+        self.delete_session()
+
+    def delete_session(self):
+        try:
+            self.session.delete()
+        except Exception as e:
+            logger.info(f'delete session error: {e}')

config_example.yml

@@ -85,7 +85,7 @@ REDIS_PORT: 6379
 # SECURITY_WATERMARK_ENABLED: False
 
 # 浏览器关闭页面后，会话过期
-# SESSION_EXPIRE_AT_BROWSER_CLOSE_FORCE: False
+# SESSION_EXPIRE_AT_BROWSER_CLOSE: False
 
 # 每次api请求，session续期
 # SESSION_SAVE_EVERY_REQUEST: True