diff mercurial/wireprotoframing.py @ 37063:39304dd63589

wireproto: explicitly track which requests are active

We previously only tracked which requests are receiving. A
misbehaving client could accidentally have multiple requests with
the same ID in flight.

We now explicitly track which request IDs are currently active.
We make it illegal to receive a frame associated with a request
ID that has already been dispatched.

Differential Revision: https://phab.mercurial-scm.org/D2901

--- a/mercurial/wireprotoframing.py	Thu Mar 15 16:09:58 2018 -0700
+++ b/mercurial/wireprotoframing.py	Thu Mar 15 18:05:49 2018 -0700
@@ -472,6 +472,10 @@
         self._bufferedframegens = []
         # request id -> dict of commands that are actively being received.
         self._receivingcommands = {}
+        # Request IDs that have been received and are actively being processed.
+        # Once all output for a request has been sent, it is removed from this
+        # set.
+        self._activecommands = set()
 
     def onframerecv(self, frame):
         """Process a frame that has been received off the wire.
@@ -496,14 +500,20 @@
 
         The raw bytes response is passed as an argument.
         """
-        framegen = createbytesresponseframesfrombytes(requestid, data)
+        def sendframes():
+            for frame in createbytesresponseframesfrombytes(requestid, data):
+                yield frame
+
+            self._activecommands.remove(requestid)
+
+        result = sendframes()
 
         if self._deferoutput:
-            self._bufferedframegens.append(framegen)
+            self._bufferedframegens.append(result)
             return 'noop', {}
         else:
             return 'sendframes', {
-                'framegen': framegen,
+                'framegen': result,
             }
 
     def oninputeof(self):
@@ -546,6 +556,9 @@
         else:
             self._state = 'idle'
 
+        assert requestid not in self._activecommands
+        self._activecommands.add(requestid)
+
         return 'runcommand', {
             'requestid': requestid,
             'command': entry['command'],
@@ -571,6 +584,11 @@
             return self._makeerrorresult(
                 _('request with ID %d already received') % frame.requestid)
 
+        if frame.requestid in self._activecommands:
+            self._state = 'errored'
+            return self._makeerrorresult((
+                _('request with ID %d is already active') % frame.requestid))
+
         expectingargs = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_ARGS)
         expectingdata = bool(frame.flags & FLAG_COMMAND_NAME_HAVE_DATA)
 
@@ -599,7 +617,13 @@
             return self._onframeidle(frame)
 
         # All other frames should be related to a command that is currently
-        # receiving.
+        # receiving but is not active.
+        if frame.requestid in self._activecommands:
+            self._state = 'errored'
+            return self._makeerrorresult(
+                _('received frame for request that is still active: %d') %
+                frame.requestid)
+
         if frame.requestid not in self._receivingcommands:
             self._state = 'errored'
             return self._makeerrorresult(