More work on porting receive kitten

kittens/transfer/receive.go

@@ -3,15 +3,18 @@
 package transfer
 
 import (
+	"bytes"
 	"compress/zlib"
 	"fmt"
 	"io"
 	"io/fs"
 	"os"
 	"strconv"
+	"strings"
 	"time"
 
 	"kitty"
+	"kitty/tools/cli/markup"
 	"kitty/tools/tui/loop"
 	"kitty/tools/utils"
 	"kitty/tools/wcswidth"
@@ -121,23 +124,27 @@ type receive_progress_tracker struct {
 }
 
 type manager struct {
-	request_id       string
-	spec             []string
-	dest             string
-	bypass           string
-	use_rsync        bool
-	failed_specs     map[int]string
-	spec_counts      map[int]int
-	remote_home      string
-	prefix, suffix   string
-	transfer_done    bool
-	files            []*remote_file
-	state            state
-	progress_tracker receive_progress_tracker
+	request_id              string
+	cli_opts                *Options
+	spec                    []string
+	dest                    string
+	bypass                  string
+	use_rsync               bool
+	failed_specs            map[int]string
+	spec_counts             map[int]int
+	remote_home             string
+	prefix, suffix          string
+	transfer_done           bool
+	files                   []*remote_file
+	files_to_be_transferred map[string]*remote_file
+	state                   state
+	progress_tracker        receive_progress_tracker
 }
 
 type handler struct {
 	lp                    *loop.Loop
+	cli_opts              *Options
+	ctx                   *markup.Context
 	manager               manager
 	quit_after_write_code int
 	check_paths_printed   bool
@@ -146,6 +153,200 @@ type handler struct {
 	max_name_length       int
 }
 
+func (self *receive_progress_tracker) start_transfer() {
+	self.started_at = time.Now()
+	self.transfers = append(self.transfers, Transfer{})
+}
+
+func (self *manager) start_transfer(send func(string) loop.IdType) {
+	s := func(c FileTransmissionCommand) {
+		send(self.prefix)
+		send(c.Serialize(false))
+		send(self.suffix)
+	}
+	s(FileTransmissionCommand{Action: Action_receive, Bypass: self.bypass, Size: int64(len(self.spec))})
+	for i, x := range self.spec {
+		s(FileTransmissionCommand{Action: Action_file, File_id: strconv.Itoa(i), Name: x})
+	}
+	self.progress_tracker.start_transfer()
+}
+
+func (self *handler) print_err(err error) {
+	self.lp.Println(self.ctx.BrightRed(err.Error()))
+}
+
+func (self *handler) do_error_quit(loop.IdType) error {
+	self.lp.Quit(1)
+	return nil
+}
+
+func (self *handler) abort_transfer(delay time.Duration) {
+	if delay <= 0 {
+		delay = time.Second * 5
+	}
+	self.lp.QueueWriteString(self.manager.prefix)
+	self.lp.QueueWriteString(FileTransmissionCommand{Action: Action_cancel}.Serialize(false))
+	self.lp.QueueWriteString(self.manager.suffix)
+	self.manager.state = state_canceled
+	self.lp.AddTimer(delay, false, self.do_error_quit)
+}
+
+func (self *manager) on_file_transfer_response(ftc *FileTransmissionCommand) (err error) {
+	switch self.state {
+	case state_waiting_for_permission:
+		if ftc.Action == Action_status {
+			if ftc.Status == `OK` {
+				self.state = state_waiting_for_file_metadata
+			} else {
+				return fmt.Errorf(`Permission for transfer denied`)
+			}
+		} else {
+			return fmt.Errorf(`Unexpected response from terminal: %s`, ftc.String())
+		}
+	case state_waiting_for_file_metadata:
+		switch ftc.Action {
+		case Action_status:
+			if ftc.File_id != "" {
+				fid, err := strconv.Atoi(ftc.File_id)
+				if err != nil {
+					return fmt.Errorf(`Unexpected response from terminal (non-integer file_id): %s`, ftc.String())
+				}
+				if fid < 0 || fid >= len(self.spec) {
+					return fmt.Errorf(`Unexpected response from terminal (out-of-range file_id): %s`, ftc.String())
+				}
+				self.failed_specs[fid] = ftc.Status
+			} else {
+				if ftc.Status == `OK` {
+					self.state = state_transferring
+					self.remote_home = ftc.Name
+					return
+				}
+				return fmt.Errorf("%s", ftc.Status)
+			}
+		case Action_file:
+			fid, err := strconv.Atoi(ftc.File_id)
+			if err != nil {
+				return fmt.Errorf(`Unexpected response from terminal (non-integer file_id): %s`, ftc.String())
+			}
+			if fid < 0 || fid >= len(self.spec) {
+				return fmt.Errorf(`Unexpected response from terminal (out-of-range file_id): %s`, ftc.String())
+			}
+			self.spec_counts[fid] += 1
+			if rf, err := new_remote_file(self.cli_opts, ftc); err == nil {
+				self.files = append(self.files, rf)
+			} else {
+				return err
+			}
+		default:
+			return fmt.Errorf(`Unexpected response from terminal (invalid action): %s`, ftc.String())
+		}
+	case state_transferring:
+		if ftc.Action == Action_data || ftc.Action == Action_end_data {
+			f, found := self.files_to_be_transferred[ftc.File_id]
+			if !found {
+				return fmt.Errorf(`Got data for unknown file id: %s`, ftc.File_id)
+			}
+			is_last := ftc.Action == Action_end_data
+			if amt_written, err := f.write_data(ftc.Data, is_last); err != nil {
+				return err
+			} else {
+				self.progress_tracker.file_written(f, amt_written, is_last)
+			}
+			if is_last {
+				delete(self.files_to_be_transferred, ftc.File_id)
+				if len(self.files_to_be_transferred) == 0 {
+					return self.finalize_transfer()
+				}
+			}
+		}
+
+	}
+	return
+}
+
+func (self *handler) on_file_transfer_response(ftc *FileTransmissionCommand) (err error) {
+	if ftc.Id != self.manager.request_id {
+		return
+	}
+	if ftc.Action == Action_status && ftc.Status == "CANCELED" {
+		self.lp.Quit(1)
+		return
+	}
+	if self.quit_after_write_code > -1 || self.manager.state == state_canceled {
+		return
+	}
+	transfer_started := self.manager.state == state_transferring
+	err = self.manager.on_file_transfer_response(ftc)
+	if err != nil {
+		self.print_err(err)
+		self.lp.Println(`Waiting to ensure terminal cancels transfer, will quit in a few seconds`)
+		self.abort_transfer(-1)
+		return
+	}
+	if !transfer_started && self.manager.state == state_transferring {
+		if len(self.manager.failed_specs) > 0 {
+			self.print_err(fmt.Errorf(`Failed to process some sources:`))
+			for spec_id, msg := range self.manager.failed_specs {
+				spec := self.manager.spec[spec_id]
+				self.lp.Println(fmt.Sprintf(`  {%s}: {%s}`, spec, msg))
+			}
+			self.lp.Quit(1)
+			return
+		}
+		zero_specs := make([]string, 0, len(self.manager.spec_counts))
+		for k, v := range self.manager.spec_counts {
+			if v == 0 {
+				zero_specs = append(zero_specs, self.manager.spec[k])
+			}
+		}
+		if len(zero_specs) > 0 {
+			self.print_err(fmt.Errorf(`No matches found for: %s`, strings.Join(zero_specs, ", ")))
+			self.lp.Quit(1)
+			return
+		}
+		self.manager.collect_files(self.cli_opts)
+		if self.cli_opts.ConfirmPaths {
+			self.confirm_paths()
+		} else {
+			self.start_transfer()
+		}
+	}
+	if self.manager.transfer_done {
+		self.lp.QueueWriteString(self.manager.prefix)
+		self.lp.QueueWriteString(FileTransmissionCommand{Action: Action_finish}.Serialize(false))
+		self.lp.QueueWriteString(self.manager.suffix)
+		self.quit_after_write_code = 0
+		self.refresh_progress()
+	} else if self.transmit_started {
+		self.refresh_progress()
+	}
+	return
+}
+
+func (self *handler) on_sigint() (handled bool, err error) {
+	handled = true
+	if self.quit_after_write_code > -1 {
+		return
+	}
+	if self.manager.state == state_canceled {
+		self.lp.Println(`Waiting for canceled acknowledgement from terminal, will abort in a few seconds if no response received`)
+		return
+	}
+	self.print_err(fmt.Errorf(`Interrupt requested, cancelling transfer, transferred files are in undefined state`))
+	self.abort_transfer(-1)
+	return
+}
+
+func (self *handler) on_sigterm() (handled bool, err error) {
+	handled = true
+	if self.quit_after_write_code > -1 {
+		return
+	}
+	self.print_err(fmt.Errorf(`Terminate requested, cancelling transfer, transferred files are in undefined state`))
+	self.abort_transfer(2 * time.Second)
+	return
+}
+
 func receive_loop(opts *Options, spec []string, dest string) (err error, rc int) {
 	lp, err := loop.New(loop.NoAlternateScreen, loop.NoRestoreColors)
 	if err != nil {
@@ -153,11 +354,11 @@ func receive_loop(opts *Options, spec []string, dest string) (err error, rc int)
 	}
 
 	handler := handler{
-		lp: lp, quit_after_write_code: -1,
+		lp: lp, quit_after_write_code: -1, cli_opts: opts,
 		manager: manager{
 			request_id: random_id(), spec: spec, dest: dest, bypass: opts.PermissionsBypass, use_rsync: opts.TransmitDeltas,
 			failed_specs: make(map[int]string, len(spec)), spec_counts: make(map[int]int, len(spec)),
-			suffix: "\x1b\\",
+			suffix: "\x1b\\", cli_opts: opts, files_to_be_transferred: make(map[string]*remote_file),
 		},
 	}
 	for i := range spec {
@@ -170,6 +371,37 @@ func receive_loop(opts *Options, spec []string, dest string) (err error, rc int)
 		}
 	}
 
+	lp.OnInitialize = func() (string, error) {
+		lp.SetCursorVisible(false)
+		lp.Println("Scanning files…")
+		handler.manager.start_transfer(lp.QueueWriteString)
+		return "", nil
+	}
+
+	lp.OnFinalize = func() string {
+		lp.SetCursorVisible(true)
+		return ""
+	}
+
+	lp.OnSIGINT = handler.on_sigint
+	lp.OnSIGTERM = handler.on_sigterm
+
+	ftc_code := strconv.Itoa(kitty.FileTransferCode)
+	lp.OnEscapeCode = func(et loop.EscapeCodeType, payload []byte) error {
+		if et == loop.OSC {
+			if idx := bytes.IndexByte(payload, ';'); idx > 0 {
+				if utils.UnsafeBytesToString(payload[:idx]) == ftc_code {
+					ftc, err := NewFileTransmissionCommand(utils.UnsafeBytesToString(payload[idx+1:]))
+					if err != nil {
+						return fmt.Errorf("Received invalid FileTransmissionCommand from terminal with error: %w", err)
+					}
+					return handler.on_file_transfer_response(ftc)
+				}
+			}
+		}
+		return nil
+	}
+
 	err = lp.Run()
 	if err != nil {
 		return err, 1