mercurial/wireproto.py

36747:4c71a26a4009

POLICY_V1_ONLY = 'v1-only'
POLICY_V2_ONLY = 'v2-only'

commands = commanddict()

def wireprotocommand(name, args='', transportpolicy=POLICY_ALL):
    """Decorator to declare a wire protocol command.

    ``name`` is the name of the wire protocol command being provided.

    def register(func):
        commands[name] = commandentry(func, args=args, transports=transports)
        return func
    return register

@wireprotocommand('batch', 'cmds *')
def batch(repo, proto, cmds, others):
    repo = repo.filtered("served")
    res = []
    for pair in cmds.split(';'):

        for a in args.split(','):
            if a:
                n, v = a.split('=')
                vals[unescapearg(n)] = unescapearg(v)
        func, spec = commands[op]
        if spec:
            keys = spec.split()
            data = {}
            for k in keys:
                if k == '*':

            result = result.data
        res.append(escapearg(result))

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

@wireprotocommand('between', 'pairs', transportpolicy=POLICY_V1_ONLY)
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 bytesresponse(''.join(r))

@wireprotocommand('branchmap')
def branchmap(repo, proto):
    branchmap = repo.branchmap()
    heads = []
    for branch, nodes in branchmap.iteritems():

        branchnodes = encodelist(nodes)
        heads.append('%s %s' % (branchname, branchnodes))

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

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

    return bytesresponse(''.join(r))

@wireprotocommand('clonebundles', '')
def clonebundles(repo, proto):
    """Server command for returning info for available bundles to seed clones.

    Clients will parse this response and determine what bundle to fetch.

    return proto.addcapabilities(repo, caps)

# If you are writing an extension and consider wrapping this function. Wrap
# `_capabilities` instead.
@wireprotocommand('capabilities')
def capabilities(repo, proto):
    return bytesresponse(' '.join(_capabilities(repo, proto)))

@wireprotocommand('changegroup', 'roots', transportpolicy=POLICY_V1_ONLY)
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 streamres(gen=gen)

@wireprotocommand('changegroupsubset', 'bases heads',
                  transportpolicy=POLICY_V1_ONLY)
def changegroupsubset(repo, proto, bases, heads):
    bases = decodelist(bases)
    heads = decodelist(heads)

                                  missingheads=heads)
    cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
    gen = iter(lambda: cg.read(32768), '')
    return streamres(gen=gen)

@wireprotocommand('debugwireargs', 'one two *')
def debugwireargs(repo, proto, one, two, others):
    # only accept optional args from the known set
    opts = options('debugwireargs', ['three', 'four'], others)
    return bytesresponse(repo.debugwireargs(one, two,
                                            **pycompat.strkwargs(opts)))

@wireprotocommand('getbundle', '*')
def getbundle(repo, proto, others):
    opts = options('getbundle', gboptsmap.keys(), others)
    for k, v in opts.iteritems():
        keytype = gboptsmap[k]

        chunks = bundler.getchunks()
        prefercompressed = False

    return streamres(gen=chunks, prefer_uncompressed=not prefercompressed)

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

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

    Returns a list of lines describing interesting things about the

        capabilities: <token0> <token1> <token2>
    """
    caps = capabilities(repo, proto).data
    return bytesresponse('capabilities: %s\n' % caps)

@wireprotocommand('listkeys', 'namespace')
def listkeys(repo, proto, namespace):
    d = sorted(repo.listkeys(encoding.tolocal(namespace)).items())
    return bytesresponse(pushkeymod.encodekeys(d))

@wireprotocommand('lookup', 'key')
def lookup(repo, proto, key):
    try:
        k = encoding.tolocal(key)
        c = repo[k]

    except Exception as inst:
        r = util.forcebytestr(inst)
        success = 0
    return bytesresponse('%d %s\n' % (success, r))

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

@wireprotocommand('pushkey', 'namespace key old new')
def pushkey(repo, proto, namespace, key, old, new):
    # compatibility with pre-1.8 clients which were accidentally
    # sending raw binary nodes rather than utf-8-encoded hex
    if len(new) == 20 and util.escapestr(new) != new:

                         encoding.tolocal(old), new) or False

    output = output.getvalue() if output else ''
    return bytesresponse('%d\n%s' % (int(r), output))

@wireprotocommand('stream_out')
def stream(repo, proto):
    '''If the server supports streaming clone, it advertises the "stream"
    capability with a value representing the version and flags of the repo
    it is serving. Client checks to see if it understands the format.
    '''
    return streamres_legacy(streamclone.generatev1wireproto(repo))

@wireprotocommand('unbundle', 'heads')
def unbundle(repo, proto, heads):
    their_heads = decodelist(heads)

    with proto.mayberedirectstdio() as output:

36760:7bf80d9d9543

POLICY_V1_ONLY = 'v1-only'
POLICY_V2_ONLY = 'v2-only'

commands = commanddict()

# Maps wire protocol name to operation type. This is used for permissions
# checking. All defined @wireiprotocommand should have an entry in this
# dict.
permissions = {}

def wireprotocommand(name, args='', transportpolicy=POLICY_ALL):
    """Decorator to declare a wire protocol command.

    ``name`` is the name of the wire protocol command being provided.

    def register(func):
        commands[name] = commandentry(func, args=args, transports=transports)
        return func
    return register

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

        for a in args.split(','):
            if a:
                n, v = a.split('=')
                vals[unescapearg(n)] = unescapearg(v)
        func, spec = commands[op]

        # If the protocol supports permissions checking, perform that
        # checking on each batched command.
        # TODO formalize permission checking as part of protocol interface.
        if util.safehasattr(proto, 'checkperm'):
            # Assume commands with no defined permissions are writes / for
            # pushes. This is the safest from a security perspective because
            # it doesn't allow commands with undefined semantics from
            # bypassing permissions checks.
            proto.checkperm(permissions.get(op, 'push'))

        if spec:
            keys = spec.split()
            data = {}
            for k in keys:
                if k == '*':

            result = result.data
        res.append(escapearg(result))

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

permissions['between'] = 'pull'
@wireprotocommand('between', 'pairs', transportpolicy=POLICY_V1_ONLY)
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 bytesresponse(''.join(r))

permissions['branchmap'] = 'pull'
@wireprotocommand('branchmap')
def branchmap(repo, proto):
    branchmap = repo.branchmap()
    heads = []
    for branch, nodes in branchmap.iteritems():

        branchnodes = encodelist(nodes)
        heads.append('%s %s' % (branchname, branchnodes))

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

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

    return bytesresponse(''.join(r))

permissions['clonebundles'] = 'pull'
@wireprotocommand('clonebundles', '')
def clonebundles(repo, proto):
    """Server command for returning info for available bundles to seed clones.

    Clients will parse this response and determine what bundle to fetch.

    return proto.addcapabilities(repo, caps)

# If you are writing an extension and consider wrapping this function. Wrap
# `_capabilities` instead.
permissions['capabilities'] = 'pull'
@wireprotocommand('capabilities')
def capabilities(repo, proto):
    return bytesresponse(' '.join(_capabilities(repo, proto)))

permissions['changegroup'] = 'pull'
@wireprotocommand('changegroup', 'roots', transportpolicy=POLICY_V1_ONLY)
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 streamres(gen=gen)

permissions['changegroupsubset'] = 'pull'
@wireprotocommand('changegroupsubset', 'bases heads',
                  transportpolicy=POLICY_V1_ONLY)
def changegroupsubset(repo, proto, bases, heads):
    bases = decodelist(bases)
    heads = decodelist(heads)

                                  missingheads=heads)
    cg = changegroupmod.makechangegroup(repo, outgoing, '01', 'serve')
    gen = iter(lambda: cg.read(32768), '')
    return streamres(gen=gen)

permissions['debugwireargs'] = 'pull'
@wireprotocommand('debugwireargs', 'one two *')
def debugwireargs(repo, proto, one, two, others):
    # only accept optional args from the known set
    opts = options('debugwireargs', ['three', 'four'], others)
    return bytesresponse(repo.debugwireargs(one, two,
                                            **pycompat.strkwargs(opts)))

permissions['getbundle'] = 'pull'
@wireprotocommand('getbundle', '*')
def getbundle(repo, proto, others):
    opts = options('getbundle', gboptsmap.keys(), others)
    for k, v in opts.iteritems():
        keytype = gboptsmap[k]

        chunks = bundler.getchunks()
        prefercompressed = False

    return streamres(gen=chunks, prefer_uncompressed=not prefercompressed)

permissions['heads'] = 'pull'
@wireprotocommand('heads')
def heads(repo, proto):
    h = repo.heads()
    return bytesresponse(encodelist(h) + '\n')

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

    Returns a list of lines describing interesting things about the

        capabilities: <token0> <token1> <token2>
    """
    caps = capabilities(repo, proto).data
    return bytesresponse('capabilities: %s\n' % caps)

permissions['listkeys'] = 'pull'
@wireprotocommand('listkeys', 'namespace')
def listkeys(repo, proto, namespace):
    d = sorted(repo.listkeys(encoding.tolocal(namespace)).items())
    return bytesresponse(pushkeymod.encodekeys(d))

permissions['lookup'] = 'pull'
@wireprotocommand('lookup', 'key')
def lookup(repo, proto, key):
    try:
        k = encoding.tolocal(key)
        c = repo[k]

    except Exception as inst:
        r = util.forcebytestr(inst)
        success = 0
    return bytesresponse('%d %s\n' % (success, r))

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

permissions['pushkey'] = 'push'
@wireprotocommand('pushkey', 'namespace key old new')
def pushkey(repo, proto, namespace, key, old, new):
    # compatibility with pre-1.8 clients which were accidentally
    # sending raw binary nodes rather than utf-8-encoded hex
    if len(new) == 20 and util.escapestr(new) != new:

                         encoding.tolocal(old), new) or False

    output = output.getvalue() if output else ''
    return bytesresponse('%d\n%s' % (int(r), output))

permissions['stream_out'] = 'pull'
@wireprotocommand('stream_out')
def stream(repo, proto):
    '''If the server supports streaming clone, it advertises the "stream"
    capability with a value representing the version and flags of the repo
    it is serving. Client checks to see if it understands the format.
    '''
    return streamres_legacy(streamclone.generatev1wireproto(repo))

permissions['unbundle'] = 'push'
@wireprotocommand('unbundle', 'heads')
def unbundle(repo, proto, heads):
    their_heads = decodelist(heads)

    with proto.mayberedirectstdio() as output: