From 1c241d268f2a4370b4cc1789fff86e4368b56f41 Mon Sep 17 00:00:00 2001
From: blallo <blallo@autistici.org>
Date: Thu, 5 Sep 2019 17:57:09 +0200
Subject: [PATCH] Improve session management in rest logic.

---
 api/rest.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/api/rest.py b/api/rest.py
index f8118d9..6b1c41e 100644
--- a/api/rest.py
+++ b/api/rest.py
@@ -11,8 +11,9 @@ import logging
 import os
 import pkg_resources
 import typing as T
+import weakref
 
-from aiohttp_session import get_session, Session
+from aiohttp_session import new_session, get_session, Session
 from passlib.hash import bcrypt
 
 from bot_z.async_operator import AsyncOperator, push_to_loop
@@ -21,7 +22,9 @@ from api import BASE_URI, DEBUG
 
 alog = logging.getLogger("api")
 routes = web.RouteTableDef()
-OPERATORS = {}  # type: T.Dict[T.Text, AsyncOperator]
+OPERATORS = weakref.WeakKeyDictionary(
+    {}
+)  # type: weakref.WeakKeyDictionary[T.Text, AsyncOperator]
 BASE_PATH = pkg_resources.resource_filename(__name__, "assets")
 EXECUTOR = ProcessPoolExecutor()
 # WARN: the default il 12 rounds; both the server and the client shall compute
@@ -54,13 +57,17 @@ async def get_set_operator(
     op = None
     if "async_operator" in session:
         op = OPERATORS.get(session["async_operator"])
-    base_uri = request.app["base_uri"]
-    debug = request.app["debug"]
-    headless = request.app["headless"]
-    if op is None:
+    else:
+        session = await new_session(request)
+
+    if op is None or session.new:
+        base_uri = request.app["base_uri"]
+        debug = request.app["debug"]
+        headless = request.app["headless"]
         op = AsyncOperator(base_uri, name=user, headless=headless, debug=debug)
         session["async_operator"] = user
         OPERATORS[user] = op
+
     return op, session
 
 
@@ -135,6 +142,7 @@ async def logout_handler(request: web.Request) -> web.Response:
             {"error": "No session", "logged_in": False}, status=401
         )
     res = await logout(op)
+    session.invalidate()
     alog.debug("logout result: %s", res)
     # FIX: assess if better to invalidate session and dump the browser instance.
     del OPERATORS[op_key]