Store appservice ID on register

2 files changed, 38 insertions, 7 deletions

diff --git a/synapse/handlers/register.py b/synapse/handlers/register.py
index e2ace6a4e5..6ffb8c0da6 100644
--- a/synapse/handlers/register.py
+++ b/synapse/handlers/register.py
@@ -182,6 +182,8 @@ class RegistrationHandler(BaseHandler):
                 errcode=Codes.EXCLUSIVE
             )
 
+        service_id = service.id if service.is_exclusive_user(user_id) else None
+
         yield self.check_user_id_not_appservice_exclusive(
             user_id, allowed_appservice=service
         )
@@ -190,7 +192,8 @@ class RegistrationHandler(BaseHandler):
         yield self.store.register(
             user_id=user_id,
             token=token,
-            password_hash=""
+            password_hash="",
+            appservice_id=service_id,
         )
         yield registered_user(self.distributor, user)
         defer.returnValue((user_id, token))
diff --git a/synapse/storage/registration.py b/synapse/storage/registration.py
index ad1157f979..aa49f53458 100644
--- a/synapse/storage/registration.py
+++ b/synapse/storage/registration.py
@@ -76,7 +76,7 @@ class RegistrationStore(SQLBaseStore):
 
     @defer.inlineCallbacks
     def register(self, user_id, token, password_hash,
-                 was_guest=False, make_guest=False):
+                 was_guest=False, make_guest=False, appservice_id=None):
         """Attempts to register an account.
 
         Args:
@@ -87,16 +87,32 @@ class RegistrationStore(SQLBaseStore):
                 upgraded to a non-guest account.
             make_guest (boolean): True if the the new user should be guest,
                 false to add a regular user account.
+            appservice_id (str): The ID of the appservice registering the user.
         Raises:
             StoreError if the user_id could not be registered.
         """
         yield self.runInteraction(
             "register",
-            self._register, user_id, token, password_hash, was_guest, make_guest
+            self._register,
+            user_id,
+            token,
+            password_hash,
+            was_guest,
+            make_guest,
+            appservice_id
         )
         self.is_guest.invalidate((user_id,))
 
-    def _register(self, txn, user_id, token, password_hash, was_guest, make_guest):
+    def _register(
+        self,
+        txn,
+        user_id,
+        token,
+        password_hash,
+        was_guest,
+        make_guest,
+        appservice_id
+    ):
         now = int(self.clock.time())
 
         next_id = self._access_tokens_id_gen.get_next()
@@ -111,9 +127,21 @@ class RegistrationStore(SQLBaseStore):
                             [password_hash, now, 1 if make_guest else 0, user_id])
             else:
                 txn.execute("INSERT INTO users "
-                            "(name, password_hash, creation_ts, is_guest) "
-                            "VALUES (?,?,?,?)",
-                            [user_id, password_hash, now, 1 if make_guest else 0])
+                            "("
+                            "   name,"
+                            "   password_hash,"
+                            "   creation_ts,"
+                            "   is_guest,"
+                            "   appservice_id"
+                            ") "
+                            "VALUES (?,?,?,?,?)",
+                            [
+                                user_id,
+                                password_hash,
+                                now,
+                                1 if make_guest else 0,
+                                appservice_id,
+                            ])
         except self.database_engine.module.IntegrityError:
             raise StoreError(
                 400, "User ID already taken.", errcode=Codes.USER_IN_USE