From df3a661e4adb7676682a5e3c298a2dfda18b08a1 Mon Sep 17 00:00:00 2001
From: Andrew Morgan <andrew@amorgan.xyz>
Date: Fri, 18 Jan 2019 10:04:47 +0000
Subject: [PATCH] Search for messages across predecessor rooms

Signed-off-by: Andrew Morgan <andrew@amorgan.xyz>
---
 synapse/api/filtering.py   |  3 ++
 synapse/handlers/search.py | 69 ++++++++++++++++++++++++++++++++++++++
 synapse/storage/state.py   |  1 +
 3 files changed, 73 insertions(+)

diff --git a/synapse/api/filtering.py b/synapse/api/filtering.py
index 16ad654864..84000e6422 100644
--- a/synapse/api/filtering.py
+++ b/synapse/api/filtering.py
@@ -444,6 +444,9 @@ class Filter(object):
     def include_redundant_members(self):
         return self.filter_json.get("include_redundant_members", False)
 
+    def add_room_ids(self, room_ids):
+        self.rooms += room_ids
+
 
 def _matches_wildcard(actual_value, filter_value):
     if filter_value.endswith("*"):
diff --git a/synapse/handlers/search.py b/synapse/handlers/search.py
index ec936bbb4e..77e7e4e0fb 100644
--- a/synapse/handlers/search.py
+++ b/synapse/handlers/search.py
@@ -37,6 +37,54 @@ class SearchHandler(BaseHandler):
     def __init__(self, hs):
         super(SearchHandler, self).__init__(hs)
 
+    @defer.inlineCallbacks
+    def get_old_rooms_from_upgraded_room(self, room_id):
+        """Retrieves room IDs of old rooms in the history of an upgraded room.
+
+        We do so by checking the m.room.create event of the room for a
+        `predecessor` key. If it exists, we add the room ID to our return
+        list and then check that room for a m.room.create event and so on
+        until we can no longer find any more previous rooms.
+
+        The full list of all found rooms in then returned.
+
+        Args:
+            room_id (str): The ID of the room to search through.
+
+        Returns:
+            dict of past room IDs as strings
+        """
+
+        historical_room_ids = []
+
+        while True:
+            state_ids = yield self.store.get_current_state_ids(room_id)
+            create_id = state_ids.get((EventTypes.Create, ""))
+
+            # If we can't find the create event, assume we've hit a dead end
+            if not create_id:
+                break
+
+            # Retrieve the room's create event
+            create_event = yield self.store.get_event(create_id)
+
+            if not create_event:
+                break
+
+            # Check if a predecessor room is present
+            predecessor = create_event.content.get("predecessor", None)
+            if not predecessor:
+                break
+
+            # Add predecessor's room ID
+            historical_room_id = predecessor["room_id"]
+            historical_room_ids.append(historical_room_id)
+
+            # Scan through the old room for further predecessors
+            room_id = historical_room_id
+
+        defer.returnValue(historical_room_ids)
+
     @defer.inlineCallbacks
     def search(self, user, content, batch=None):
         """Performs a full text search for a user.
@@ -139,6 +187,27 @@ class SearchHandler(BaseHandler):
 
         room_ids = search_filter.filter_rooms(room_ids)
 
+        # If doing a subset of all rooms seearch, check if any of the rooms
+        # are from an upgraded room, and search their contents as well
+        # XXX: There is the possibility that we don't have a create event for
+        # the room in question, in which case we can't return all the results
+        # we want to.
+        # Ideally we would just return the results we can get now, and
+        # try to get more results from other servers in the background.
+        if search_filter.rooms:
+            historical_room_ids = []
+            for room_id in room_ids:
+                # Add any previous rooms to the search if they exist
+                ids = yield self.get_old_rooms_from_upgraded_room(room_id)
+                historical_room_ids += ids
+
+            # Add any found rooms to the list to search
+            for historical_room_id in historical_room_ids:
+                room_ids.add(historical_room_id)
+
+            # Prevent any historical events from being filtered
+            search_filter.add_room_ids(historical_room_ids)
+
         if batch_group == "room_id":
             room_ids.intersection_update({batch_group_key})
 
diff --git a/synapse/storage/state.py b/synapse/storage/state.py
index a134e9b3e8..49b3ff4a71 100644
--- a/synapse/storage/state.py
+++ b/synapse/storage/state.py
@@ -448,6 +448,7 @@ class StateGroupWorkerStore(EventsWorkerStore, SQLBaseStore):
         Returns:
             deferred: dict of (type, state_key) -> event_id
         """
+
         def _get_current_state_ids_txn(txn):
             txn.execute(
                 """SELECT type, state_key, event_id FROM current_state_events