Implement writing of MIME data on drop

kittens/dnd/main.go

@@ -4,6 +4,7 @@ package dnd
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"maps"
@@ -18,6 +19,7 @@ import (
 	"github.com/kovidgoyal/kitty/tools/tty"
 	"github.com/kovidgoyal/kitty/tools/tui/loop"
 	"github.com/kovidgoyal/kitty/tools/utils"
+	"github.com/kovidgoyal/kitty/tools/utils/streaming_base64"
 	"github.com/kovidgoyal/kitty/tools/wcswidth"
 )
 
@@ -52,6 +54,69 @@ type drop_dest struct {
 	dest             io.WriteCloser
 	mime_type        string
 	completed        bool
+	close_on_finish  bool
+	b64_decoder      streaming_base64.StreamingBase64Decoder
+}
+
+func open_file_for_writing(path string) (*os.File, error) {
+	f, err := os.Create(path)
+	if errors.Is(err, os.ErrNotExist) {
+		dir := filepath.Dir(path)
+		if err := os.MkdirAll(dir, 0755); err != nil {
+			return nil, err
+		}
+		return os.Create(path)
+	}
+	return f, err
+}
+
+func (d *drop_dest) write(chunk []byte) (err error) {
+	if d.dest == nil {
+		d.dest, err = open_file_for_writing(d.path)
+		d.close_on_finish = true
+		if err != nil {
+			return
+		}
+
+	}
+	_, err = d.dest.Write(chunk)
+	return
+}
+
+func (d *drop_dest) finish() error {
+	defer func() {
+		d.completed = true
+		if d.dest != nil && d.close_on_finish {
+			d.dest.Close()
+			d.dest = nil
+		}
+	}()
+	if chunk, err := d.b64_decoder.Finish(); err != nil {
+		return err
+	} else if len(chunk) > 0 {
+		return d.write(chunk)
+	}
+	return nil
+}
+
+func (d *drop_dest) add_data(x []byte, output_buf []byte, has_more bool) error {
+	d.completed = false
+	for chunk, err := range d.b64_decoder.Decode(x, output_buf) {
+		if err == nil {
+			err = d.write(chunk)
+		}
+		if err != nil {
+			return err
+		}
+	}
+	if !has_more {
+		if chunk, err := d.b64_decoder.Finish(); err != nil {
+			return err
+		} else if len(chunk) > 0 {
+			return d.write(chunk)
+		}
+	}
+	return nil
 }
 
 type button_region struct {
@@ -322,20 +387,25 @@ func run_loop(opts *Options, drop_dests map[string]*drop_dest, drag_sources map[
 		return nil
 	}
 
+	var drop_buf []byte
+
 	on_drop_data := func(cmd DC) error {
 		if drop_status.remote_phase_started {
 			return on_remote_drop_data(cmd)
 		}
 		idx := cmd.X - 1
-		if idx < 0 || idx > len(drop_status.accepted_mimes) {
+		if idx < 0 || idx > len(drop_status.offered_mimes) {
 			return fmt.Errorf("terminal sent drop data for a index outside the list of accepted MIMEs")
 		}
-		mime := drop_status.accepted_mimes[idx]
+		mime := drop_status.offered_mimes[idx]
 		dest := drop_dests[mime]
 		if cmd.Xp == 1 && mime == "text/uri-list" {
 			drop_status.is_remote_client = true
 		}
 		if !cmd.Has_more && len(cmd.Payload) == 0 {
+			if err := dest.finish(); err != nil {
+				return err
+			}
 			dest.completed = true
 			pending := false
 			for _, d := range drop_dests {
@@ -349,8 +419,10 @@ func run_loop(opts *Options, drop_dests map[string]*drop_dest, drag_sources map[
 			}
 			return nil
 		}
-		// TODO: Implement this
-		return nil
+		if sz := max(4096, len(cmd.Payload)+4); len(drop_buf) < sz {
+			drop_buf = make([]byte, sz)
+		}
+		return dest.add_data(cmd.Payload, drop_buf, cmd.Has_more)
 	}
 	// }}}
 

kitty/dnd.c

@@ -2094,7 +2094,7 @@ dnd_test_fake_drop_data(PyObject *self UNUSED, PyObject *args) {
     const char *mime;
     RAII_PY_BUFFER(data);
     int error_code = 0, no_eod = 0;
-    if (!PyArg_ParseTuple(args, "Ksy*|ii", &window_id, &mime, &data, &error_code, &no_eod)) return NULL;
+    if (!PyArg_ParseTuple(args, "Ksy*|ip", &window_id, &mime, &data, &error_code, &no_eod)) return NULL;
     Window *w = window_for_window_id((id_type)window_id);
     if (!w) { PyErr_SetString(PyExc_ValueError, "Window not found"); return NULL; }
     if (error_code > 0) {

kitty_tests/dnd_kitten.py

@@ -3,6 +3,7 @@
 
 import os
 import random
+import shutil
 import tempfile
 from base64 import standard_b64encode
 from functools import partial
@@ -37,7 +38,7 @@ def create_fs(base):
         if sz == 0:
             sz = random.randint(5713, 9879)
         with open(join(*path), 'wb') as f:
-            f.write(os.urandom(sz))
+            f.write(b'x' * sz)
     os.makedirs(join('d1', 'sd', 'ssd'))
     os.mkdir(join('d2'))
     os.symlink('/does-not-exist', join('s1'))
@@ -160,7 +161,8 @@ class TestDnDKitten(BaseTest):
         self.dnd_kitten_drop(True)
 
     def dnd_kitten_drop(self, remote_client):
-        self.finish_setup(remote_client=remote_client, cli_args=('--drop=image/png:images/image.png',))
+        img_drop_path = 'images/image.png'
+        self.finish_setup(remote_client=remote_client, cli_args=(f'--drop=image/png:{img_drop_path}',))
         copy, move = self.get_button_geometry()
         all_mimes = 'text/uri-list a/b c/d'
         for b, expected in ((copy, GLFW_DRAG_OPERATION_COPY), (move, GLFW_DRAG_OPERATION_MOVE)):
@@ -200,5 +202,24 @@ class TestDnDKitten(BaseTest):
         uri_list.insert(3, 'ignore://me')
         dnd_test_fake_drop_data(self.capture.window_id, 'text/uri-list', '\r\n'.join(uri_list).encode())
         self.assertEqual('image/png', self.probe_state('drop_getting_data_for_mime'))
+
+        def send_file_in_chunks(f, mime, sz=3072):
+            while True:
+                chunk = f.read(sz)
+                if not chunk:
+                    break
+                dnd_test_fake_drop_data(self.capture.window_id, mime, chunk, 0, True)
+            dnd_test_fake_drop_data(self.capture.window_id, mime, b'')
+
+        with open(os.path.join(self.src_data_dir, 'some-image.png'), 'rb') as f:
+            send_file_in_chunks(f, 'image/png', 1117)
+
         self.send_dnd_command_to_kitten('DROP_IS_REMOTE')
         self.wait_for_responses(str(remote_client))
+        jn = os.path.join
+        self.assert_files_have_same_content(jn(self.src_data_dir, 'some-image.png'), jn(self.kitten_wd, img_drop_path))
+        shutil.rmtree(os.path.dirname(jn(self.kitten_wd, img_drop_path)))
+
+    def assert_files_have_same_content(self, a, b):
+        with open(a, 'rb') as fa, open(b, 'rb') as fb:
+            self.assertEqual(fa.read(), fb.read(), f'{a} ({os.path.getsize(a)}) != {b} ({os.path.getsize(b)})')