comparison mercurial/wireproto.py @ 37612:5e71dea79aae

wireproto: move value encoding functions to wireprototypes (API)

These functions should live in the same place. I plan to separate
client from server code in upcoming commits. wireprototypes is
where we are putting shared code like this.

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

37611:ae8730877371

# Forward a couple of names from peer to make wireproto interactions
# slightly more sensible.
batchable = peer.batchable
future = peer.future

# list of nodes encoding / decoding

def decodelist(l, sep=' '):
    if l:
        return [bin(v) for v in  l.split(sep)]
    return []

def encodelist(l, sep=' '):
    try:
        return sep.join(map(hex, l))
    except TypeError:
        raise

# batched call argument encoding

def escapearg(plain):
    return (plain
            .replace(':', ':c')
            .replace(',', ':o')
            .replace(';', ':s')
            .replace('=', ':e'))

def unescapearg(escaped):
    return (escaped
            .replace(':e', '=')
            .replace(':s', ';')
            .replace(':o', ',')
            .replace(':c', ':'))

def encodebatchcmds(req):
    """Return a ``cmds`` argument value for the ``batch`` command."""
    cmds = []
    for op, argsdict in req:
        # Old servers didn't properly unescape argument names. So prevent
        # the sending of argument names that may not be decoded properly by
        # servers.

    def heads(self):
        f = future()
        yield {}, f
        d = f.value
        try:
            yield decodelist(d[:-1])
        except ValueError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

    @batchable
    def known(self, nodes):
        f = future()
        yield {'nodes': encodelist(nodes)}, f
        d = f.value
        try:
            yield [bool(int(b)) for b in d]
        except ValueError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

        try:
            branchmap = {}
            for branchpart in d.splitlines():
                branchname, branchheads = branchpart.split(' ', 1)
                branchname = encoding.tolocal(urlreq.unquote(branchname))
                branchheads = decodelist(branchheads)
                branchmap[branchname] = branchheads
            yield branchmap
        except TypeError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

            keytype = gboptsmap.get(key)
            if keytype is None:
                raise error.ProgrammingError(
                    'Unexpectedly None keytype for key %s' % key)
            elif keytype == 'nodes':
                value = encodelist(value)
            elif keytype == 'csv':
                value = ','.join(value)
            elif keytype == 'scsv':
                value = ','.join(sorted(value))
            elif keytype == 'boolean':

        `url` is the url the client thinks it's pushing to, which is
        visible to hooks.
        '''

        if heads != ['force'] and self.capable('unbundlehash'):
            heads = encodelist(['hashed',
                                hashlib.sha1(''.join(sorted(heads))).digest()])
        else:
            heads = encodelist(heads)

        if util.safehasattr(cg, 'deltaheader'):
            # this a bundle10, do the old style call sequence
            ret, output = self._callpush("unbundle", cg, heads=heads)
            if ret == "":

    # End of ipeercommands interface.

    # Begin of ipeerlegacycommands interface.

    def branches(self, nodes):
        n = encodelist(nodes)
        d = self._call("branches", nodes=n)
        try:
            br = [tuple(decodelist(b)) for b in d.splitlines()]
            return br
        except ValueError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

    def between(self, pairs):
        batch = 8 # avoid giant requests
        r = []
        for i in xrange(0, len(pairs), batch):
            n = " ".join([encodelist(p, '-') for p in pairs[i:i + batch]])
            d = self._call("between", pairs=n)
            try:
                r.extend(l and decodelist(l) or [] for l in d.splitlines())
            except ValueError:
                self._abort(error.ResponseError(_("unexpected response:"), d))
        return r

    def changegroup(self, nodes, kind):
        n = encodelist(nodes)
        f = self._callcompressable("changegroup", roots=n)
        return changegroupmod.cg1unpacker(f, 'UN')

    def changegroupsubset(self, bases, heads, kind):
        self.requirecap('changegroupsubset', _('look up remote changes'))
        bases = encodelist(bases)
        heads = encodelist(heads)
        f = self._callcompressable("changegroupsubset",
                                   bases=bases, heads=heads)
        return changegroupmod.cg1unpacker(f, 'UN')

    # End of ipeerlegacycommands interface.

        if ui.debugflag and ui.configbool('devel', 'debug.peer-request'):
            ui.debug('devel-peer-request: batched-content\n')
            for op, args in req:
                msg = 'devel-peer-request:    - %s (%d arguments)\n'
                ui.debug(msg % (op, len(args)))

        rsp = self._callstream("batch", cmds=encodebatchcmds(req))
        chunk = rsp.read(1024)
        work = [chunk]
        while chunk:

# TODO define a more appropriate permissions type to use for this.
@wireprotocommand('batch', 'cmds *', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def batch(repo, proto, cmds, others):
    repo = repo.filtered("served")
    res = []
    for pair in cmds.split(';'):
        op, args = pair.split(' ', 1)
        vals = {}

        # For now, all batchable commands must return bytesresponse or
        # raw bytes (for backwards compatibility).
        assert isinstance(result, (wireprototypes.bytesresponse, bytes))
        if isinstance(result, wireprototypes.bytesresponse):
            result = result.data
        res.append(escapearg(result))

    return wireprototypes.bytesresponse(';'.join(res))

@wireprotocommand('between', 'pairs', transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def between(repo, proto, pairs):
    pairs = [decodelist(p, '-') for p in pairs.split(" ")]
    r = []
    for b in repo.between(pairs):
        r.append(encodelist(b) + "\n")

    return wireprototypes.bytesresponse(''.join(r))

@wireprotocommand('branchmap', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def branchmap(repo, proto):
    branchmap = repo.branchmap()
    heads = []
    for branch, nodes in branchmap.iteritems():
        branchname = urlreq.quote(encoding.fromlocal(branch))
        branchnodes = encodelist(nodes)
        heads.append('%s %s' % (branchname, branchnodes))

    return wireprototypes.bytesresponse('\n'.join(heads))

@wireprotocommand('branches', 'nodes', transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def branches(repo, proto, nodes):
    nodes = decodelist(nodes)
    r = []
    for b in repo.branches(nodes):
        r.append(encodelist(b) + "\n")

    return wireprototypes.bytesresponse(''.join(r))

@wireprotocommand('clonebundles', '', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)

    return wireprototypes.bytesresponse(' '.join(sorted(caps)))

@wireprotocommand('changegroup', 'roots', transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def changegroup(repo, proto, roots):
    nodes = decodelist(roots)
    outgoing = discovery.outgoing(repo, missingroots=nodes,
                                  missingheads=repo.heads())
    cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
    gen = iter(lambda: cg.read(32768), '')
    return wireprototypes.streamres(gen=gen)

@wireprotocommand('changegroupsubset', 'bases heads',
                  transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def changegroupsubset(repo, proto, bases, heads):
    bases = decodelist(bases)
    heads = decodelist(heads)
    outgoing = discovery.outgoing(repo, missingroots=bases,
                                  missingheads=heads)
    cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
    gen = iter(lambda: cg.read(32768), '')
    return wireprototypes.streamres(gen=gen)

def getbundle(repo, proto, others):
    opts = options('getbundle', gboptsmap.keys(), others)
    for k, v in opts.iteritems():
        keytype = gboptsmap[k]
        if keytype == 'nodes':
            opts[k] = decodelist(v)
        elif keytype == 'csv':
            opts[k] = list(v.split(','))
        elif keytype == 'scsv':
            opts[k] = set(v.split(','))
        elif keytype == 'boolean':

        gen=chunks, prefer_uncompressed=not prefercompressed)

@wireprotocommand('heads', permission='pull', transportpolicy=POLICY_V1_ONLY)
def heads(repo, proto):
    h = repo.heads()
    return wireprototypes.bytesresponse(encodelist(h) + '\n')

@wireprotocommand('hello', permission='pull', transportpolicy=POLICY_V1_ONLY)
def hello(repo, proto):
    """Called as part of SSH handshake to obtain server info.

    return wireprototypes.bytesresponse('%d %s\n' % (success, r))

@wireprotocommand('known', 'nodes *', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def known(repo, proto, nodes, others):
    v = ''.join(b and '1' or '0' for b in repo.known(decodelist(nodes)))
    return wireprototypes.bytesresponse(v)

@wireprotocommand('protocaps', 'caps', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def protocaps(repo, proto, caps):

        streamclone.generatev1wireproto(repo))

@wireprotocommand('unbundle', 'heads', permission='push',
                  transportpolicy=POLICY_V1_ONLY)
def unbundle(repo, proto, heads):
    their_heads = decodelist(heads)

    with proto.mayberedirectstdio() as output:
        try:
            exchange.check_heads(repo, their_heads, 'preparing changes')
            cleanup = lambda: None

37612:5e71dea79aae

# Forward a couple of names from peer to make wireproto interactions
# slightly more sensible.
batchable = peer.batchable
future = peer.future


def encodebatchcmds(req):
    """Return a ``cmds`` argument value for the ``batch`` command."""
    escapearg = wireprototypes.escapebatcharg

    cmds = []
    for op, argsdict in req:
        # Old servers didn't properly unescape argument names. So prevent
        # the sending of argument names that may not be decoded properly by
        # servers.

    def heads(self):
        f = future()
        yield {}, f
        d = f.value
        try:
            yield wireprototypes.decodelist(d[:-1])
        except ValueError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

    @batchable
    def known(self, nodes):
        f = future()
        yield {'nodes': wireprototypes.encodelist(nodes)}, f
        d = f.value
        try:
            yield [bool(int(b)) for b in d]
        except ValueError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

        try:
            branchmap = {}
            for branchpart in d.splitlines():
                branchname, branchheads = branchpart.split(' ', 1)
                branchname = encoding.tolocal(urlreq.unquote(branchname))
                branchheads = wireprototypes.decodelist(branchheads)
                branchmap[branchname] = branchheads
            yield branchmap
        except TypeError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

            keytype = gboptsmap.get(key)
            if keytype is None:
                raise error.ProgrammingError(
                    'Unexpectedly None keytype for key %s' % key)
            elif keytype == 'nodes':
                value = wireprototypes.encodelist(value)
            elif keytype == 'csv':
                value = ','.join(value)
            elif keytype == 'scsv':
                value = ','.join(sorted(value))
            elif keytype == 'boolean':

        `url` is the url the client thinks it's pushing to, which is
        visible to hooks.
        '''

        if heads != ['force'] and self.capable('unbundlehash'):
            heads = wireprototypes.encodelist(
                ['hashed', hashlib.sha1(''.join(sorted(heads))).digest()])
        else:
            heads = wireprototypes.encodelist(heads)

        if util.safehasattr(cg, 'deltaheader'):
            # this a bundle10, do the old style call sequence
            ret, output = self._callpush("unbundle", cg, heads=heads)
            if ret == "":

    # End of ipeercommands interface.

    # Begin of ipeerlegacycommands interface.

    def branches(self, nodes):
        n = wireprototypes.encodelist(nodes)
        d = self._call("branches", nodes=n)
        try:
            br = [tuple(wireprototypes.decodelist(b)) for b in d.splitlines()]
            return br
        except ValueError:
            self._abort(error.ResponseError(_("unexpected response:"), d))

    def between(self, pairs):
        batch = 8 # avoid giant requests
        r = []
        for i in xrange(0, len(pairs), batch):
            n = " ".join([wireprototypes.encodelist(p, '-')
                          for p in pairs[i:i + batch]])
            d = self._call("between", pairs=n)
            try:
                r.extend(l and wireprototypes.decodelist(l) or []
                         for l in d.splitlines())
            except ValueError:
                self._abort(error.ResponseError(_("unexpected response:"), d))
        return r

    def changegroup(self, nodes, kind):
        n = wireprototypes.encodelist(nodes)
        f = self._callcompressable("changegroup", roots=n)
        return changegroupmod.cg1unpacker(f, 'UN')

    def changegroupsubset(self, bases, heads, kind):
        self.requirecap('changegroupsubset', _('look up remote changes'))
        bases = wireprototypes.encodelist(bases)
        heads = wireprototypes.encodelist(heads)
        f = self._callcompressable("changegroupsubset",
                                   bases=bases, heads=heads)
        return changegroupmod.cg1unpacker(f, 'UN')

    # End of ipeerlegacycommands interface.

        if ui.debugflag and ui.configbool('devel', 'debug.peer-request'):
            ui.debug('devel-peer-request: batched-content\n')
            for op, args in req:
                msg = 'devel-peer-request:    - %s (%d arguments)\n'
                ui.debug(msg % (op, len(args)))

        unescapearg = wireprototypes.unescapebatcharg

        rsp = self._callstream("batch", cmds=encodebatchcmds(req))
        chunk = rsp.read(1024)
        work = [chunk]
        while chunk:

# TODO define a more appropriate permissions type to use for this.
@wireprotocommand('batch', 'cmds *', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def batch(repo, proto, cmds, others):
    unescapearg = wireprototypes.unescapebatcharg
    repo = repo.filtered("served")
    res = []
    for pair in cmds.split(';'):
        op, args = pair.split(' ', 1)
        vals = {}

        # For now, all batchable commands must return bytesresponse or
        # raw bytes (for backwards compatibility).
        assert isinstance(result, (wireprototypes.bytesresponse, bytes))
        if isinstance(result, wireprototypes.bytesresponse):
            result = result.data
        res.append(wireprototypes.escapebatcharg(result))

    return wireprototypes.bytesresponse(';'.join(res))

@wireprotocommand('between', 'pairs', transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def between(repo, proto, pairs):
    pairs = [wireprototypes.decodelist(p, '-') for p in pairs.split(" ")]
    r = []
    for b in repo.between(pairs):
        r.append(wireprototypes.encodelist(b) + "\n")

    return wireprototypes.bytesresponse(''.join(r))

@wireprotocommand('branchmap', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def branchmap(repo, proto):
    branchmap = repo.branchmap()
    heads = []
    for branch, nodes in branchmap.iteritems():
        branchname = urlreq.quote(encoding.fromlocal(branch))
        branchnodes = wireprototypes.encodelist(nodes)
        heads.append('%s %s' % (branchname, branchnodes))

    return wireprototypes.bytesresponse('\n'.join(heads))

@wireprotocommand('branches', 'nodes', transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def branches(repo, proto, nodes):
    nodes = wireprototypes.decodelist(nodes)
    r = []
    for b in repo.branches(nodes):
        r.append(wireprototypes.encodelist(b) + "\n")

    return wireprototypes.bytesresponse(''.join(r))

@wireprotocommand('clonebundles', '', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)

    return wireprototypes.bytesresponse(' '.join(sorted(caps)))

@wireprotocommand('changegroup', 'roots', transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def changegroup(repo, proto, roots):
    nodes = wireprototypes.decodelist(roots)
    outgoing = discovery.outgoing(repo, missingroots=nodes,
                                  missingheads=repo.heads())
    cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
    gen = iter(lambda: cg.read(32768), '')
    return wireprototypes.streamres(gen=gen)

@wireprotocommand('changegroupsubset', 'bases heads',
                  transportpolicy=POLICY_V1_ONLY,
                  permission='pull')
def changegroupsubset(repo, proto, bases, heads):
    bases = wireprototypes.decodelist(bases)
    heads = wireprototypes.decodelist(heads)
    outgoing = discovery.outgoing(repo, missingroots=bases,
                                  missingheads=heads)
    cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
    gen = iter(lambda: cg.read(32768), '')
    return wireprototypes.streamres(gen=gen)

def getbundle(repo, proto, others):
    opts = options('getbundle', gboptsmap.keys(), others)
    for k, v in opts.iteritems():
        keytype = gboptsmap[k]
        if keytype == 'nodes':
            opts[k] = wireprototypes.decodelist(v)
        elif keytype == 'csv':
            opts[k] = list(v.split(','))
        elif keytype == 'scsv':
            opts[k] = set(v.split(','))
        elif keytype == 'boolean':

        gen=chunks, prefer_uncompressed=not prefercompressed)

@wireprotocommand('heads', permission='pull', transportpolicy=POLICY_V1_ONLY)
def heads(repo, proto):
    h = repo.heads()
    return wireprototypes.bytesresponse(wireprototypes.encodelist(h) + '\n')

@wireprotocommand('hello', permission='pull', transportpolicy=POLICY_V1_ONLY)
def hello(repo, proto):
    """Called as part of SSH handshake to obtain server info.

    return wireprototypes.bytesresponse('%d %s\n' % (success, r))

@wireprotocommand('known', 'nodes *', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def known(repo, proto, nodes, others):
    v = ''.join(b and '1' or '0'
                for b in repo.known(wireprototypes.decodelist(nodes)))
    return wireprototypes.bytesresponse(v)

@wireprotocommand('protocaps', 'caps', permission='pull',
                  transportpolicy=POLICY_V1_ONLY)
def protocaps(repo, proto, caps):

        streamclone.generatev1wireproto(repo))

@wireprotocommand('unbundle', 'heads', permission='push',
                  transportpolicy=POLICY_V1_ONLY)
def unbundle(repo, proto, heads):
    their_heads = wireprototypes.decodelist(heads)

    with proto.mayberedirectstdio() as output:
        try:
            exchange.check_heads(repo, their_heads, 'preparing changes')
            cleanup = lambda: None