comparison mercurial/wireprotoserver.py @ 43076:2372284d9457

formatting: blacken the codebase

This is using my patch to black
(https://github.com/psf/black/pull/826) so we don't un-wrap collection
literals.

Done with:

  hg files 'set:**.py - mercurial/thirdparty/** - "contrib/python-zstandard/**"' | xargs black -S

# skip-blame mass-reformatting only

# no-check-commit reformats foo_bar functions

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

43075:57875cf423c9

    util,
    wireprototypes,
    wireprotov1server,
    wireprotov2server,
)
from .interfaces import (
    util as interfaceutil,
)
from .utils import (
    cborutil,
    compression,
)

HGTYPE2 = 'application/mercurial-0.2'
HGERRTYPE = 'application/hg-error'

SSHV1 = wireprototypes.SSHV1
SSHV2 = wireprototypes.SSHV2

def decodevaluefromheaders(req, headerprefix):
    """Decode a long value from multiple HTTP request headers.

    Returns the value as a bytes, not a str.

            break
        chunks.append(pycompat.bytesurl(v))
        i += 1

    return ''.join(chunks)

@interfaceutil.implementer(wireprototypes.baseprotocolhandler)
class httpv1protocolhandler(object):
    def __init__(self, req, ui, checkperm):
        self._req = req

    def _args(self):
        args = self._req.qsparams.asdictoflists()
        postlen = int(self._req.headers.get(b'X-HgArgs-Post', 0))
        if postlen:
            args.update(urlreq.parseqs(
                self._req.bodyfh.read(postlen), keep_blank_values=True))
            return args

        argvalue = decodevaluefromheaders(self._req, b'X-HgArg')
        args.update(urlreq.parseqs(argvalue, keep_blank_values=True))
        return args

    def client(self):
        return 'remote:%s:%s:%s' % (
            self._req.urlscheme,
            urlreq.quote(self._req.remotehost or ''),
            urlreq.quote(self._req.remoteuser or ''))

    def addcapabilities(self, repo, caps):
        caps.append(b'batch')

        caps.append('httpheader=%d' %
                    repo.ui.configint('server', 'maxhttpheaderlen'))
        if repo.ui.configbool('experimental', 'httppostargs'):
            caps.append('httppostargs')

        # FUTURE advertise 0.2rx once support is implemented
        # FUTURE advertise minrx and mintx after consulting config option
        caps.append('httpmediatype=0.1rx,0.1tx,0.2tx')

        compengines = wireprototypes.supportedcompengines(repo.ui,
            compression.SERVERROLE)
        if compengines:
            comptypes = ','.join(urlreq.quote(e.wireprotosupport().name)
                                 for e in compengines)
            caps.append('compression=%s' % comptypes)

        return caps

    def checkperm(self, perm):
        return self._checkperm(perm)

# This method exists mostly so that extensions like remotefilelog can
# disable a kludgey legacy method only over http. As of early 2018,
# there are no other known users, so with any luck we can discard this
# hook if remotefilelog becomes a first-party extension.
def iscmd(cmd):
    return cmd in wireprotov1server.commands

def handlewsgirequest(rctx, req, res, checkperm):
    """Possibly process a wire protocol request.

    If the current request is a wire protocol request, the request is

        # TODO This is not a good response to issue for this request. This
        # is mostly for BC for now.
        res.setbodybytes('0\n%s\n' % b'Not Found')
        return True

    proto = httpv1protocolhandler(req, repo.ui,
                                  lambda perm: checkperm(rctx, req, perm))

    # The permissions checker should be the only thing that can raise an
    # ErrorResponse. It is kind of a layer violation to catch an hgweb
    # exception here. So consider refactoring into a exception type that
    # is associated with the wire protocol.

        # "unbundle." That assumption is not always valid.
        res.setbodybytes('0\n%s\n' % pycompat.bytestr(e))

    return True

def _availableapis(repo):
    apis = set()

    # Registered APIs are made available via config options of the name of
    # the protocol.

        section, option = v['config']
        if repo.ui.configbool(section, option):
            apis.add(k)

    return apis

def handlewsgiapirequest(rctx, req, res, checkperm):
    """Handle requests to /api/*."""
    assert req.dispatchparts[0] == b'api'

    # Requests to /api/ list available APIs.
    if req.dispatchparts == [b'api']:
        res.status = b'200 OK'
        res.headers[b'Content-Type'] = b'text/plain'
        lines = [_('APIs can be accessed at /api/<name>, where <name> can be '
                   'one of the following:\n')]
        if availableapis:
            lines.extend(sorted(availableapis))
        else:
            lines.append(_('(no available APIs)\n'))
        res.setbodybytes(b'\n'.join(lines))

    proto = req.dispatchparts[1]

    if proto not in API_HANDLERS:
        res.status = b'404 Not Found'
        res.headers[b'Content-Type'] = b'text/plain'
        res.setbodybytes(_('Unknown API: %s\nKnown APIs: %s') % (
            proto, b', '.join(sorted(availableapis))))
        return

    if proto not in availableapis:
        res.status = b'404 Not Found'
        res.headers[b'Content-Type'] = b'text/plain'
        res.setbodybytes(_('API %s not enabled\n') % proto)
        return

    API_HANDLERS[proto]['handler'](rctx, req, res, checkperm,
                                   req.dispatchparts[2:])

# Maps API name to metadata so custom API can be registered.
# Keys are:
#
# config

        'handler': wireprotov2server.handlehttpv2request,
        'apidescriptor': wireprotov2server.httpv2apidescriptor,
    },
}

def _httpresponsetype(ui, proto, prefer_uncompressed):
    """Determine the appropriate response type and compression settings.

    Returns a tuple of (mediatype, compengine, engineopts).
    """

        if prefer_uncompressed:
            return HGTYPE2, compression._noopengine(), {}

        # Now find an agreed upon compression format.
        compformats = wireprotov1server.clientcompressionsupport(proto)
        for engine in wireprototypes.supportedcompengines(ui,
                compression.SERVERROLE):
            if engine.wireprotosupport().name in compformats:
                opts = {}
                level = ui.configint('server', '%slevel' % engine.name())
                if level is not None:
                    opts['level'] = level

    # Don't allow untrusted settings because disabling compression or
    # setting a very high compression level could lead to flooding
    # the server's network or CPU.
    opts = {'level': ui.configint('server', 'zliblevel')}
    return HGTYPE, util.compengines['zlib'], opts

def processcapabilitieshandshake(repo, req, res, proto):
    """Called during a ?cmd=capabilities request.

    If the client is advertising support for a newer protocol, we send

    res.headers[b'Content-Type'] = b'application/mercurial-cbor'
    res.setbodybytes(b''.join(cborutil.streamencode(m)))

    return True

def _callhttp(repo, req, res, proto, cmd):
    # Avoid cycle involving hg module.
    from .hgweb import common as hgwebcommon

    def genversion2(gen, engine, engineopts):

            res.setbodybytes(bodybytes)
        if bodygen is not None:
            res.setbodygen(bodygen)

    if not wireprotov1server.commands.commandavailable(cmd, proto):
        setresponse(HTTP_OK, HGERRTYPE,
                    _('requested wire protocol command is not available over '
                      'HTTP'))
        return

    proto.checkperm(wireprotov1server.commands[cmd].permission)

    # Possibly handle a modern client wanting to switch protocols.
    if (cmd == 'capabilities' and
        processcapabilitieshandshake(repo, req, res, proto)):

        return

    rsp = wireprotov1server.dispatch(repo, proto, cmd)

        gen = rsp.gen

        # This code for compression should not be streamres specific. It
        # is here because we only compress streamres at the moment.
        mediatype, engine, engineopts = _httpresponsetype(
            repo.ui, proto, rsp.prefer_uncompressed)
        gen = engine.compressstream(gen, engineopts)

        if mediatype == HGTYPE2:
            gen = genversion2(gen, engine, engineopts)

    elif isinstance(rsp, wireprototypes.ooberror):
        setresponse(HTTP_OK, HGERRTYPE, bodybytes=rsp.message)
    else:
        raise error.ProgrammingError('hgweb.protocol internal failure', rsp)

def _sshv1respondbytes(fout, value):
    """Send a bytes response for protocol version 1."""
    fout.write('%d\n' % len(value))
    fout.write(value)
    fout.flush()

def _sshv1respondstream(fout, source):
    write = fout.write
    for chunk in source.gen:
        write(chunk)
    fout.flush()

def _sshv1respondooberror(fout, ferr, rsp):
    ferr.write(b'%s\n-\n' % rsp)
    ferr.flush()
    fout.write(b'\n')
    fout.flush()

@interfaceutil.implementer(wireprototypes.baseprotocolhandler)
class sshv1protocolhandler(object):
    """Handler for requests services via version 1 of SSH protocol."""
    def __init__(self, ui, fin, fout):
        self._ui = ui
        self._fin = fin
        self._fout = fout
        self._protocaps = set()

        return caps

    def checkperm(self, perm):
        pass

class sshv2protocolhandler(sshv1protocolhandler):
    """Protocol handler for version 2 of the SSH protocol."""

    @property
    def name(self):
        return wireprototypes.SSHV2

    def addcapabilities(self, repo, caps):
        return caps

def _runsshserver(ui, repo, fin, fout, ev):
    # This function operates like a state machine of sorts. The following
    # states are defined:
    #

            # It looks like a protocol upgrade request. Transition state to
            # handle it.
            if request.startswith(b'upgrade '):
                if protoswitched:
                    _sshv1respondooberror(fout, ui.ferr,
                                          b'cannot upgrade protocols multiple '
                                          b'times')
                    state = 'shutdown'
                    continue

                state = 'upgrade-initial'
                continue

            available = wireprotov1server.commands.commandavailable(
                request, proto)

            # This command isn't available. Send an empty response and go
            # back to waiting for a new command.
            if not available:
                _sshv1respondbytes(fout, b'')

            elif isinstance(rsp, wireprototypes.pusherr):
                _sshv1respondbytes(fout, rsp.res)
            elif isinstance(rsp, wireprototypes.ooberror):
                _sshv1respondooberror(fout, ui.ferr, rsp.message)
            else:
                raise error.ProgrammingError('unhandled response type from '
                                             'wire protocol command: %s' % rsp)

        # For now, protocol version 2 serving just goes back to version 1.
        elif state == 'protov2-serving':
            state = 'protov1-serving'
            continue

            ok = True
            for line in (b'hello', b'between', b'pairs 81'):
                request = fin.readline()[:-1]

                if request != line:
                    _sshv1respondooberror(fout, ui.ferr,
                                          b'malformed handshake protocol: '
                                          b'missing %s' % line)
                    ok = False
                    state = 'shutdown'
                    break

            if not ok:
                continue

            request = fin.read(81)
            if request != b'%s-%s' % (b'0' * 40, b'0' * 40):
                _sshv1respondooberror(fout, ui.ferr,
                                      b'malformed handshake protocol: '
                                      b'missing between argument value')
                state = 'shutdown'
                continue

            state = 'upgrade-v2-finish'
            continue

        elif state == 'shutdown':
            break

        else:
            raise error.ProgrammingError('unhandled ssh server state: %s' %
                                         state)

class sshserver(object):
    def __init__(self, ui, repo, logfh=None):
        self._ui = ui
        self._repo = repo
        self._fin, self._fout = ui.protectfinout()

        # Log write I/O to stdout and stderr if configured.
        if logfh:
            self._fout = util.makeloggingfileobject(
                logfh, self._fout, 'o', logdata=True)
            ui.ferr = util.makeloggingfileobject(
                logfh, ui.ferr, 'e', logdata=True)

    def serve_forever(self):
        self.serveuntil(threading.Event())
        self._ui.restorefinout(self._fin, self._fout)
        sys.exit(0)

43076:2372284d9457

    util,
    wireprototypes,
    wireprotov1server,
    wireprotov2server,
)
from .interfaces import util as interfaceutil
from .utils import (
    cborutil,
    compression,
)

HGTYPE2 = 'application/mercurial-0.2'
HGERRTYPE = 'application/hg-error'

SSHV1 = wireprototypes.SSHV1
SSHV2 = wireprototypes.SSHV2


def decodevaluefromheaders(req, headerprefix):
    """Decode a long value from multiple HTTP request headers.

    Returns the value as a bytes, not a str.

            break
        chunks.append(pycompat.bytesurl(v))
        i += 1

    return ''.join(chunks)


@interfaceutil.implementer(wireprototypes.baseprotocolhandler)
class httpv1protocolhandler(object):
    def __init__(self, req, ui, checkperm):
        self._req = req

    def _args(self):
        args = self._req.qsparams.asdictoflists()
        postlen = int(self._req.headers.get(b'X-HgArgs-Post', 0))
        if postlen:
            args.update(
                urlreq.parseqs(
                    self._req.bodyfh.read(postlen), keep_blank_values=True
                )
            )
            return args

        argvalue = decodevaluefromheaders(self._req, b'X-HgArg')
        args.update(urlreq.parseqs(argvalue, keep_blank_values=True))
        return args

    def client(self):
        return 'remote:%s:%s:%s' % (
            self._req.urlscheme,
            urlreq.quote(self._req.remotehost or ''),
            urlreq.quote(self._req.remoteuser or ''),
        )

    def addcapabilities(self, repo, caps):
        caps.append(b'batch')

        caps.append(
            'httpheader=%d' % repo.ui.configint('server', 'maxhttpheaderlen')
        )
        if repo.ui.configbool('experimental', 'httppostargs'):
            caps.append('httppostargs')

        # FUTURE advertise 0.2rx once support is implemented
        # FUTURE advertise minrx and mintx after consulting config option
        caps.append('httpmediatype=0.1rx,0.1tx,0.2tx')

        compengines = wireprototypes.supportedcompengines(
            repo.ui, compression.SERVERROLE
        )
        if compengines:
            comptypes = ','.join(
                urlreq.quote(e.wireprotosupport().name) for e in compengines
            )
            caps.append('compression=%s' % comptypes)

        return caps

    def checkperm(self, perm):
        return self._checkperm(perm)


# This method exists mostly so that extensions like remotefilelog can
# disable a kludgey legacy method only over http. As of early 2018,
# there are no other known users, so with any luck we can discard this
# hook if remotefilelog becomes a first-party extension.
def iscmd(cmd):
    return cmd in wireprotov1server.commands


def handlewsgirequest(rctx, req, res, checkperm):
    """Possibly process a wire protocol request.

    If the current request is a wire protocol request, the request is

        # TODO This is not a good response to issue for this request. This
        # is mostly for BC for now.
        res.setbodybytes('0\n%s\n' % b'Not Found')
        return True

    proto = httpv1protocolhandler(
        req, repo.ui, lambda perm: checkperm(rctx, req, perm)
    )

    # The permissions checker should be the only thing that can raise an
    # ErrorResponse. It is kind of a layer violation to catch an hgweb
    # exception here. So consider refactoring into a exception type that
    # is associated with the wire protocol.

        # "unbundle." That assumption is not always valid.
        res.setbodybytes('0\n%s\n' % pycompat.bytestr(e))

    return True


def _availableapis(repo):
    apis = set()

    # Registered APIs are made available via config options of the name of
    # the protocol.

        section, option = v['config']
        if repo.ui.configbool(section, option):
            apis.add(k)

    return apis


def handlewsgiapirequest(rctx, req, res, checkperm):
    """Handle requests to /api/*."""
    assert req.dispatchparts[0] == b'api'

    # Requests to /api/ list available APIs.
    if req.dispatchparts == [b'api']:
        res.status = b'200 OK'
        res.headers[b'Content-Type'] = b'text/plain'
        lines = [
            _(
                'APIs can be accessed at /api/<name>, where <name> can be '
                'one of the following:\n'
            )
        ]
        if availableapis:
            lines.extend(sorted(availableapis))
        else:
            lines.append(_('(no available APIs)\n'))
        res.setbodybytes(b'\n'.join(lines))

    proto = req.dispatchparts[1]

    if proto not in API_HANDLERS:
        res.status = b'404 Not Found'
        res.headers[b'Content-Type'] = b'text/plain'
        res.setbodybytes(
            _('Unknown API: %s\nKnown APIs: %s')
            % (proto, b', '.join(sorted(availableapis)))
        )
        return

    if proto not in availableapis:
        res.status = b'404 Not Found'
        res.headers[b'Content-Type'] = b'text/plain'
        res.setbodybytes(_('API %s not enabled\n') % proto)
        return

    API_HANDLERS[proto]['handler'](
        rctx, req, res, checkperm, req.dispatchparts[2:]
    )


# Maps API name to metadata so custom API can be registered.
# Keys are:
#
# config

        'handler': wireprotov2server.handlehttpv2request,
        'apidescriptor': wireprotov2server.httpv2apidescriptor,
    },
}


def _httpresponsetype(ui, proto, prefer_uncompressed):
    """Determine the appropriate response type and compression settings.

    Returns a tuple of (mediatype, compengine, engineopts).
    """

        if prefer_uncompressed:
            return HGTYPE2, compression._noopengine(), {}

        # Now find an agreed upon compression format.
        compformats = wireprotov1server.clientcompressionsupport(proto)
        for engine in wireprototypes.supportedcompengines(
            ui, compression.SERVERROLE
        ):
            if engine.wireprotosupport().name in compformats:
                opts = {}
                level = ui.configint('server', '%slevel' % engine.name())
                if level is not None:
                    opts['level'] = level

    # Don't allow untrusted settings because disabling compression or
    # setting a very high compression level could lead to flooding
    # the server's network or CPU.
    opts = {'level': ui.configint('server', 'zliblevel')}
    return HGTYPE, util.compengines['zlib'], opts


def processcapabilitieshandshake(repo, req, res, proto):
    """Called during a ?cmd=capabilities request.

    If the client is advertising support for a newer protocol, we send

    res.headers[b'Content-Type'] = b'application/mercurial-cbor'
    res.setbodybytes(b''.join(cborutil.streamencode(m)))

    return True


def _callhttp(repo, req, res, proto, cmd):
    # Avoid cycle involving hg module.
    from .hgweb import common as hgwebcommon

    def genversion2(gen, engine, engineopts):

            res.setbodybytes(bodybytes)
        if bodygen is not None:
            res.setbodygen(bodygen)

    if not wireprotov1server.commands.commandavailable(cmd, proto):
        setresponse(
            HTTP_OK,
            HGERRTYPE,
            _('requested wire protocol command is not available over ' 'HTTP'),
        )
        return

    proto.checkperm(wireprotov1server.commands[cmd].permission)

    # Possibly handle a modern client wanting to switch protocols.
    if cmd == 'capabilities' and processcapabilitieshandshake(
        repo, req, res, proto
    ):

        return

    rsp = wireprotov1server.dispatch(repo, proto, cmd)

        gen = rsp.gen

        # This code for compression should not be streamres specific. It
        # is here because we only compress streamres at the moment.
        mediatype, engine, engineopts = _httpresponsetype(
            repo.ui, proto, rsp.prefer_uncompressed
        )
        gen = engine.compressstream(gen, engineopts)

        if mediatype == HGTYPE2:
            gen = genversion2(gen, engine, engineopts)

    elif isinstance(rsp, wireprototypes.ooberror):
        setresponse(HTTP_OK, HGERRTYPE, bodybytes=rsp.message)
    else:
        raise error.ProgrammingError('hgweb.protocol internal failure', rsp)


def _sshv1respondbytes(fout, value):
    """Send a bytes response for protocol version 1."""
    fout.write('%d\n' % len(value))
    fout.write(value)
    fout.flush()


def _sshv1respondstream(fout, source):
    write = fout.write
    for chunk in source.gen:
        write(chunk)
    fout.flush()


def _sshv1respondooberror(fout, ferr, rsp):
    ferr.write(b'%s\n-\n' % rsp)
    ferr.flush()
    fout.write(b'\n')
    fout.flush()


@interfaceutil.implementer(wireprototypes.baseprotocolhandler)
class sshv1protocolhandler(object):
    """Handler for requests services via version 1 of SSH protocol."""

    def __init__(self, ui, fin, fout):
        self._ui = ui
        self._fin = fin
        self._fout = fout
        self._protocaps = set()

        return caps

    def checkperm(self, perm):
        pass


class sshv2protocolhandler(sshv1protocolhandler):
    """Protocol handler for version 2 of the SSH protocol."""

    @property
    def name(self):
        return wireprototypes.SSHV2

    def addcapabilities(self, repo, caps):
        return caps


def _runsshserver(ui, repo, fin, fout, ev):
    # This function operates like a state machine of sorts. The following
    # states are defined:
    #

            # It looks like a protocol upgrade request. Transition state to
            # handle it.
            if request.startswith(b'upgrade '):
                if protoswitched:
                    _sshv1respondooberror(
                        fout,
                        ui.ferr,
                        b'cannot upgrade protocols multiple ' b'times',
                    )
                    state = 'shutdown'
                    continue

                state = 'upgrade-initial'
                continue

            available = wireprotov1server.commands.commandavailable(
                request, proto
            )

            # This command isn't available. Send an empty response and go
            # back to waiting for a new command.
            if not available:
                _sshv1respondbytes(fout, b'')

            elif isinstance(rsp, wireprototypes.pusherr):
                _sshv1respondbytes(fout, rsp.res)
            elif isinstance(rsp, wireprototypes.ooberror):
                _sshv1respondooberror(fout, ui.ferr, rsp.message)
            else:
                raise error.ProgrammingError(
                    'unhandled response type from '
                    'wire protocol command: %s' % rsp
                )

        # For now, protocol version 2 serving just goes back to version 1.
        elif state == 'protov2-serving':
            state = 'protov1-serving'
            continue

            ok = True
            for line in (b'hello', b'between', b'pairs 81'):
                request = fin.readline()[:-1]

                if request != line:
                    _sshv1respondooberror(
                        fout,
                        ui.ferr,
                        b'malformed handshake protocol: ' b'missing %s' % line,
                    )
                    ok = False
                    state = 'shutdown'
                    break

            if not ok:
                continue

            request = fin.read(81)
            if request != b'%s-%s' % (b'0' * 40, b'0' * 40):
                _sshv1respondooberror(
                    fout,
                    ui.ferr,
                    b'malformed handshake protocol: '
                    b'missing between argument value',
                )
                state = 'shutdown'
                continue

            state = 'upgrade-v2-finish'
            continue

        elif state == 'shutdown':
            break

        else:
            raise error.ProgrammingError(
                'unhandled ssh server state: %s' % state
            )


class sshserver(object):
    def __init__(self, ui, repo, logfh=None):
        self._ui = ui
        self._repo = repo
        self._fin, self._fout = ui.protectfinout()

        # Log write I/O to stdout and stderr if configured.
        if logfh:
            self._fout = util.makeloggingfileobject(
                logfh, self._fout, 'o', logdata=True
            )
            ui.ferr = util.makeloggingfileobject(
                logfh, ui.ferr, 'e', logdata=True
            )

    def serve_forever(self):
        self.serveuntil(threading.Event())
        self._ui.restorefinout(self._fin, self._fout)
        sys.exit(0)