diff --git a/kitty/clipboard.py b/kitty/clipboard.py index 8776244f6..53a490c13 100644 --- a/kitty/clipboard.py +++ b/kitty/clipboard.py @@ -279,10 +279,6 @@ class WriteRequest: self.currently_writing_mime = mime self.write_base64_data(data) - @property - def current_leftover_bytes(self) -> memoryview: - return self.decoder.leftover_bytes() - def flush_base64_data(self) -> None: if self.currently_writing_mime: self.decoder.flush() diff --git a/kitty/data-types.c b/kitty/data-types.c index 6b5a7bb16..d210a1ec6 100644 --- a/kitty/data-types.c +++ b/kitty/data-types.c @@ -108,8 +108,8 @@ pybase64_decode(PyObject UNUSED *self, PyObject *args) { typedef struct StreamingBase64Decoder { PyObject_HEAD PyObject *output; - size_t output_sz, output_capacity, num_leftover_bytes, initial_capacity; - unsigned char leftover_bytes[8]; + size_t output_sz, output_capacity, initial_capacity; + struct base64_state state; } StreamingBase64Decoder; static int @@ -119,6 +119,7 @@ StreamingBase64Decoder_init(PyObject *s, PyObject *args, PyObject *kwds) { if (!PyArg_ParseTupleAndKeywords(args, kwds, "|k", kwlist, &initial_capacity)) return -1; StreamingBase64Decoder *self = (StreamingBase64Decoder*)s; self->initial_capacity = initial_capacity; + base64_stream_decode_init(&self->state, 0); return 0; } @@ -130,16 +131,17 @@ StreamingBase64Decoder_dealloc(PyObject *self) { } static bool -write_base64_data(StreamingBase64Decoder *self, const void *data, size_t len) { +write_base64_data(StreamingBase64Decoder *self, const void *data, const size_t len) { if (!len) return true; size_t sz = required_buffer_size_for_base64_decode(len); if ((self->output_sz + sz) > self->output_capacity) { - size_t cap = MAX(self->output_capacity * 2, self->output_sz + sz + self->initial_capacity); + size_t cap = MAX(self->output_capacity * 2, MAX(self->output_sz + sz, self->initial_capacity)); if (self->output) { if (_PyBytes_Resize(&self->output, cap) != 0) return false; } else { self->output = PyBytes_FromStringAndSize(NULL, cap); if (!self->output) return false; } self->output_capacity = cap; } - if (!base64_decode8(data, len, (unsigned char*)(PyBytes_AS_STRING(self->output) + self->output_sz), &sz)) { + bool ok = base64_stream_decode(&self->state, data, len, PyBytes_AS_STRING(self->output) + self->output_sz, &sz); + if (!ok) { PyErr_SetString(PyExc_ValueError, "Invalid base64 input data"); return false; } @@ -147,55 +149,22 @@ write_base64_data(StreamingBase64Decoder *self, const void *data, size_t len) { return true; } -static bool -write_saving_leftover_bytes(StreamingBase64Decoder *self, const unsigned char *data, size_t len) { - size_t extra = len % 4; - if (!write_base64_data(self, data, len - extra)) return false; - self->num_leftover_bytes = extra; - if (extra) memcpy(self->leftover_bytes, data + len - extra, extra); - return true; -} - static PyObject* StreamingBase64Decoder_add(StreamingBase64Decoder *self, PyObject *a) { RAII_PY_BUFFER(data); if (PyObject_GetBuffer(a, &data, PyBUF_SIMPLE) != 0) return NULL; if (!data.buf || !data.len) return PyLong_FromLong(0); - unsigned char *d = data.buf; size_t dlen = data.len; size_t before = self->output_sz; - if (self->num_leftover_bytes) { - size_t extra = 4 - self->num_leftover_bytes; - if (dlen >= extra) { - memcpy(self->leftover_bytes + self->num_leftover_bytes, d, extra); - if (!write_base64_data(self, self->leftover_bytes, self->num_leftover_bytes + extra)) return NULL; - self->num_leftover_bytes = 0; - d += extra; dlen -= extra; - if (!write_saving_leftover_bytes(self, d, dlen)) return NULL; - } else { - memcpy(self->leftover_bytes + self->num_leftover_bytes, d, dlen); - self->num_leftover_bytes += dlen; - } - } else if (!write_saving_leftover_bytes(self, d, dlen)) return NULL; + if (!write_base64_data(self, data.buf, data.len)) return NULL; return PyLong_FromSize_t(self->output_sz - before); } static Py_ssize_t StreamingBase64Decoder_len(PyObject *s) { return ((StreamingBase64Decoder*)s)->output_sz; } -static PyObject* -StreamingBase64Decoder_leftover_bytes(StreamingBase64Decoder *self, PyObject *a UNUSED) { - return PyMemoryView_FromMemory((char*)self->leftover_bytes, self->num_leftover_bytes, PyBUF_READ); -} - static PyObject* StreamingBase64Decoder_flush(StreamingBase64Decoder *self, PyObject *args UNUSED) { - size_t padding = 4 - self->num_leftover_bytes; - switch(padding) { - case 1: self->leftover_bytes[self->num_leftover_bytes++] = '='; break; - case 2: self->leftover_bytes[self->num_leftover_bytes++] = '='; self->leftover_bytes[self->num_leftover_bytes++] = '='; break; - } - write_base64_data(self, self->leftover_bytes, self->num_leftover_bytes); - self->num_leftover_bytes = 0; + base64_stream_decode_init(&self->state, 0); Py_RETURN_NONE; } @@ -207,11 +176,9 @@ StreamingBase64Decoder_copy_output(StreamingBase64Decoder *self, PyObject *args static PyObject* StreamingBase64Decoder_take_output(StreamingBase64Decoder *self, PyObject *args UNUSED) { if (!self->output_sz) return PyBytes_FromStringAndSize(NULL, 0); - RAII_PyObject(newbuf, PyBytes_FromStringAndSize(NULL, self->initial_capacity)); - if (!newbuf) return NULL; if (_PyBytes_Resize(&self->output, self->output_sz) != 0) return NULL; PyObject *ans = self->output; - self->output = Py_NewRef(newbuf); self->output_sz = 0; self->output_capacity = self->initial_capacity; + self->output = NULL; self->output_sz = 0; self->output_capacity = 0; return ans; } @@ -227,7 +194,6 @@ static PyTypeObject StreamingBase64Decoder_Type = { {"flush", (PyCFunction)StreamingBase64Decoder_flush, METH_NOARGS, ""}, {"take_output", (PyCFunction)StreamingBase64Decoder_take_output, METH_NOARGS, ""}, {"copy_output", (PyCFunction)StreamingBase64Decoder_copy_output, METH_NOARGS, ""}, - {"leftover_bytes", (PyCFunction)StreamingBase64Decoder_leftover_bytes, METH_NOARGS, ""}, {NULL, NULL, 0, NULL}, }, .tp_new = PyType_GenericNew, diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi index 26192985e..6077325fb 100644 --- a/kitty/fast_data_types.pyi +++ b/kitty/fast_data_types.pyi @@ -1711,7 +1711,6 @@ class StreamingBase64Decoder: def take_output(self) -> bytes: ... # take the output so far. The decoder no longer references this output def copy_output(self) -> bytes: ... # copy the output so far def __len__(self) -> int: ... # return the length of the current output - def leftover_bytes(self) -> memoryview: ... # return the currently leftover bytes that will be consumed by flush() class DiskCache: diff --git a/kitty_tests/clipboard.py b/kitty_tests/clipboard.py index a0b6edf3b..a1521ea4d 100644 --- a/kitty_tests/clipboard.py +++ b/kitty_tests/clipboard.py @@ -10,14 +10,17 @@ from . import BaseTest class TestClipboard(BaseTest): def test_clipboard_write_request(self): - wr = WriteRequest(max_size=64) - wr.add_base64_data('bGlnaHQgd29yaw') - self.ae(bytes(wr.current_leftover_bytes), b'aw') - wr.flush_base64_data() - self.ae(wr.data_for(), b'light work') - wr = WriteRequest(max_size=64) - wr.add_base64_data('bGlnaHQgd29yaw==') - self.ae(wr.data_for(), b'light work') + def t(data, expected): + wr = WriteRequest(max_size=64) + wr.add_base64_data(data) + self.ae(wr.data_for(), expected) + t('dGl0bGU=', b'title') + t('dGl0bGU', b'title') + t('dGl0bG', b'titl') + t('dGl0bG==', b'titl') + t('dGl0b', b'tit') + t('bGlnaHQgd29yaw', b'light work') + t('bGlnaHQgd29yaw==', b'light work') wr = WriteRequest(max_size=64) wr.add_base64_data('bGlnaHQgd29') for x in b'y', b'a', b'y', b'4', b'=':