1from __future__ import annotations
 2
 3import time
 4
 5from opentelemetry import trace
 6from opentelemetry.semconv._incubating.attributes.session_attributes import SESSION_ID
 7
 8from plain.http import HttpMiddleware, Request, Response
 9from plain.runtime import settings
10from plain.utils.cache import patch_vary_headers
11from plain.utils.http import http_date
12
13from .core import SessionStore
14from .requests import get_request_session, set_request_session
15
16__all__ = ["SessionMiddleware"]
17
18
19class SessionMiddleware(HttpMiddleware):
20    def before_request(self, request: Request) -> Response | None:
21        session_key = request.cookies.get(settings.SESSION_COOKIE_NAME)
22
23        session = SessionStore(session_key)
24        set_request_session(request, session)
25
26        if session.model_instance:
27            trace.get_current_span().set_attribute(
28                SESSION_ID, session.model_instance.id
29            )
30
31        return None
32
33    def after_response(self, request: Request, response: Response) -> Response:
34        """
35        If request.session was modified, or if the configuration is to save the
36        session every time, save the changes and set a session cookie or delete
37        the session cookie if the session has been emptied.
38        """
39        session = get_request_session(request)
40        accessed = session.accessed
41        modified = session.modified
42        empty = session.is_empty()
43
44        # First check if we need to delete this cookie.
45        # The session should be deleted only if the session is entirely empty.
46        if settings.SESSION_COOKIE_NAME in request.cookies and empty:
47            response.delete_cookie(
48                settings.SESSION_COOKIE_NAME,
49                path=settings.SESSION_COOKIE_PATH,
50                domain=settings.SESSION_COOKIE_DOMAIN,
51                samesite=settings.SESSION_COOKIE_SAMESITE,
52            )
53            patch_vary_headers(response, ["Cookie"])
54        else:
55            if accessed:
56                patch_vary_headers(response, ["Cookie"])
57            if (modified or settings.SESSION_SAVE_EVERY_REQUEST) and not empty:
58                if settings.SESSION_EXPIRE_AT_BROWSER_CLOSE:
59                    max_age = None
60                    expires = None
61                else:
62                    max_age = settings.SESSION_COOKIE_AGE
63                    expires_time = time.time() + max_age
64                    expires = http_date(expires_time)
65                # Save the session data and refresh the client cookie.
66                # Skip session save for 5xx responses.
67                if response.status_code < 500:
68                    session.save()
69                    # session_key must exist after save()
70                    assert session.session_key is not None, (
71                        "Session key should exist after save()"
72                    )
73                    response.set_cookie(
74                        settings.SESSION_COOKIE_NAME,
75                        session.session_key,
76                        max_age=max_age,
77                        expires=expires,
78                        domain=settings.SESSION_COOKIE_DOMAIN,
79                        path=settings.SESSION_COOKIE_PATH,
80                        secure=bool(settings.SESSION_COOKIE_SECURE),
81                        httponly=bool(settings.SESSION_COOKIE_HTTPONLY),
82                        samesite=settings.SESSION_COOKIE_SAMESITE,
83                    )
84        return response