wireproto: move version 2 command handlers to wireprotov2server
authorGregory Szorc <gregory.szorc@gmail.com>
Mon, 09 Apr 2018 19:35:39 -0700
changeset 37546 3a2367e6c6f2
parent 37545 93397c4633f6
child 37547 835ccc2a5ef1
wireproto: move version 2 command handlers to wireprotov2server This is relatively straightforward. As part of this, we introduced a local @wireprotocommand that wraps the main one and defines a v2 only policy by default. Because the hacky HTTPv2 peer isn't using capabilities response yet, we had to move some code around to force import of wireprotov2server so commands are registered. This is super hacky. But this code will go away once the HTTPv2 peer is using the capabilities response to derive permissions. Differential Revision: https://phab.mercurial-scm.org/D3231
mercurial/httppeer.py
mercurial/wireproto.py
mercurial/wireprotov2server.py
--- a/mercurial/httppeer.py	Mon Apr 09 19:35:04 2018 -0700
+++ b/mercurial/httppeer.py	Mon Apr 09 19:35:39 2018 -0700
@@ -493,6 +493,15 @@
     def _call(self, name, **args):
         """Call a wire protocol command with arguments."""
 
+        # Having this early has a side-effect of importing wireprotov2server,
+        # which has the side-effect of ensuring commands are registered.
+
+        # TODO modify user-agent to reflect v2.
+        headers = {
+            r'Accept': wireprotov2server.FRAMINGTYPE,
+            r'Content-Type': wireprotov2server.FRAMINGTYPE,
+        }
+
         # TODO permissions should come from capabilities results.
         permission = wireproto.commandsv2[name].permission
         if permission not in ('push', 'pull'):
@@ -507,12 +516,6 @@
         url = '%s/api/%s/%s/%s' % (self.url, wireprotov2server.HTTPV2,
                                    permission, name)
 
-        # TODO modify user-agent to reflect v2.
-        headers = {
-            r'Accept': wireprotov2server.FRAMINGTYPE,
-            r'Content-Type': wireprotov2server.FRAMINGTYPE,
-        }
-
         # TODO this should be part of a generic peer for the frame-based
         # protocol.
         reactor = wireprotoframing.clientreactor(hasmultiplesend=False,
--- a/mercurial/wireproto.py	Mon Apr 09 19:35:04 2018 -0700
+++ b/mercurial/wireproto.py	Mon Apr 09 19:35:39 2018 -0700
@@ -1305,113 +1305,3 @@
                 bundler.newpart('error:pushraced',
                                 [('message', stringutil.forcebytestr(exc))])
             return wireprototypes.streamreslegacy(gen=bundler.getchunks())
-
-# Wire protocol version 2 commands only past this point.
-
-def _capabilitiesv2(repo, proto):
-    """Obtain the set of capabilities for version 2 transports.
-
-    These capabilities are distinct from the capabilities for version 1
-    transports.
-    """
-    compression = []
-    for engine in supportedcompengines(repo.ui, util.SERVERROLE):
-        compression.append({
-            b'name': engine.wireprotosupport().name,
-        })
-
-    caps = {
-        'commands': {},
-        'compression': compression,
-    }
-
-    for command, entry in commandsv2.items():
-        caps['commands'][command] = {
-            'args': entry.args,
-            'permissions': [entry.permission],
-        }
-
-    return proto.addcapabilities(repo, caps)
-
-@wireprotocommand('branchmap', permission='pull',
-                  transportpolicy=POLICY_V2_ONLY)
-def branchmapv2(repo, proto):
-    branchmap = {encoding.fromlocal(k): v
-                 for k, v in repo.branchmap().iteritems()}
-
-    return wireprototypes.cborresponse(branchmap)
-
-@wireprotocommand('capabilities', permission='pull',
-                  transportpolicy=POLICY_V2_ONLY)
-def capabilitiesv2(repo, proto):
-    caps = _capabilitiesv2(repo, proto)
-
-    return wireprototypes.cborresponse(caps)
-
-@wireprotocommand('heads',
-                  args={
-                      'publiconly': False,
-                  },
-                  permission='pull',
-                  transportpolicy=POLICY_V2_ONLY)
-def headsv2(repo, proto, publiconly=False):
-    if publiconly:
-        repo = repo.filtered('immutable')
-
-    return wireprototypes.cborresponse(repo.heads())
-
-@wireprotocommand('known',
-                  args={
-                      'nodes': [b'deadbeef'],
-                  },
-                  permission='pull',
-                  transportpolicy=POLICY_V2_ONLY)
-def knownv2(repo, proto, nodes=None):
-    nodes = nodes or []
-    result = b''.join(b'1' if n else b'0' for n in repo.known(nodes))
-    return wireprototypes.cborresponse(result)
-
-@wireprotocommand('listkeys',
-                  args={
-                      'namespace': b'ns',
-                  },
-                  permission='pull',
-                  transportpolicy=POLICY_V2_ONLY)
-def listkeysv2(repo, proto, namespace=None):
-    keys = repo.listkeys(encoding.tolocal(namespace))
-    keys = {encoding.fromlocal(k): encoding.fromlocal(v)
-            for k, v in keys.iteritems()}
-
-    return wireprototypes.cborresponse(keys)
-
-@wireprotocommand('lookup',
-                  args={
-                      'key': b'foo',
-                  },
-                  permission='pull',
-                  transportpolicy=POLICY_V2_ONLY)
-def lookupv2(repo, proto, key):
-    key = encoding.tolocal(key)
-
-    # TODO handle exception.
-    node = repo.lookup(key)
-
-    return wireprototypes.cborresponse(node)
-
-@wireprotocommand('pushkey',
-                  args={
-                      'namespace': b'ns',
-                      'key': b'key',
-                      'old': b'old',
-                      'new': b'new',
-                  },
-                  permission='push',
-                  transportpolicy=POLICY_V2_ONLY)
-def pushkeyv2(repo, proto, namespace, key, old, new):
-    # TODO handle ui output redirection
-    r = repo.pushkey(encoding.tolocal(namespace),
-                     encoding.tolocal(key),
-                     encoding.tolocal(old),
-                     encoding.tolocal(new))
-
-    return wireprototypes.cborresponse(r)
--- a/mercurial/wireprotov2server.py	Mon Apr 09 19:35:04 2018 -0700
+++ b/mercurial/wireprotov2server.py	Mon Apr 09 19:35:39 2018 -0700
@@ -16,8 +16,10 @@
     interface as zi,
 )
 from . import (
+    encoding,
     error,
     pycompat,
+    util,
     wireproto,
     wireprotoframing,
     wireprototypes,
@@ -362,3 +364,111 @@
 
     def checkperm(self, perm):
         raise NotImplementedError
+
+def _capabilitiesv2(repo, proto):
+    """Obtain the set of capabilities for version 2 transports.
+
+    These capabilities are distinct from the capabilities for version 1
+    transports.
+    """
+    compression = []
+    for engine in wireproto.supportedcompengines(repo.ui, util.SERVERROLE):
+        compression.append({
+            b'name': engine.wireprotosupport().name,
+        })
+
+    caps = {
+        'commands': {},
+        'compression': compression,
+    }
+
+    for command, entry in wireproto.commandsv2.items():
+        caps['commands'][command] = {
+            'args': entry.args,
+            'permissions': [entry.permission],
+        }
+
+    return proto.addcapabilities(repo, caps)
+
+def wireprotocommand(*args, **kwargs):
+    def register(func):
+        return wireproto.wireprotocommand(
+            *args, transportpolicy=wireproto.POLICY_V2_ONLY, **kwargs)(func)
+
+    return register
+
+@wireprotocommand('branchmap', permission='pull')
+def branchmapv2(repo, proto):
+    branchmap = {encoding.fromlocal(k): v
+                 for k, v in repo.branchmap().iteritems()}
+
+    return wireprototypes.cborresponse(branchmap)
+
+@wireprotocommand('capabilities', permission='pull')
+def capabilitiesv2(repo, proto):
+    caps = _capabilitiesv2(repo, proto)
+
+    return wireprototypes.cborresponse(caps)
+
+@wireprotocommand('heads',
+                  args={
+                      'publiconly': False,
+                  },
+                  permission='pull')
+def headsv2(repo, proto, publiconly=False):
+    if publiconly:
+        repo = repo.filtered('immutable')
+
+    return wireprototypes.cborresponse(repo.heads())
+
+@wireprotocommand('known',
+                  args={
+                      'nodes': [b'deadbeef'],
+                  },
+                  permission='pull')
+def knownv2(repo, proto, nodes=None):
+    nodes = nodes or []
+    result = b''.join(b'1' if n else b'0' for n in repo.known(nodes))
+    return wireprototypes.cborresponse(result)
+
+@wireprotocommand('listkeys',
+                  args={
+                      'namespace': b'ns',
+                  },
+                  permission='pull')
+def listkeysv2(repo, proto, namespace=None):
+    keys = repo.listkeys(encoding.tolocal(namespace))
+    keys = {encoding.fromlocal(k): encoding.fromlocal(v)
+            for k, v in keys.iteritems()}
+
+    return wireprototypes.cborresponse(keys)
+
+@wireprotocommand('lookup',
+                  args={
+                      'key': b'foo',
+                  },
+                  permission='pull')
+def lookupv2(repo, proto, key):
+    key = encoding.tolocal(key)
+
+    # TODO handle exception.
+    node = repo.lookup(key)
+
+    return wireprototypes.cborresponse(node)
+
+@wireprotocommand('pushkey',
+                  args={
+                      'namespace': b'ns',
+                      'key': b'key',
+                      'old': b'old',
+                      'new': b'new',
+                  },
+                  permission='push')
+def pushkeyv2(repo, proto, namespace, key, old, new):
+    # TODO handle ui output redirection
+    r = repo.pushkey(encoding.tolocal(namespace),
+                     encoding.tolocal(key),
+                     encoding.tolocal(old),
+                     encoding.tolocal(new))
+
+    return wireprototypes.cborresponse(r)