diff --git a/max/exceptions.py b/max/exceptions.py index 8d94784..708e4c4 100644 --- a/max/exceptions.py +++ b/max/exceptions.py @@ -10,6 +10,10 @@ class NotFound(MaxException): pass +class SessionWithoutUser(MaxException): + pass + + class PermissionDenied(MaxException): def __init__(self, message): super().__init__() diff --git a/max/sessions.py b/max/sessions.py index f8c9732..ac30b1d 100644 --- a/max/sessions.py +++ b/max/sessions.py @@ -6,6 +6,7 @@ from itsdangerous import BadSignature, Signer, want_bytes from werkzeug.datastructures import CallbackDict from max.crypto import get_random_string +from max.exceptions import SessionWithoutUser # in seconds SESSION_IDLE_TIMEOUT = 30 * 60 @@ -16,48 +17,42 @@ CREATED_AT_KEY = "_created_at" LAST_SEEN_KEY = "_last_seen" +SESSION_MAX_PER_USER = 10 + + class SessionStorage: - def __init__(self, maxcap=100): - self.maxcap = maxcap - self.lru_ids = [] # newest ids are on the end (so do append to add a fresh id) + def __init__(self): self.data = {} def __getitem__(self, id): - if id not in self.data: - return None - - self._refresh_id(id) - return self.data[id] + return self.data.get(id) def __delitem__(self, id): - if id in self.data: - del self.data[id] - self._remove_id(id) + self.data.pop(id, None) def __setitem__(self, id, data): - self._refresh_id(id) + if data.get("user_id") is None: + raise SessionWithoutUser("Refusing to store a session without user_id") + + if id not in self.data: + self._limit_user_sessions(data["user_id"]) self.data[id] = data - while len(self.lru_ids) > self.maxcap: - removed_id = self.lru_ids.pop(0) # remove first element - del self.data[removed_id] - - def _remove_id(self, id): - self.lru_ids = [_id for _id in self.lru_ids if _id != id] - - def _refresh_id(self, id): - self._remove_id(id) - self.lru_ids.append(id) + def _limit_user_sessions(self, user_id): + """Makes room for one new session by removing the least recently used ones""" + user_ids = [id for id, data in self.data.items() if data["user_id"] == user_id] + user_ids.sort(key=lambda id: self.data[id][LAST_SEEN_KEY]) + for id in user_ids[: max(0, len(user_ids) - SESSION_MAX_PER_USER + 1)]: + del self.data[id] def remove_user_sessions(self, user_id, except_id=None): ids_to_remove = [] for id, data in self.data.items(): - if data.get("user_id", None) == user_id and id != except_id: + if data["user_id"] == user_id and id != except_id: ids_to_remove.append(id) for id in ids_to_remove: del self.data[id] - self._remove_id(id) session_storage = SessionStorage() @@ -161,11 +156,16 @@ class SessionInterface(FlaskSessionInterface): session.last_seen = time.monotonic() if session.created_at is None: session.created_at = session.last_seen - session_storage[session.id] = { - **session, - CREATED_AT_KEY: session.created_at, - LAST_SEEN_KEY: session.last_seen, - } + try: + session_storage[session.id] = { + **session, + CREATED_AT_KEY: session.created_at, + LAST_SEEN_KEY: session.last_seen, + } + except SessionWithoutUser: + # Flask saves the session again for the error response + session.clear() + raise session_id = ( self._get_signer(app).sign(want_bytes(session.id)).decode("utf-8") )