From b5db4ed5f68ee81557393e94436d768b955b1aa0 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 31 May 2017 13:30:12 +0100
Subject: [PATCH] Update room column when room becomes unpublic

---
 synapse/handlers/user_directory.py | 23 +++++++++++++++++++++--
 synapse/storage/user_directory.py  | 10 ++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py
index e9488ce55..0cf403f59 100644
--- a/synapse/handlers/user_directory.py
+++ b/synapse/handlers/user_directory.py
@@ -121,12 +121,13 @@ class UserDirectoyHandler(object):
 
         # TODO: Make this faster?
         rooms = yield self.store.get_rooms_for_user(user_id)
-        for room_id in rooms:
+        for j_room_id in rooms:
             is_public = yield self.store.is_room_world_readable_or_publicly_joinable(
-                room_id
+                j_room_id
             )
 
             if is_public:
+                yield self.store.update_user_in_user_dir(user_id, j_room_id)
                 return
 
         yield self.store.remove_from_user_dir(user_id)
@@ -149,6 +150,15 @@ class UserDirectoyHandler(object):
                 if change is None:
                     continue
 
+                is_public = yield self.store.is_room_world_readable_or_publicly_joinable(
+                    room_id
+                )
+
+                if change and is_public:
+                    continue
+                elif not change and not is_public:
+                    continue
+
                 users_with_profile = yield self.state.get_current_user_in_room(room_id)
                 for user_id, profile in users_with_profile.iteritems():
                     if change:
@@ -164,6 +174,15 @@ class UserDirectoyHandler(object):
                 if change is None:
                     continue
 
+                is_public = yield self.store.is_room_world_readable_or_publicly_joinable(
+                    room_id
+                )
+
+                if change and is_public:
+                    continue
+                elif not change and not is_public:
+                    continue
+
                 users_with_profile = yield self.state.get_current_user_in_room(room_id)
                 for user_id, profile in users_with_profile.iteritems():
                     if change:
diff --git a/synapse/storage/user_directory.py b/synapse/storage/user_directory.py
index 6c7c8c4be..d72b93b58 100644
--- a/synapse/storage/user_directory.py
+++ b/synapse/storage/user_directory.py
@@ -79,6 +79,16 @@ class UserDirectoryStore(SQLBaseStore):
             "add_profiles_to_user_dir", _add_profiles_to_user_dir_txn
         )
 
+    @defer.inlineCallbacks
+    def update_user_in_user_dir(self, user_id, room_id):
+        yield self._simple_update_one(
+            table="user_directory",
+            keyvalues={"user_id": user_id},
+            updatevalues={"room_id": room_id},
+            desc="update_user_in_user_dir",
+        )
+        self.get_user_in_directory.invalidate((user_id,))
+
     @defer.inlineCallbacks
     def remove_from_user_dir(self, user_id):
         yield self._simple_delete(