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