diff --git a/kitty/kittens.c b/kitty/kittens.c index 0ef743704..7e3196da3 100644 --- a/kitty/kittens.c +++ b/kitty/kittens.c @@ -7,6 +7,92 @@ #include "data-types.h" +#define CMD_BUF_SZ 2048 + + +static inline bool +append_buf(char buf[CMD_BUF_SZ], size_t *pos, PyObject *ans) { + if (*pos) { + PyObject *bytes = PyBytes_FromStringAndSize(buf, *pos); + if (!bytes) { PyErr_NoMemory(); return false; } + int ret = PyList_Append(ans, bytes); + Py_CLEAR(bytes); + if (ret != 0) return false; + *pos = 0; + } + return true; +} + +static inline bool +add_char(char buf[CMD_BUF_SZ], size_t *pos, char ch, PyObject *ans) { + if (*pos >= CMD_BUF_SZ) { + if (!append_buf(buf, pos, ans)) return false; + } + buf[*pos] = ch; + *pos += 1; + return true; +} + +static inline bool +read_response(int fd, double timeout, char buf[CMD_BUF_SZ], size_t *pos, PyObject *ans) { + enum ReadState {START, STARTING_ESC, P, AT, K, I, T, T2, Y, HYPHEN, C, M, BODY, TRAILING_ESC}; + enum ReadState state = START; + char ch; + double end_time = monotonic() + timeout; + while(monotonic() <= end_time) { + ssize_t len = read(fd, &ch, 1); + if (len == 0) continue; + if (len < 0) { + if (errno == EINTR || errno == EAGAIN) continue; + PyErr_SetFromErrno(PyExc_OSError); + return false; + continue; + } + end_time = monotonic() + timeout; + switch(state) { + case START: + if (ch == 0x1b) state = STARTING_ESC; + break; +#define CASE(curr, q, next) case curr: state = ch == q ? next : START; break; + CASE(STARTING_ESC, 'P', P); + CASE(P, '@', AT); + CASE(AT, 'k', K); + CASE(K, 'i', I); + CASE(I, 't', T); + CASE(T, 't', T2); + CASE(T2, 'y', Y); + CASE(Y, '-', HYPHEN); + CASE(HYPHEN, 'c', C); + CASE(C, 'm', M); + CASE(M, 'd', BODY); + case BODY: + if (ch == 0x1b) { state = TRAILING_ESC; } + else { + if (!add_char(buf, pos, ch, ans)) return false; + } + break; + case TRAILING_ESC: + if (ch == '\\') return append_buf(buf, pos, ans); + state = BODY; + break; + } + } + PyErr_SetString(PyExc_TimeoutError, "Timed out while waiting to read cmd response"); + return false; +} + +static PyObject* +read_command_response(PyObject *self UNUSED, PyObject *args) { + double timeout; + int fd; + PyObject *ans; + if (!PyArg_ParseTuple(args, "idO!", &fd, &timeout, &PyList_Type, &ans)) return NULL; + static char buf[CMD_BUF_SZ]; + size_t pos = 0; + if (!read_response(fd, timeout, buf, &pos, ans)) return NULL; + Py_RETURN_NONE; +} + static PyObject* parse_input_from_terminal(PyObject *self UNUSED, PyObject *args) { enum State { NORMAL, ESC, CSI, ST, ESC_ST }; @@ -105,6 +191,7 @@ parse_input_from_terminal(PyObject *self UNUSED, PyObject *args) { static PyMethodDef module_methods[] = { METHODB(parse_input_from_terminal, METH_VARARGS), + METHODB(read_command_response, METH_VARARGS), {NULL, NULL, 0, NULL} /* Sentinel */ }; diff --git a/kitty/remote_control.py b/kitty/remote_control.py index 87a19f2be..1ed56ec0f 100644 --- a/kitty/remote_control.py +++ b/kitty/remote_control.py @@ -11,6 +11,7 @@ from functools import partial from .cli import emph, parse_args from .cmds import cmap, parse_subcommand_cli from .constants import appname, version +from .fast_data_types import read_command_response from .utils import TTYIO, parse_address_spec @@ -73,13 +74,23 @@ class SocketIO: out.flush() self.socket.shutdown(socket.SHUT_WR) - def recv(self, more_needed, timeout): - # We dont bother with more_needed since the server will close the - # connection after transmission + def recv(self, timeout): + dcs = re.compile(br'\x1bP@kitty-cmd([^\x1b]+)\x1b\\') self.socket.settimeout(timeout) with self.socket.makefile('rb') as src: data = src.read() - more_needed(data) + m = dcs.search(data) + if m is None: + raise TimeoutError('Timed out while waiting to read cmd response') + return m.group(1) + + +class RCIO(TTYIO): + + def recv(self, timeout): + ans = [] + read_command_response(self.tty_fd, timeout, ans) + return b''.join(ans) def do_io(to, send, no_response): @@ -93,28 +104,14 @@ def do_io(to, send, no_response): yield encode_send(send) send_data = send_generator() - io = SocketIO(to) if to else TTYIO() + io = SocketIO(to) if to else RCIO() with io: io.send(send_data) if no_response: return {'ok': True} + received = io.recv(timeout=10) - dcs = re.compile(br'\x1bP@kitty-cmd([^\x1b]+)\x1b\\') - - received = b'' - match = None - - def more_needed(data): - nonlocal received, match - received += data - match = dcs.search(received) - return match is None - - io.recv(more_needed, timeout=10) - - if match is None: - raise SystemExit('Failed to receive response from ' + appname) - response = json.loads(match.group(1).decode('ascii')) + response = json.loads(received.decode('ascii')) return response