Implement MSC3231: Token authenticated registration

synapse/rest/admin/registration_tokens.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import logging
-import random
 import string
 from typing import TYPE_CHECKING, Tuple
 
@@ -153,7 +152,9 @@ async def on_POST(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
                 )
 
             # Generate token
-            token = "".join(random.choices(self.allowed_chars, k=length))
+            token = await self.store.generate_registration_token(
+                length, self.allowed_chars
+            )
 
         uses_allowed = body.get("uses_allowed", None)
         if not (
@@ -179,28 +180,7 @@ async def on_POST(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
         created = await self.store.create_registration_token(
             token, uses_allowed, expiry_time
         )
-
-        if "token" not in body:
-            # The token was generated. If it could not be created because
-            # that token already exists, then try a few more times before
-            # reporting a failure.
-            i = 0
-            while not created and i < 3:
-                token = "".join(random.choices(self.allowed_chars, k=length))
-                created = await self.store.create_registration_token(
-                    token, uses_allowed, expiry_time
-                )
-                i += 1
-            if not created:
-                raise SynapseError(
-                    500,
-                    "The generated token already exists. Try again with a greater length.",
-                    Codes.UNKNOWN,
-                )
-
-        elif not created:
-            # The token was specified in the request, but it already exists
-            # so could not be created.
+        if not created:
             raise SynapseError(
                 400, f"Token already exists: {token}", Codes.INVALID_PARAM
             )

synapse/storage/databases/main/registration.py

@@ -1326,6 +1326,47 @@ async def get_one_registration_token(self, token: str) -> Optional[Dict[str, Any
             desc="get_one_registration_token",
         )
 
+    async def generate_registration_token(
+        self, length: int, chars: str
+    ) -> Optional[str]:
+        """Generate a random registration token. Used by the admin API.
+
+        Args:
+            length: The length of the token to generate.
+            chars: A string of the characters allowed in the generated token.
+
+        Returns:
+            The generated token.
+
+        Raises:
+            SynapseError if a unique registration token could still not be
+            generated after a few tries.
+        """
+        # Make a few attempts at generating a unique token of the required
+        # length before failing.
+        for _i in range(3):
+            # Generate token
+            token = "".join(random.choices(chars, k=length))
+
+            # Check if the token already exists
+            existing_token = await self.db_pool.simple_select_one_onecol(
+                "registration_tokens",
+                keyvalues={"token": token},
+                retcol="token",
+                allow_none=True,
+                desc="check_if_registration_token_exists",
+            )
+
+            if existing_token is None:
+                # The generated token doesn't exist yet, return it
+                return token
+
+        raise SynapseError(
+            500,
+            "Unable to generate a registration token. Try again with a greater length",
+            Codes.UNKNOWN,
+        )
+
     async def create_registration_token(
         self, token: str, uses_allowed: Optional[int], expiry_time: Optional[int]
     ) -> bool:

tests/rest/admin/test_registration_tokens.py

@@ -189,6 +189,37 @@ def test_create_token_already_exists(self):
         self.assertEqual(400, int(channel2.result["code"]), msg=channel2.result["body"])
         self.assertEqual(channel2.json_body["errcode"], Codes.INVALID_PARAM)
 
+    def test_create_unable_to_generate_token(self):
+        """Check right error is raised when server can't generate unique token."""
+        # Create all possible single character tokens
+        tokens = []
+        for c in string.ascii_letters + string.digits + "-_":
+            tokens.append(
+                {
+                    "token": c,
+                    "uses_allowed": None,
+                    "pending": 0,
+                    "completed": 0,
+                    "expiry_time": None,
+                }
+            )
+        self.get_success(
+            self.store.db_pool.simple_insert_many(
+                "registration_tokens",
+                tokens,
+                "create_all_registration_tokens",
+            )
+        )
+
+        # Check creating a single character token fails with a 500 status code
+        channel = self.make_request(
+            "POST",
+            self.url + "/new",
+            {"length": 1},
+            access_token=self.admin_user_tok,
+        )
+        self.assertEqual(500, int(channel.result["code"]), msg=channel.result["body"])
+
     def test_create_uses_allowed(self):
         """Check you can only create a token with good values for uses_allowed."""
         # Should work with 0 (token is invalid from the start)