123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268 |
- import datetime
- import time
- from importlib import import_module
-
- from django.conf import settings
- from django.contrib.sessions.backends.base import UpdateError
- from django.core.exceptions import SuspiciousOperation
- from django.http import parse_cookie
- from django.http.cookie import SimpleCookie
- from django.utils import timezone
- from django.utils.encoding import force_str
- from django.utils.functional import LazyObject
-
- from channels.db import database_sync_to_async
-
- try:
- from django.utils.http import http_date
- except ImportError:
- from django.utils.http import cookie_date as http_date
-
-
- class CookieMiddleware:
- """
- Extracts cookies from HTTP or WebSocket-style scopes and adds them as a
- scope["cookies"] entry with the same format as Django's request.COOKIES.
- """
-
- def __init__(self, inner):
- self.inner = inner
-
- async def __call__(self, scope, receive, send):
- # Check this actually has headers. They're a required scope key for HTTP and WS.
- if "headers" not in scope:
- raise ValueError(
- "CookieMiddleware was passed a scope that did not have a headers key "
- + "(make sure it is only passed HTTP or WebSocket connections)"
- )
- # Go through headers to find the cookie one
- for name, value in scope.get("headers", []):
- if name == b"cookie":
- cookies = parse_cookie(value.decode("latin1"))
- break
- else:
- # No cookie header found - add an empty default.
- cookies = {}
- # Return inner application
- return await self.inner(dict(scope, cookies=cookies), receive, send)
-
- @classmethod
- def set_cookie(
- cls,
- message,
- key,
- value="",
- max_age=None,
- expires=None,
- path="/",
- domain=None,
- secure=False,
- httponly=False,
- samesite="lax",
- ):
- """
- Sets a cookie in the passed HTTP response message.
-
- ``expires`` can be:
- - a string in the correct format,
- - a naive ``datetime.datetime`` object in UTC,
- - an aware ``datetime.datetime`` object in any time zone.
- If it is a ``datetime.datetime`` object then ``max_age`` will be calculated.
- """
- value = force_str(value)
- cookies = SimpleCookie()
- cookies[key] = value
- if expires is not None:
- if isinstance(expires, datetime.datetime):
- if timezone.is_aware(expires):
- expires = timezone.make_naive(expires, timezone.utc)
- delta = expires - expires.utcnow()
- # Add one second so the date matches exactly (a fraction of
- # time gets lost between converting to a timedelta and
- # then the date string).
- delta = delta + datetime.timedelta(seconds=1)
- # Just set max_age - the max_age logic will set expires.
- expires = None
- max_age = max(0, delta.days * 86400 + delta.seconds)
- else:
- cookies[key]["expires"] = expires
- else:
- cookies[key]["expires"] = ""
- if max_age is not None:
- cookies[key]["max-age"] = max_age
- # IE requires expires, so set it if hasn't been already.
- if not expires:
- cookies[key]["expires"] = http_date(time.time() + max_age)
- if path is not None:
- cookies[key]["path"] = path
- if domain is not None:
- cookies[key]["domain"] = domain
- if secure:
- cookies[key]["secure"] = True
- if httponly:
- cookies[key]["httponly"] = True
- if samesite is not None:
- assert samesite.lower() in [
- "strict",
- "lax",
- "none",
- ], "samesite must be either 'strict', 'lax' or 'none'"
- cookies[key]["samesite"] = samesite
- # Write out the cookies to the response
- for c in cookies.values():
- message.setdefault("headers", []).append(
- (b"Set-Cookie", bytes(c.output(header=""), encoding="utf-8"))
- )
-
- @classmethod
- def delete_cookie(cls, message, key, path="/", domain=None):
- """
- Deletes a cookie in a response.
- """
- return cls.set_cookie(
- message,
- key,
- max_age=0,
- path=path,
- domain=domain,
- expires="Thu, 01-Jan-1970 00:00:00 GMT",
- )
-
-
- class InstanceSessionWrapper:
- """
- Populates the session in application instance scope, and wraps send to save
- the session.
- """
-
- # Message types that trigger a session save if it's modified
- save_message_types = ["http.response.start"]
-
- # Message types that can carry session cookies back
- cookie_response_message_types = ["http.response.start"]
-
- def __init__(self, scope, send):
- self.cookie_name = settings.SESSION_COOKIE_NAME
- self.session_store = import_module(settings.SESSION_ENGINE).SessionStore
-
- self.scope = dict(scope)
-
- if "session" in self.scope:
- # There's already session middleware of some kind above us, pass
- # that through
- self.activated = False
- else:
- # Make sure there are cookies in the scope
- if "cookies" not in self.scope:
- raise ValueError(
- "No cookies in scope - SessionMiddleware needs to run "
- "inside of CookieMiddleware."
- )
- # Parse the headers in the scope into cookies
- self.scope["session"] = LazyObject()
- self.activated = True
-
- # Override send
- self.real_send = send
-
- async def resolve_session(self):
- session_key = self.scope["cookies"].get(self.cookie_name)
- self.scope["session"]._wrapped = await database_sync_to_async(
- self.session_store
- )(session_key)
-
- async def send(self, message):
- """
- Overridden send that also does session saves/cookies.
- """
- # Only save session if we're the outermost session middleware
- if self.activated:
- modified = self.scope["session"].modified
- empty = self.scope["session"].is_empty()
- # If this is a message type that we want to save on, and there's
- # changed data, save it. We also save if it's empty as we might
- # not be able to send a cookie-delete along with this message.
- if (
- message["type"] in self.save_message_types
- and message.get("status", 200) != 500
- and (modified or settings.SESSION_SAVE_EVERY_REQUEST)
- ):
- await database_sync_to_async(self.save_session)()
- # If this is a message type that can transport cookies back to the
- # client, then do so.
- if message["type"] in self.cookie_response_message_types:
- if empty:
- # Delete cookie if it's set
- if settings.SESSION_COOKIE_NAME in self.scope["cookies"]:
- CookieMiddleware.delete_cookie(
- message,
- settings.SESSION_COOKIE_NAME,
- path=settings.SESSION_COOKIE_PATH,
- domain=settings.SESSION_COOKIE_DOMAIN,
- )
- else:
- # Get the expiry data
- if self.scope["session"].get_expire_at_browser_close():
- max_age = None
- expires = None
- else:
- max_age = self.scope["session"].get_expiry_age()
- expires_time = time.time() + max_age
- expires = http_date(expires_time)
- # Set the cookie
- CookieMiddleware.set_cookie(
- message,
- self.cookie_name,
- self.scope["session"].session_key,
- max_age=max_age,
- expires=expires,
- domain=settings.SESSION_COOKIE_DOMAIN,
- path=settings.SESSION_COOKIE_PATH,
- secure=settings.SESSION_COOKIE_SECURE or None,
- httponly=settings.SESSION_COOKIE_HTTPONLY or None,
- samesite=settings.SESSION_COOKIE_SAMESITE,
- )
- # Pass up the send
- return await self.real_send(message)
-
- def save_session(self):
- """
- Saves the current session.
- """
- try:
- self.scope["session"].save()
- except UpdateError:
- raise SuspiciousOperation(
- "The request's session was deleted before the "
- "request completed. The user may have logged "
- "out in a concurrent request, for example."
- )
-
-
- class SessionMiddleware:
- """
- Class that adds Django sessions (from HTTP cookies) to the
- scope. Works with HTTP or WebSocket protocol types (or anything that
- provides a "headers" entry in the scope).
-
- Requires the CookieMiddleware to be higher up in the stack.
- """
-
- def __init__(self, inner):
- self.inner = inner
-
- async def __call__(self, scope, receive, send):
- """
- Instantiate a session wrapper for this scope, resolve the session and
- call the inner application.
- """
- wrapper = InstanceSessionWrapper(scope, send)
-
- await wrapper.resolve_session()
-
- return await self.inner(wrapper.scope, receive, wrapper.send)
-
-
- # Shortcut to include cookie middleware
- def SessionMiddlewareStack(inner):
- return CookieMiddleware(SessionMiddleware(inner))
|