From 4b33819004d8e3c7c377d0ace7625c47034e2907 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20P=C3=A9ters?= <fpeters@entrouvert.com>
Date: Thu, 8 Dec 2022 20:32:52 +0100
Subject: [PATCH] misc: use database to store form tokens (#71455)

---
 tests/test_sessions.py  |  8 ++++---
 tests/test_sql.py       | 19 +++++++++++++++++
 wcs/qommon/publisher.py | 19 -----------------
 wcs/qommon/sessions.py  | 46 +----------------------------------------
 wcs/sql.py              | 29 ++++++++++++++++++++++----
 5 files changed, 50 insertions(+), 71 deletions(-)

diff --git a/tests/test_sessions.py b/tests/test_sessions.py
index c187206fc..9bf445aa2 100644
--- a/tests/test_sessions.py
+++ b/tests/test_sessions.py
@@ -239,14 +239,16 @@ def test_transient_data_removal(pub, app):
     resp = resp.form.submit('submit')
 
     assert sql.Session.count() == 1
-    assert sql.TransientData.count() == 1
-    transient_data = sql.TransientData.select()[0]
+    assert sql.TransientData.count() == 2
+    transient_data1, transient_data2 = sql.TransientData.select(order_by='last_update_time')
+    assert transient_data1.data is None  # form_token
+    assert transient_data2.data == {'1': 'test'}  # magictoken
 
     app.get('/logout')
     assert sql.Session.count() == 0
     assert sql.TransientData.count() == 0
 
-    transient_data.store()  # session_id not found, should not fail
+    transient_data2.store()  # session_id not found, should not fail
 
 
 def test_magictoken_migration(pub, app):
diff --git a/tests/test_sql.py b/tests/test_sql.py
index 526dcfba9..af55cf1bc 100644
--- a/tests/test_sql.py
+++ b/tests/test_sql.py
@@ -1,5 +1,6 @@
 import datetime
 import io
+import os
 import pickle
 import random
 import string
@@ -2449,3 +2450,21 @@ def test_python_datasource_migration(pub):
     # check it's no longer needed afterwards
     sql.migrate()
     assert sql.is_reindex_needed('python_ds_migration', conn=conn, cur=cur) is False
+
+
+def test_form_tokens_migration(pub):
+    conn, cur = sql.get_connection_and_cursor()
+    cur.execute('UPDATE wcs_meta SET value = 70 WHERE key = %s', ('sql_level',))
+    conn.commit()
+    cur.close()
+
+    form_tokens_dir = os.path.join(pub.app_dir, 'form_tokens')
+    if not os.path.exists(form_tokens_dir):
+        os.mkdir(form_tokens_dir)
+    with open(os.path.join(form_tokens_dir, '1234'), 'w'):
+        pass
+
+    assert os.path.exists(os.path.join(form_tokens_dir, '1234'))
+    sql.migrate()
+    assert not os.path.exists(os.path.join(form_tokens_dir, '1234'))
+    assert not os.path.exists(form_tokens_dir)
diff --git a/wcs/qommon/publisher.py b/wcs/qommon/publisher.py
index f63edd84b..6045b74f4 100644
--- a/wcs/qommon/publisher.py
+++ b/wcs/qommon/publisher.py
@@ -130,10 +130,6 @@ class QommonPublisher(Publisher):
     app_dir = None
     _i18n_catalog = None
 
-    @property
-    def form_tokens_dir(self):
-        return os.path.join(self.app_dir, 'form_tokens')
-
     def get_root_url(self):
         if self.get_request():
             return self.get_request().environ['SCRIPT_NAME'] + '/'
@@ -478,11 +474,6 @@ class QommonPublisher(Publisher):
                 raise ImmediateRedirectException(self.missing_appdir_redirect)
             raise Http404()
 
-        try:
-            os.mkdir(self.form_tokens_dir)
-        except OSError:  # already exists
-            pass
-
     def init_publish(self, request):
         self.set_app_dir(request)
 
@@ -614,16 +605,6 @@ class QommonPublisher(Publisher):
                         except KeyError:
                             pass
                         continue
-                # also delete obsolete form_tokens that would have be missed when
-                # cleaning sessions.
-                form_tokens_dir = self.form_tokens_dir
-                if os.path.exists(form_tokens_dir):
-                    for filename in os.listdir(form_tokens_dir):
-                        try:
-                            if os.stat(os.path.join(form_tokens_dir, filename)).st_mtime < creation_limit:
-                                os.unlink(os.path.join(form_tokens_dir, filename))
-                        except FileNotFoundError:
-                            pass
         except locket.LockError:
             pass
 
diff --git a/wcs/qommon/sessions.py b/wcs/qommon/sessions.py
index 544fe090d..97446720b 100644
--- a/wcs/qommon/sessions.py
+++ b/wcs/qommon/sessions.py
@@ -166,40 +166,6 @@ class Session(QommonSession, CaptchaSession, StorableObject):
 
     session_id = property(get_session_id, set_session_id)
 
-    def get_form_token_filepath(self, token):
-        return os.path.join(get_publisher().form_tokens_dir, token)
-
-    def create_form_token(self):
-        token = super().create_form_token()
-        with open(self.get_form_token_filepath(token), 'wb'):
-            # create empty file
-            pass
-        return token
-
-    def has_form_token(self, token):
-        if not token:
-            return False
-        has_form_token = super().has_form_token(token)
-        if not os.path.exists(self.get_form_token_filepath(token)):
-            has_form_token = False
-        return has_form_token
-
-    def remove_form_token(self, token):
-        super().remove_form_token(token)
-        self.store()
-        try:
-            os.unlink(self.get_form_token_filepath(token))
-        except OSError:
-            pass
-
-    def clean_form_tokens(self):
-        dirname = os.path.join(get_publisher().app_dir, 'form_tokens')
-        for token in self._form_tokens:
-            try:
-                os.unlink(os.path.join(dirname, token))
-            except OSError:
-                pass
-
     def has_user(self):
         user_id = QuixoteSession.get_user(self)
         return bool(user_id)
@@ -400,17 +366,7 @@ class StorageSessionManager(QommonSessionManager):
     def __delitem__(self, session_id):
         if not session_id:
             return
-        try:
-            session = self.session_class.get(session_id)
-        except KeyError:
-            session = None
-        try:
-            self.session_class.remove_object(session_id)
-        except OSError:
-            raise KeyError
-
-        if session:
-            session.clean_form_tokens()
+        self.session_class.remove_object(session_id)
 
     def get_sessions_for_saml(self, name_identifier=Ellipsis, session_indexes=()):
         return self.session_class.get_sessions_for_saml(name_identifier, session_indexes)
diff --git a/wcs/sql.py b/wcs/sql.py
index 317c2414f..08fd0671f 100644
--- a/wcs/sql.py
+++ b/wcs/sql.py
@@ -19,7 +19,10 @@ import datetime
 import io
 import itertools
 import json
+import os
 import re
+import secrets
+import shutil
 import time
 
 import psycopg2
@@ -3518,7 +3521,7 @@ class Role(SqlMixin, wcs.roles.Role):
 
 
 class TransientData(SqlMixin):
-    # table to keep some transient submission data out of global session dictionary
+    # table to keep some transient submission data and form tokens out of session object
     _table_name = 'transient_data'
     _table_static_fields = [
         ('id', 'varchar'),
@@ -3538,7 +3541,7 @@ class TransientData(SqlMixin):
         sql_dict = {
             'id': self.id,
             'session_id': self.session_id,
-            'data': bytearray(pickle.dumps(self.data, protocol=2)),
+            'data': bytearray(pickle.dumps(self.data, protocol=2)) if self.data is not None else None,
             'last_update_time': now(),
         }
 
@@ -3565,7 +3568,7 @@ class TransientData(SqlMixin):
         o = cls.__new__(cls)
         o.id = str_encode(row[0])
         o.session_id = row[1]
-        o.data = pickle_loads(row[2])
+        o.data = pickle_loads(row[2]) if row[2] else None
         return o
 
     @classmethod
@@ -3720,6 +3723,17 @@ class Session(SqlMixin, wcs.sessions.BasicSession):
         super().remove_magictoken(token)
         TransientData.remove_object(token)
 
+    def create_form_token(self):
+        token = TransientData(id=secrets.token_urlsafe(16), session_id=self.id, data=None)
+        token.store()
+        return token.id
+
+    def has_form_token(self, token):
+        return TransientData.exists([Equal('id', token)])
+
+    def remove_form_token(self, token):
+        TransientData.remove_object(token)
+
 
 class TrackingCode(SqlMixin, wcs.tracking_code.TrackingCode):
     _table_name = 'tracking_codes'
@@ -4660,7 +4674,7 @@ def get_period_total(
 # latest migration, number + description (description is not used
 # programmaticaly but will make sure git conflicts if two migrations are
 # separately added with the same number)
-SQL_LEVEL = (71, 'python datasource migration')
+SQL_LEVEL = (72, 'form tokens to db')
 
 
 def migrate_global_views(conn, cur):
@@ -4948,6 +4962,13 @@ def migrate():
         # 71: python datasource migration
         set_reindex('python_ds_migration', 'needed', conn=conn, cur=cur)
 
+    if sql_level < 72:
+        # 72: form tokens to db
+        # it uses the existing tokens table, this "migration" is just to remove old files.
+        form_tokens_dir = os.path.join(get_publisher().app_dir, 'form_tokens')
+        if os.path.exists(form_tokens_dir):
+            shutil.rmtree(form_tokens_dir, ignore_errors=True)
+
     if sql_level != SQL_LEVEL[0]:
         cur.execute(
             '''UPDATE wcs_meta SET value = %s, updated_at=NOW() WHERE key = %s''',
-- 
2.39.0

