Fix a regression that broke writing to the clipboard using the OSC 5522 protocol

Fixes #7858

docs/changelog.rst

@@ -85,6 +85,8 @@ Detailed list of changes
 
 - Fix a regression when tinting of background images was introduced that caused window borders to have :opt:`background_opacity` applied to them (:iss:`7850`)
 
+- Fix a regression that broke writing to the clipboard using the OSC 5522 protocol (:iss:`7858`)
+
 0.36.2 [2024-09-06]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

kitty/clipboard.py

@@ -282,6 +282,8 @@ class WriteRequest:
 
     def flush_base64_data(self) -> None:
         if self.currently_writing_mime:
+            if self.decoder.needs_more_data():
+                log_error('Received incomplete data for clipboard')
             self.decoder.reset()
             start = self.mime_map[self.currently_writing_mime][0]
             self.mime_map[self.currently_writing_mime] = MimePos(start, self.tempfile.tell() - start)
@@ -291,8 +293,8 @@ class WriteRequest:
         if not self.max_size_exceeded:
             try:
                 decoded = self.decoder.decode(b)
-            except ValueError:
-                log_error('Clipboard write request has invalid data, ignoring this chunk of data')
+            except ValueError as e:
+                log_error(f'Clipboard write request has invalid data, ignoring this chunk of data. Error: {e}')
                 self.decoder.reset()
                 decoded = b''
             if decoded:

kitty/data-types.c

@@ -137,9 +137,15 @@ base64_decode_into(PyObject UNUSED *self, PyObject *args) {
 typedef struct StreamingBase64Decoder {
     PyObject_HEAD
     struct base64_state state;
-    bool add_trailing_bytes;
+    bool add_trailing_bytes, needs_more_data;
 } StreamingBase64Decoder;
 
+static void
+StreamingBase64Decoder_reset_(StreamingBase64Decoder *self) {
+    base64_stream_decode_init(&self->state, 0);
+    self->needs_more_data = false;
+}
+
 static int
 StreamingBase64Decoder_init(PyObject *s, PyObject *args, PyObject *kwds UNUSED) {
     if (PyTuple_GET_SIZE(args)) { PyErr_SetString(PyExc_TypeError, "constructor takes no arguments"); return -1; }
@@ -156,10 +162,17 @@ StreamingBase64Decoder_decode(StreamingBase64Decoder *self, PyObject *a) {
     size_t sz = required_buffer_size_for_base64_decode(data.len);
     RAII_PyObject(ans, PyBytes_FromStringAndSize(NULL, sz));
     if (!ans) return NULL;
-    if (!base64_stream_decode(&self->state, data.buf, data.len, PyBytes_AS_STRING(ans), &sz)) {
+    int ret;
+    Py_BEGIN_ALLOW_THREADS
+    ret = base64_stream_decode(&self->state, data.buf, data.len, PyBytes_AS_STRING(ans), &sz);
+    Py_END_ALLOW_THREADS;
+    if (!ret) {
+        StreamingBase64Decoder_reset_(self);
         PyErr_SetString(PyExc_ValueError, "Invalid base64 input data");
         return NULL;
     }
+    if (self->state.eof) StreamingBase64Decoder_reset_(self);
+    else self->needs_more_data = self->state.carry != 0 || self->state.bytes != 0;
     if (_PyBytes_Resize(&ans, sz) != 0) return NULL;
     return Py_NewRef(ans);
 }
@@ -175,19 +188,30 @@ StreamingBase64Decoder_decode_into(StreamingBase64Decoder *self, PyObject *const
     if (!src.buf || !src.len) return PyLong_FromLong(0);
     size_t sz = required_buffer_size_for_base64_decode(src.len);
     if ((Py_ssize_t)sz > data.len) { PyErr_SetString(PyExc_BufferError, "output buffer too small"); return NULL; }
-    if (!base64_stream_decode(&self->state, src.buf, src.len, data.buf, &sz)) {
+    int ret;
+    Py_BEGIN_ALLOW_THREADS
+    ret = base64_stream_decode(&self->state, src.buf, src.len, data.buf, &sz);
+    Py_END_ALLOW_THREADS
+    if (!ret) {
+        StreamingBase64Decoder_reset_(self);
         PyErr_SetString(PyExc_ValueError, "Invalid base64 input data");
         return NULL;
     }
+    if (self->state.eof) StreamingBase64Decoder_reset_(self); else self->needs_more_data = true;
     return PyLong_FromSize_t(sz);
 }
 
 static PyObject*
 StreamingBase64Decoder_reset(StreamingBase64Decoder *self, PyObject *args UNUSED) {
-    base64_stream_decode_init(&self->state, 0);
+    StreamingBase64Decoder_reset_(self);
     Py_RETURN_NONE;
 }
 
+static PyObject*
+StreamingBase64Decoder_needs_more_data(StreamingBase64Decoder *self, PyObject *args UNUSED) {
+    return Py_NewRef(self->needs_more_data ? Py_True : Py_False);
+}
+
 static PyTypeObject StreamingBase64Decoder_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     .tp_name = "kitty.fast_data_types.StreamingBase64Decoder",
@@ -198,6 +222,7 @@ static PyTypeObject StreamingBase64Decoder_Type = {
         {"decode", (PyCFunction)StreamingBase64Decoder_decode, METH_O, ""},
         {"decode_into", (PyCFunction)(void(*)(void))StreamingBase64Decoder_decode_into, METH_FASTCALL, ""},
         {"reset", (PyCFunction)StreamingBase64Decoder_reset, METH_NOARGS, ""},
+        {"needs_more_data", (PyCFunction)StreamingBase64Decoder_needs_more_data, METH_NOARGS, ""},
         {NULL, NULL, 0, NULL},
     },
     .tp_new = PyType_GenericNew,
@@ -225,7 +250,9 @@ StreamingBase64Encoder_encode(StreamingBase64Decoder *self, PyObject *a) {
     size_t sz = required_buffer_size_for_base64_encode(data.len);
     RAII_PyObject(ans, PyBytes_FromStringAndSize(NULL, sz));
     if (!ans) return NULL;
+    Py_BEGIN_ALLOW_THREADS
     base64_stream_encode(&self->state, data.buf, data.len, PyBytes_AS_STRING(ans), &sz);
+    Py_END_ALLOW_THREADS
     if (_PyBytes_Resize(&ans, sz) != 0) return NULL;
     return Py_NewRef(ans);
 }
@@ -241,7 +268,9 @@ StreamingBase64Encoder_encode_into(StreamingBase64Decoder *self, PyObject *const
     if (!src.buf || !src.len) return PyLong_FromLong(0);
     size_t sz = required_buffer_size_for_base64_encode(src.len);
     if ((Py_ssize_t)sz > data.len) { PyErr_SetString(PyExc_BufferError, "output buffer too small"); return NULL; }
+    Py_BEGIN_ALLOW_THREADS
     base64_stream_encode(&self->state, src.buf, src.len, data.buf, &sz);
+    Py_END_ALLOW_THREADS
     return PyLong_FromSize_t(sz);
 }
 

kitty/fast_data_types.pyi

@@ -1719,7 +1719,8 @@ class StreamingBase64Decoder:
     def decode(self, data: ReadableBuffer) -> bytes: ...
     # decode the specified data, return number of bytes written dest should be as large as src (technically 3/4 src + 2)
     def decode_into(self, dest: WriteableBuffer, src: ReadableBuffer) -> int: ...
-
+    # whether the data stream decoded so far is complete or not
+    def needs_more_data(self) -> bool: ...
 
 
 class StreamingBase64Encodeer:

kitty_tests/clipboard.py

@@ -1,8 +1,10 @@
 #!/usr/bin/env python
 # License: GPLv3 Copyright: 2022, Kovid Goyal <kovid at kovidgoyal.net>
 
+from base64 import standard_b64decode, standard_b64encode
 
 from kitty.clipboard import WriteRequest
+from kitty.fast_data_types import StreamingBase64Decoder
 
 from . import BaseTest
 
@@ -30,3 +32,23 @@ class TestClipboard(BaseTest):
         for x in 'bGlnaHQgd29y':
             wr.add_base64_data(x)
         self.ae(wr.data_for(), b'light wor')
+
+    def test_base64_streaming_decoder(self):
+        d = StreamingBase64Decoder()
+        c = standard_b64encode(b'abcdef')
+        self.ae(b'abcdef', d.decode(c))
+        self.assertFalse(d.needs_more_data())
+        a = d.decode(c[:4])
+        self.assertFalse(d.needs_more_data())
+        self.ae(b'abcdef', a + d.decode(c[4:]))
+        self.assertFalse(d.needs_more_data())
+        a = d.decode(c[:1])
+        self.assertTrue(d.needs_more_data())
+        self.ae(b'abcdef', a + d.decode(c[1:4]) + d.decode(c[4:]))
+        self.assertFalse(d.needs_more_data())
+        c = standard_b64encode(b'abcd')
+        self.ae(b'abcd', d.decode(c[:2]) + d.decode(c[2:]))
+        c1 = standard_b64encode(b'1' * 4096)
+        c2 = standard_b64encode(b'2' * 4096)
+        self.ae(standard_b64decode(c1) + standard_b64decode(c2), d.decode(c1) + d.decode(c2))
+        self.assertFalse(d.needs_more_data())