# commandserver.py - communicate with Mercurial's API over a pipe

#

#  Copyright Olivia Mackall <olivia@selenic.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

import gc

import os

import random

import selectors

import signal

import socket

import struct

import traceback

from .i18n import _

from . import (

    encoding,

    error,

    loggingutil,

    pycompat,

    repocache,

    util,

    vfs as vfsmod,

)

from .utils import (

    cborutil,

    procutil,

)

class channeledoutput:

    """

    Write data to out in the following format:

    data length (unsigned int),

    data

    """

    def __init__(self, out, channel):

        self.out = out

        self.channel = channel

    @property

    def name(self):

        return b'<%c-channel>' % self.channel

    def write(self, data):

        if not data:

            return

        # single write() to guarantee the same atomicity as the underlying file

        self.out.write(struct.pack(b'>cI', self.channel, len(data)) + data)

        self.out.flush()

    def __getattr__(self, attr):

        if attr in ('isatty', 'fileno', 'tell', 'seek'):

            raise AttributeError(attr)

        return getattr(self.out, attr)

class channeledmessage:

    """

    Write encoded message and metadata to out in the following format:

    data length (unsigned int),

    encoded message and metadata, as a flat key-value dict.

    Each message should have 'type' attribute. Messages of unknown type

    should be ignored.

    """

    # teach ui that write() can take **opts

    structured = True

    def __init__(self, out, channel, encodename, encodefn):

        self._cout = channeledoutput(out, channel)

        self.encoding = encodename

        self._encodefn = encodefn

    def write(self, data, **opts):

        opts = pycompat.byteskwargs(opts)

        if data is not None:

            opts[b'data'] = data

        self._cout.write(self._encodefn(opts))

    def __getattr__(self, attr):

        return getattr(self._cout, attr)

class channeledinput:

    """

    Read data from in_.

    Requests for input are written to out in the following format:

    channel identifier - 'I' for plain input, 'L' line based (1 byte)

    how many bytes to send at most (unsigned int),

    The client replies with:

    data length (unsigned int), 0 meaning EOF

    data

    """

    maxchunksize = 4 * 1024

    def __init__(self, in_, out, channel):

        self.in_ = in_

        self.out = out

        self.channel = channel

    @property

    def name(self):

        return b'<%c-channel>' % self.channel

    def read(self, size=-1):

        if size < 0:

            # if we need to consume all the clients input, ask for 4k chunks

            # so the pipe doesn't fill up risking a deadlock

            size = self.maxchunksize

            s = self._read(size, self.channel)

            buf = s

            while s:

                s = self._read(size, self.channel)

                buf += s

            return buf

        else:

            return self._read(size, self.channel)

    def _read(self, size, channel):

        if not size:

            return b''

        assert size > 0

        # tell the client we need at most size bytes

        self.out.write(struct.pack(b'>cI', channel, size))

        self.out.flush()

        length = self.in_.read(4)

        length = struct.unpack(b'>I', length)[0]

        if not length:

            return b''

        else:

            return self.in_.read(length)

    def readline(self, size=-1):

        if size < 0:

            size = self.maxchunksize

            s = self._read(size, b'L')

            buf = s

            # keep asking for more until there's either no more or

            # we got a full line

            while s and not s.endswith(b'\n'):

                s = self._read(size, b'L')

                buf += s

            return buf

        else:

            return self._read(size, b'L')

    def __iter__(self):

        return self

    def next(self):

        l = self.readline()

        if not l:

            raise StopIteration

        return l

    __next__ = next

    def __getattr__(self, attr):

        if attr in ('isatty', 'fileno', 'tell', 'seek'):

            raise AttributeError(attr)

        return getattr(self.in_, attr)

_messageencoders = {

    b'cbor': lambda v: b''.join(cborutil.streamencode(v)),

}

def _selectmessageencoder(ui):

    encnames = ui.configlist(b'cmdserver', b'message-encodings')

    for n in encnames:

        f = _messageencoders.get(n)

        if f:

            return n, f

    raise error.Abort(

        b'no supported message encodings: %s' % b' '.join(encnames)

    )

class server:

    """

    Listens for commands on fin, runs them and writes the output on a channel

    based stream to fout.

    """

    def __init__(self, ui, repo, fin, fout, prereposetups=None):

        self.cwd = encoding.getcwd()

        if repo:

            # the ui here is really the repo ui so take its baseui so we don't

            # end up with its local configuration

            self.ui = repo.baseui

            self.repo = repo

            self.repoui = repo.ui

        else:

            self.ui = ui

            self.repo = self.repoui = None

        self._prereposetups = prereposetups

        self.cdebug = channeledoutput(fout, b'd')

        self.cerr = channeledoutput(fout, b'e')

        self.cout = channeledoutput(fout, b'o')

        self.cin = channeledinput(fin, fout, b'I')

        self.cresult = channeledoutput(fout, b'r')

        if self.ui.config(b'cmdserver', b'log') == b'-':

            # switch log stream of server's ui to the 'd' (debug) channel

            # (don't touch repo.ui as its lifetime is longer than the server)

            self.ui = self.ui.copy()

            setuplogging(self.ui, repo=None, fp=self.cdebug)

        self.cmsg = None

        if ui.config(b'ui', b'message-output') == b'channel':

            encname, encfn = _selectmessageencoder(ui)

            self.cmsg = channeledmessage(fout, b'm', encname, encfn)

        self.client = fin

        # If shutdown-on-interrupt is off, the default SIGINT handler is

        # removed so that client-server communication wouldn't be interrupted.

        # For example, 'runcommand' handler will issue three short read()s.

        # If one of the first two read()s were interrupted, the communication

        # channel would be left at dirty state and the subsequent request

        # wouldn't be parsed. So catching KeyboardInterrupt isn't enough.

        self._shutdown_on_interrupt = ui.configbool(

            b'cmdserver', b'shutdown-on-interrupt'

        )

        self._old_inthandler = None

        if not self._shutdown_on_interrupt:

            self._old_inthandler = signal.signal(signal.SIGINT, signal.SIG_IGN)

    def cleanup(self):

        """release and restore resources taken during server session"""

        if not self._shutdown_on_interrupt:

            signal.signal(signal.SIGINT, self._old_inthandler)

    def _read(self, size):

        if not size:

            return b''

        data = self.client.read(size)

        # is the other end closed?

        if not data:

            raise EOFError

        return data

    def _readstr(self):

        """read a string from the channel

        format:

        data length (uint32), data

        """

        length = struct.unpack(b'>I', self._read(4))[0]

        if not length:

            return b''

        return self._read(length)

    def _readlist(self):

        """read a list of NULL separated strings from the channel"""

        s = self._readstr()

        if s:

            return s.split(b'\0')

        else:

            return []

    def _dispatchcommand(self, req):

        from . import dispatch  # avoid cycle

        if self._shutdown_on_interrupt:

            # no need to restore SIGINT handler as it is unmodified.

            return dispatch.dispatch(req)

        try:

            signal.signal(signal.SIGINT, self._old_inthandler)

            return dispatch.dispatch(req)

        except error.SignalInterrupt:

            # propagate SIGBREAK, SIGHUP, or SIGTERM.

            raise

        except KeyboardInterrupt:

            # SIGINT may be received out of the try-except block of dispatch(),

            # so catch it as last ditch. Another KeyboardInterrupt may be

            # raised while handling exceptions here, but there's no way to

            # avoid that except for doing everything in C.

            pass

        finally:

            signal.signal(signal.SIGINT, signal.SIG_IGN)

        # On KeyboardInterrupt, print error message and exit *after* SIGINT

        # handler removed.

        req.ui.error(_(b'interrupted!\n'))

        return -1

    def runcommand(self):

        """reads a list of \0 terminated arguments, executes

        and writes the return code to the result channel"""

        from . import dispatch  # avoid cycle

        args = self._readlist()

        # copy the uis so changes (e.g. --config or --verbose) don't

        # persist between requests

        copiedui = self.ui.copy()

        uis = [copiedui]

        if self.repo:

            self.repo.baseui = copiedui

            # clone ui without using ui.copy because this is protected

            repoui = self.repoui.__class__(self.repoui)

            repoui.copy = copiedui.copy  # redo copy protection

            uis.append(repoui)

            self.repo.ui = self.repo.dirstate._ui = repoui

            self.repo.invalidateall()

        for ui in uis:

            ui.resetstate()

            # any kind of interaction must use server channels, but chg may

            # replace channels by fully functional tty files. so nontty is

            # enforced only if cin is a channel.

            if not hasattr(self.cin, 'fileno'):

                ui.setconfig(b'ui', b'nontty', b'true', b'commandserver')

        req = dispatch.request(

            args[:],

            copiedui,

            self.repo,

            self.cin,

            self.cout,

            self.cerr,

            self.cmsg,

            prereposetups=self._prereposetups,

        )

        try:

            ret = self._dispatchcommand(req) & 255

            # If shutdown-on-interrupt is off, it's important to write the

            # result code *after* SIGINT handler removed. If the result code

            # were lost, the client wouldn't be able to continue processing.

            self.cresult.write(struct.pack(b'>i', int(ret)))

        finally:

            # restore old cwd

            if b'--cwd' in args:

                os.chdir(self.cwd)

    def getencoding(self):

        """writes the current encoding to the result channel"""

        self.cresult.write(encoding.encoding)

    def serveone(self):

        cmd = self.client.readline()[:-1]

        if cmd:

            handler = self.capabilities.get(cmd)

            if handler:

                handler(self)

            else:

                # clients are expected to check what commands are supported by

                # looking at the servers capabilities

                raise error.Abort(_(b'unknown command %s') % cmd)

        return cmd != b''

    capabilities = {b'runcommand': runcommand, b'getencoding': getencoding}

    def serve(self):

        hellomsg = b'capabilities: ' + b' '.join(sorted(self.capabilities))

        hellomsg += b'\n'

        hellomsg += b'encoding: ' + encoding.encoding

        hellomsg += b'\n'

        if self.cmsg:

            hellomsg += b'message-encoding: %s\n' % self.cmsg.encoding

        hellomsg += b'pid: %d' % procutil.getpid()

        if hasattr(os, 'getpgid'):

            hellomsg += b'\n'

            hellomsg += b'pgid: %d' % os.getpgid(0)

        # write the hello msg in -one- chunk

        self.cout.write(hellomsg)

        try: