From 8cb74976e191264a2d9bd97680ba065436b1b86b Mon Sep 17 00:00:00 2001 From: feng <1304903146@qq.com> Date: Wed, 31 Jan 2024 17:17:34 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E7=94=A8=E6=88=B7ses?= =?UTF-8?q?sion=20=E4=BC=9A=E8=AF=9D=E8=BF=87=E6=9C=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/audits/api.py | 4 +- apps/audits/filters.py | 11 ++--- apps/audits/models.py | 7 ++-- apps/authentication/forms.py | 2 +- apps/authentication/middleware.py | 16 -------- apps/authentication/signal_handlers.py | 3 -- apps/common/sessions/__init__.py | 0 apps/common/sessions/cache.py | 56 ++++++++++++++++++++++++++ apps/jumpserver/conf.py | 1 - apps/jumpserver/middleware.py | 5 --- apps/jumpserver/settings/base.py | 6 +-- apps/notifications/ws.py | 46 +++++++++++++++++---- config_example.yml | 2 +- 13 files changed, 107 insertions(+), 52 deletions(-) create mode 100644 apps/common/sessions/__init__.py create mode 100644 apps/common/sessions/cache.py diff --git a/apps/audits/api.py b/apps/audits/api.py index 209bc30fd..5a665777b 100644 --- a/apps/audits/api.py +++ b/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) diff --git a/apps/audits/filters.py b/apps/audits/filters.py index d24cf72fd..078e68c2c 100644 --- a/apps/audits/filters.py +++ b/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: diff --git a/apps/audits/models.py b/apps/audits/models.py index 64a0ebc5b..c4b4486a5 100644 --- a/apps/audits/models.py +++ b/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): diff --git a/apps/authentication/forms.py b/apps/authentication/forms.py index 90429bc82..5f56a4a98 100644 --- a/apps/authentication/forms.py +++ b/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( diff --git a/apps/authentication/middleware.py b/apps/authentication/middleware.py index eabb90263..e6182e9bb 100644 --- a/apps/authentication/middleware.py +++ b/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 diff --git a/apps/authentication/signal_handlers.py b/apps/authentication/signal_handlers.py index 3ac92411f..943d751dd 100644 --- a/apps/authentication/signal_handlers.py +++ b/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): diff --git a/apps/common/sessions/__init__.py b/apps/common/sessions/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/common/sessions/cache.py b/apps/common/sessions/cache.py new file mode 100644 index 000000000..805d6738a --- /dev/null +++ b/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() diff --git a/apps/jumpserver/conf.py b/apps/jumpserver/conf.py index d2f99e85f..9f869ace6 100644 --- a/apps/jumpserver/conf.py +++ b/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, diff --git a/apps/jumpserver/middleware.py b/apps/jumpserver/middleware.py index 9a49189c8..ad6bb496d 100644 --- a/apps/jumpserver/middleware.py +++ b/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 diff --git a/apps/jumpserver/settings/base.py b/apps/jumpserver/settings/base.py index 10188a1a7..f05293155 100644 --- a/apps/jumpserver/settings/base.py +++ b/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 diff --git a/apps/notifications/ws.py b/apps/notifications/ws.py index 019b0b409..b172c4ed4 100644 --- a/apps/notifications/ws.py +++ b/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}') diff --git a/config_example.yml b/config_example.yml index 03cc0aa03..232ac32d8 100644 --- a/config_example.yml +++ b/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