From 1fe3cc2c9c59001a6d3f7b28f81bd6681c3c03ac Mon Sep 17 00:00:00 2001
From: Andrew Morgan <andrew@amorgan.xyz>
Date: Wed, 6 Nov 2019 14:54:24 +0000
Subject: [PATCH] Address review comments

---
 synapse/handlers/register.py         | 24 ++++++++++++------------
 synapse/replication/http/register.py |  2 +-
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/synapse/handlers/register.py b/synapse/handlers/register.py
index 8be82e375..47b9ae8d7 100644
--- a/synapse/handlers/register.py
+++ b/synapse/handlers/register.py
@@ -24,7 +24,6 @@ from synapse.api.errors import (
     AuthError,
     Codes,
     ConsentNotGivenError,
-    LimitExceededError,
     RegistrationError,
     SynapseError,
 )
@@ -218,8 +217,8 @@ class RegistrationHandler(BaseHandler):
 
         else:
             # autogen a sequential user ID
-            user = None
-            while not user:
+            # Fail after being unable to find a suitable ID a few times
+            for x in range(10):
                 localpart = yield self._generate_user_id()
                 user = UserID(localpart, self.hs.hostname)
                 user_id = user.to_string()
@@ -234,10 +233,12 @@ class RegistrationHandler(BaseHandler):
                         create_profile_with_displayname=default_display_name,
                         address=address,
                     )
+
+                    # Successfully registered
+                    break
                 except SynapseError:
                     # if user id is taken, just generate another
-                    user = None
-                    user_id = None
+                    pass
 
         if not self.hs.config.user_consent_at_registration:
             yield self._auto_join_rooms(user_id)
@@ -420,25 +421,24 @@ class RegistrationHandler(BaseHandler):
         for a given IP address
 
         Args:
-            address (str): the IP address used to perform the registration.
+            address (str|None): the IP address used to perform the registration. If this is
+                None, no ratelimiting will be performed.
 
         Raises:
             LimitExceededError: If the rate limit has been exceeded.
         """
+        if not address:
+            return
+
         time_now = self.clock.time()
 
-        allowed, time_allowed = self.ratelimiter.can_do_action(
+        self.ratelimiter.ratelimit(
             address,
             time_now_s=time_now,
             rate_hz=self.hs.config.rc_registration.per_second,
             burst_count=self.hs.config.rc_registration.burst_count,
         )
 
-        if not allowed:
-            raise LimitExceededError(
-                retry_after_ms=int(1000 * (time_allowed - time_now))
-            )
-
     def register_with_store(
         self,
         user_id,
diff --git a/synapse/replication/http/register.py b/synapse/replication/http/register.py
index 6f4bba7aa..0c4aca129 100644
--- a/synapse/replication/http/register.py
+++ b/synapse/replication/http/register.py
@@ -75,7 +75,7 @@ class ReplicationRegisterServlet(ReplicationEndpoint):
     async def _handle_request(self, request, user_id):
         content = parse_json_object_from_request(request)
 
-        await self.registration_handler.check_registration_ratelimit(content["address"])
+        self.registration_handler.check_registration_ratelimit(content["address"])
 
         await self.registration_handler.register_with_store(
             user_id=user_id,