diff --git a/kittens/transfer/receive.go b/kittens/transfer/receive.go index 758e285e4..a756099ac 100644 --- a/kittens/transfer/receive.go +++ b/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 diff --git a/tools/tui/loop/api.go b/tools/tui/loop/api.go index f66397c4a..33863e489 100644 --- a/tools/tui/loop/api.go +++ b/tools/tui/loop/api.go @@ -98,6 +98,12 @@ type Loop struct { // Called when main loop is woken up OnWakeup func() error + + // Called on SIGINT return true if you wish to handle it yourself + OnSIGINT func() (bool, error) + + // Called on SIGTERM return true if you wish to handle it yourself + OnSIGTERM func() (bool, error) } func New(options ...func(self *Loop)) (*Loop, error) { diff --git a/tools/tui/loop/run.go b/tools/tui/loop/run.go index 4b3e4fd5d..1c65bd46f 100644 --- a/tools/tui/loop/run.go +++ b/tools/tui/loop/run.go @@ -198,12 +198,22 @@ func (self *Loop) handle_end_of_bracketed_paste() { func (self *Loop) on_signal(s unix.Signal) error { switch s { case unix.SIGINT: + if self.OnSIGINT != nil { + if handled, err := self.OnSIGINT(); handled { + return err + } + } return self.on_SIGINT() case unix.SIGPIPE: return self.on_SIGPIPE() case unix.SIGWINCH: return self.on_SIGWINCH() case unix.SIGTERM: + if self.OnSIGTERM != nil { + if handled, err := self.OnSIGTERM(); handled { + return err + } + } return self.on_SIGTERM() case unix.SIGTSTP: return self.on_SIGTSTP()