Make it possible to paginate forwards from stream tokens

In order that we can fill the gap after a /sync, make it possible to paginate
forwards from a stream token.

2 files changed, 35 insertions, 24 deletions

diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index ff800f8af1..b73ad62147 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -105,8 +105,6 @@ class MessageHandler(BaseHandler):
             room_token = pagin_config.from_token.room_key
 
         room_token = RoomStreamToken.parse(room_token)
-        if room_token.topological is None:
-            raise SynapseError(400, "Invalid token")
 
         pagin_config.from_token = pagin_config.from_token.copy_and_replace(
             "room_key", str(room_token)
@@ -117,27 +115,28 @@ class MessageHandler(BaseHandler):
         membership, member_event_id = yield self._check_in_room_or_world_readable(
             room_id, user_id
         )
-        if membership == Membership.LEAVE:
-            # If they have left the room then clamp the token to be before
-            # they left the room.
-            leave_token = yield self.store.get_topological_token_for_event(
-                member_event_id
+
+        if source_config.direction == 'b':
+            # if we're going backwards, we might need to backfill. This
+            # requires that we have a topo token.
+            if room_token.topological is None:
+                raise SynapseError(400, "Invalid token: cannot paginate "
+                                        "backwards from a stream token")
+
+            if membership == Membership.LEAVE:
+                # If they have left the room then clamp the token to be before
+                # they left the room, to save the effort of loading from the
+                # database.
+                leave_token = yield self.store.get_topological_token_for_event(
+                    member_event_id
+                )
+                leave_token = RoomStreamToken.parse(leave_token)
+                if leave_token.topological < room_token.topological:
+                    source_config.from_key = str(leave_token)
+
+            yield self.hs.get_handlers().federation_handler.maybe_backfill(
+                room_id, room_token.topological
             )
-            leave_token = RoomStreamToken.parse(leave_token)
-            if leave_token.topological < room_token.topological:
-                source_config.from_key = str(leave_token)
-
-            if source_config.direction == "f":
-                if source_config.to_key is None:
-                    source_config.to_key = str(leave_token)
-                else:
-                    to_token = RoomStreamToken.parse(source_config.to_key)
-                    if leave_token.topological < to_token.topological:
-                        source_config.to_key = str(leave_token)
-
-        yield self.hs.get_handlers().federation_handler.maybe_backfill(
-            room_id, room_token.topological
-        )
 
         events, next_key = yield data_source.get_pagination_rows(
             requester.user, source_config, room_id
diff --git a/tests/rest/client/v1/test_rooms.py b/tests/rest/client/v1/test_rooms.py
index cd03106e88..2fe6f695f5 100644
--- a/tests/rest/client/v1/test_rooms.py
+++ b/tests/rest/client/v1/test_rooms.py
@@ -1045,8 +1045,20 @@ class RoomMessageListTestCase(RestTestCase):
         self.assertTrue("end" in response)
 
     @defer.inlineCallbacks
-    def test_stream_token_is_rejected(self):
+    def test_stream_token_is_rejected_for_back_pagination(self):
         (code, response) = yield self.mock_resource.trigger_get(
-            "/rooms/%s/messages?access_token=x&from=s0_0_0_0" %
+            "/rooms/%s/messages?access_token=x&from=s0_0_0_0_0&dir=b" %
             self.room_id)
         self.assertEquals(400, code)
+
+    @defer.inlineCallbacks
+    def test_stream_token_is_accepted_for_fwd_pagianation(self):
+        token = "s0_0_0_0_0"
+        (code, response) = yield self.mock_resource.trigger_get(
+            "/rooms/%s/messages?access_token=x&from=%s" %
+            (self.room_id, token))
+        self.assertEquals(200, code)
+        self.assertTrue("start" in response)
+        self.assertEquals(token, response['start'])
+        self.assertTrue("chunk" in response)
+        self.assertTrue("end" in response)
\ No newline at end of file