Add storage support for checking and updating a user's approval status

Author: babolivier

synapse/storage/databases/main/registration.py

@@ -181,6 +181,7 @@ async def get_user_by_id(self, user_id: str) -> Optional[Dict[str, Any]]:
                 "user_type",
                 "deactivated",
                 "shadow_banned",
+                "approved",
             ],
             allow_none=True,
             desc="get_user_by_id",
@@ -1778,6 +1779,26 @@ async def is_guest(self, user_id: str) -> bool:
 
         return res if res else False
 
+    @cached()
+    async def is_user_approved(self, user_id: str) -> bool:
+        """Checks if a user is approved and therefore can be allowed to log in.
+
+        Args:
+            user_id: the user to check the approval status of.
+
+        Returns:
+            A boolean that is True if the user is approved, False otherwise.
+        """
+        ret = await self.db_pool.simple_select_one_onecol(
+            table="users",
+            keyvalues={"name": user_id},
+            retcol="approved",
+            allow_none=True,
+            desc="is_user_pending_approval",
+        )
+
+        return bool(ret)
+
 
 class RegistrationBackgroundUpdateStore(RegistrationWorkerStore):
     def __init__(
@@ -1817,6 +1838,10 @@ def __init__(
             unique=False,
         )
 
+        self.db_pool.updates.register_background_update_handler(
+            "users_set_approved_flag", self._background_update_set_approved_flag
+        )
+
     async def _background_update_set_deactivated_flag(
         self, progress: JsonDict, batch_size: int
     ) -> int:
@@ -1915,6 +1940,75 @@ def set_user_deactivated_status_txn(
         self._invalidate_cache_and_stream(txn, self.get_user_by_id, (user_id,))
         txn.call_after(self.is_guest.invalidate, (user_id,))
 
+    async def _background_update_set_approved_flag(
+        self, progress: JsonDict, batch_size: int
+    ) -> int:
+        """Set the 'approved' flag for all already existing users. We want to set it to
+        true systematically because we don't want to suddenly prevent already existing
+        users from logging in if the option to block registration on approval is turned
+        on.
+        """
+        last_user = progress.get("user_id", "")
+
+        def _background_update_set_approved_flag_txn(txn: LoggingTransaction) -> int:
+            txn.execute(
+                """
+                SELECT name
+                FROM users
+                WHERE
+                    approved IS NULL
+                    AND name > ?
+                ORDER BY name ASC
+                LIMIT ?
+                """,
+                (last_user, batch_size),
+            )
+            rows = self.db_pool.cursor_to_dict(txn)
+
+            if len(rows) == 0:
+                return 0
+
+            for user in rows:
+                self.update_user_approval_status_txn(txn, user["name"], True)
+
+            self.db_pool.updates._background_update_progress_txn(
+                txn, "users_set_approved_flag", {"user_id": rows[-1]["name"]}
+            )
+
+            return len(rows)
+
+        nb_processed = await self.db_pool.runInteraction(
+            "users_set_approved_flag", _background_update_set_approved_flag_txn
+        )
+
+        if nb_processed < batch_size:
+            await self.db_pool.updates._end_background_update("users_set_approved_flag")
+
+        return nb_processed
+
+    def update_user_approval_status_txn(
+        self, txn: LoggingTransaction, user_id: str, approved: bool
+    ) -> None:
+        """Set the user's 'approved' flag to the given value.
+
+        The boolean is turned into an int because the column is a smallint.
+
+        Args:
+            txn: the current database transaction.
+            user_id: the user to update the flag for.
+            approved: the value to set the flag to.
+        """
+        self.db_pool.simple_update_one_txn(
+            txn=txn,
+            table="users",
+            keyvalues={"name": user_id},
+            updatevalues={"approved": int(approved)},
+        )
+
+        # Invalidate the caches of methods that read the value of the 'approved' flag.
+        self._invalidate_cache_and_stream(txn, self.get_user_by_id, (user_id,))
+        self._invalidate_cache_and_stream(txn, self.is_user_approved, (user_id,))
+
 
 class RegistrationStore(StatsStore, RegistrationBackgroundUpdateStore):
     def __init__(
@@ -1932,6 +2026,13 @@ def __init__(
         self._access_tokens_id_gen = IdGenerator(db_conn, "access_tokens", "id")
         self._refresh_tokens_id_gen = IdGenerator(db_conn, "refresh_tokens", "id")
 
+        # If support for MSC3866 is enabled and configured to require approval for new
+        # account, we will create new users with an 'approved' flag set to false.
+        self._require_approval = (
+            hs.config.experimental.msc3866.enabled
+            and hs.config.experimental.msc3866.require_approval_for_new_accounts
+        )
+
     async def add_access_token_to_user(
         self,
         user_id: str,
@@ -2064,6 +2165,7 @@ async def register_user(
         admin: bool = False,
         user_type: Optional[str] = None,
         shadow_banned: bool = False,
+        approved: bool = False,
     ) -> None:
         """Attempts to register an account.
 
@@ -2082,6 +2184,8 @@ async def register_user(
                 or None for a normal user.
             shadow_banned: Whether the user is shadow-banned, i.e. they may be
                 told their requests succeeded but we ignore them.
+            approved: Whether to consider the user has already been approved by an
+                administrator.
 
         Raises:
             StoreError if the user_id could not be registered.
@@ -2098,6 +2202,7 @@ async def register_user(
             admin,
             user_type,
             shadow_banned,
+            approved,
         )
 
     def _register_user(
@@ -2112,11 +2217,14 @@ def _register_user(
         admin: bool,
         user_type: Optional[str],
         shadow_banned: bool,
+        approved: bool,
     ) -> None:
         user_id_obj = UserID.from_string(user_id)
 
         now = int(self._clock.time())
 
+        pending_approval = self._require_approval and not approved
+
         try:
             if was_guest:
                 # Ensure that the guest user actually exists
@@ -2142,6 +2250,7 @@ def _register_user(
                         "admin": 1 if admin else 0,
                         "user_type": user_type,
                         "shadow_banned": shadow_banned,
+                        "approved": 0 if pending_approval else 1,
                     },
                 )
             else:
@@ -2157,6 +2266,7 @@ def _register_user(
                         "admin": 1 if admin else 0,
                         "user_type": user_type,
                         "shadow_banned": shadow_banned,
+                        "approved": 0 if pending_approval else 1,
                     },
                 )
 
@@ -2499,6 +2609,25 @@ def start_or_continue_validation_session_txn(txn: LoggingTransaction) -> None:
             start_or_continue_validation_session_txn,
         )
 
+    async def update_user_approval_status(
+        self, user_id: UserID, approved: bool
+    ) -> None:
+        """Set the user's 'approved' flag to the given value.
+
+        The boolean will be turned into an int (in update_user_approval_status_txn)
+        because the column is a smallint.
+
+        Args:
+            user_id: the user to update the flag for.
+            approved: the value to set the flag to.
+        """
+        await self.db_pool.runInteraction(
+            "update_user_approval_status",
+            self.update_user_approval_status_txn,
+            user_id.to_string(),
+            approved,
+        )
+
 
 def find_max_generated_user_id_localpart(cur: Cursor) -> int:
     """

synapse/storage/schema/main/delta/72/04users_approved_column.sql

@@ -0,0 +1,22 @@
+/* Copyright 2022 The Matrix.org Foundation C.I.C
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+-- Add a column to the users table to track whether the user needs to be approved by an
+-- administrator.
+ALTER TABLE users ADD COLUMN approved SMALLINT;
+
+-- Run a background update to set the approved flag on already existing users.
+INSERT INTO background_updates (ordering, update_name, progress_json) VALUES
+  (7204, 'users_set_approved_flag', '{}');

tests/storage/test_registration.py

@@ -11,11 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from twisted.test.proto_helpers import MemoryReactor
 
 from synapse.api.constants import UserTypes
 from synapse.api.errors import ThreepidValidationError
+from synapse.server import HomeServer
+from synapse.types import JsonDict, UserID
+from synapse.util import Clock
 
-from tests.unittest import HomeserverTestCase
+from tests.unittest import HomeserverTestCase, override_config
 
 
 class RegistrationStoreTestCase(HomeserverTestCase):
@@ -44,6 +48,7 @@ def test_register(self):
                 "user_type": None,
                 "deactivated": 0,
                 "shadow_banned": 0,
+                "approved": 1,
             },
             (self.get_success(self.store.get_user_by_id(self.user_id))),
         )
@@ -147,3 +152,101 @@ def test_3pid_inhibit_invalid_validation_session_error(self):
             ThreepidValidationError,
         )
         self.assertEqual(e.value.msg, "Validation token not found or has expired", e)
+
+
+class ApprovalRequiredRegistrationTestCase(HomeserverTestCase):
+    def default_config(self) -> JsonDict:
+        config = super().default_config()
+
+        # If there's already some config for this feature in the default config, it
+        # means we're overriding it with @override_config. In this case we don't want
+        # to do anything more with it.
+        msc3866_config = config.get("experimental_features", {}).get("msc3866")
+        if msc3866_config is not None:
+            return config
+
+        # Require approval for all new accounts.
+        config["experimental_features"] = {
+            "msc3866": {
+                "enabled": True,
+                "require_approval_for_new_accounts": True,
+            }
+        }
+        return config
+
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+        self.store = hs.get_datastores().main
+        self.user_id = "@my-user:test"
+        self.pwhash = "{xx1}123456789"
+
+    @override_config(
+        {
+            "experimental_features": {
+                "msc3866": {
+                    "enabled": True,
+                    "require_approval_for_new_accounts": False,
+                }
+            }
+        }
+    )
+    def test_approval_not_required(self) -> None:
+        """Tests that if we don't require approval for new accounts, newly created
+        accounts are automatically marked as approved.
+        """
+        self.get_success(self.store.register_user(self.user_id, self.pwhash))
+
+        user = self.get_success(self.store.get_user_by_id(self.user_id))
+        assert user is not None
+        self.assertEqual(user["approved"], 1)
+
+        approved = self.get_success(self.store.is_user_approved(self.user_id))
+        self.assertTrue(approved)
+
+    def test_approval_required(self) -> None:
+        """Tests that if we require approval for new accounts, newly created accounts
+        are not automatically marked as approved.
+        """
+        self.get_success(self.store.register_user(self.user_id, self.pwhash))
+
+        user = self.get_success(self.store.get_user_by_id(self.user_id))
+        assert user is not None
+        self.assertEqual(user["approved"], 0)
+
+        approved = self.get_success(self.store.is_user_approved(self.user_id))
+        self.assertFalse(approved)
+
+    def test_override(self) -> None:
+        """Tests that if we require approval for new accounts, but we explicitly say the
+        new user should be considered approved, they're marked as approved.
+        """
+        self.get_success(
+            self.store.register_user(
+                self.user_id,
+                self.pwhash,
+                approved=True,
+            )
+        )
+
+        user = self.get_success(self.store.get_user_by_id(self.user_id))
+        self.assertIsNotNone(user)
+        assert user is not None
+        self.assertEqual(user["approved"], 1)
+
+        approved = self.get_success(self.store.is_user_approved(self.user_id))
+        self.assertTrue(approved)
+
+    def test_approve_user(self) -> None:
+        """Tests that approving the user updates their approval status."""
+        self.get_success(self.store.register_user(self.user_id, self.pwhash))
+
+        approved = self.get_success(self.store.is_user_approved(self.user_id))
+        self.assertFalse(approved)
+
+        self.get_success(
+            self.store.update_user_approval_status(
+                UserID.from_string(self.user_id), True
+            )
+        )
+
+        approved = self.get_success(self.store.is_user_approved(self.user_id))
+        self.assertTrue(approved)