From 15205be67b82f6b906aefc45bd5d4d9ae2b2775e Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Sun, 23 Jul 2023 20:28:17 +0530 Subject: [PATCH] More work on porting receive kitten --- kittens/transfer/receive.go | 325 ++++++++++++++++++++++++++++++++++-- 1 file changed, 312 insertions(+), 13 deletions(-) diff --git a/kittens/transfer/receive.go b/kittens/transfer/receive.go index a756099ac..88191d8f1 100644 --- a/kittens/transfer/receive.go +++ b/kittens/transfer/receive.go @@ -9,6 +9,7 @@ import ( "io" "io/fs" "os" + "path/filepath" "strconv" "strings" "time" @@ -17,6 +18,7 @@ import ( "kitty/tools/cli/markup" "kitty/tools/tui/loop" "kitty/tools/utils" + "kitty/tools/utils/humanize" "kitty/tools/wcswidth" ) @@ -32,7 +34,7 @@ const ( ) type output_file interface { - write([]byte) error + write([]byte) (int, error) close() error tell() (int64, error) } @@ -41,21 +43,20 @@ type filesystem_file struct { f *os.File } -func (ff *filesystem_file) tell() int64 { - pos, _ := ff.f.Seek(0, os.SEEK_CUR) - return pos +func (ff *filesystem_file) tell() (int64, error) { + return ff.f.Seek(0, os.SEEK_CUR) } func (ff *filesystem_file) close() error { return ff.f.Close() } -func (ff *filesystem_file) write(data []byte) error { +func (ff *filesystem_file) write(data []byte) (int, error) { n, err := ff.f.Write(data) if err == nil && n < len(data) { err = io.ErrShortWrite } - return err + return n, err } type remote_file struct { @@ -89,6 +90,63 @@ func (self *remote_file) close() error { return nil } +func (self *remote_file) write_data(data []byte, is_last bool) (amt_written int64, err error) { + self.received_bytes += int64(len(data)) + self.decompressor(data, is_last, func(data []byte) (err error) { + switch self.ftype { + case FileType_symlink: + self.remote_symlink_value += string(data) + return + case FileType_regular: + if self.actual_file == nil { + parent := filepath.Dir(self.expanded_local_path) + if parent != "" { + os.MkdirAll(parent, 0o755) + } + if self.expect_diff { + panic(`TODO: create PatchFile for rsync`) + } else { + if ff, err := os.Create(self.expanded_local_path); err != nil { + return err + } else { + f := filesystem_file{f: ff} + self.actual_file = &f + } + } + base, err := self.actual_file.tell() + if err != nil { + return err + } + for len(data) > 0 { + n, werr := self.actual_file.write(data) + data = data[n:] + if werr != nil && werr != io.ErrShortWrite { + return werr + } + } + pos, err := self.actual_file.tell() + if err != nil { + return err + } + amt_written = pos - base + if is_last { + self.actual_file.close() + self.actual_file = nil + } + } + } + return + }) + return +} + +func (self *remote_file) apply_metadata() { + if self.ftype != FileType_symlink { + os.Chmod(self.expanded_local_path, self.permissions) + os.Chtimes(self.expanded_local_path, time.Unix(0, int64(self.mtime)), time.Unix(0, int64(self.mtime))) + } +} + func new_remote_file(opts *Options, ftc *FileTransmissionCommand) (*remote_file, error) { spec_id, err := strconv.Atoi(ftc.File_id) if err != nil { @@ -116,13 +174,47 @@ type receive_progress_tracker struct { total_bytes_to_transfer int64 total_transferred int64 transfered_stats_amt int64 - transfered_stats_interval float64 + transfered_stats_interval time.Duration started_at time.Time transfers []Transfer active_file *remote_file done_files []*remote_file } +func (self *receive_progress_tracker) change_active_file(nf *remote_file) { + now := time.Now() + self.active_file = nf + nf.transmit_started_at = now +} + +func (self *receive_progress_tracker) start_transfer() { + self.started_at = time.Now() + self.transfers = append(self.transfers, Transfer{at: time.Now()}) +} + +func (self *receive_progress_tracker) file_written(af *remote_file, amt int64, is_done bool) { + if self.active_file != af { + self.change_active_file(af) + } + af.written_bytes += amt + self.total_transferred += amt + now := time.Now() + self.transfers = append(self.transfers, Transfer{amt: amt, at: now}) + for len(self.transfers) > 2 && self.transfers[0].is_too_old(now) { + utils.ShiftLeft(self.transfers, 1) + } + self.transfered_stats_interval = now.Sub(self.transfers[0].at) + self.transfered_stats_amt = 0 + for _, t := range self.transfers { + self.transfered_stats_amt += t.amt + } + if is_done { + af.done_at = now + self.done_files = append(self.done_files, af) + } + +} + type manager struct { request_id string cli_opts *Options @@ -153,11 +245,6 @@ 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) @@ -191,6 +278,54 @@ func (self *handler) abort_transfer(delay time.Duration) { self.lp.AddTimer(delay, false, self.do_error_quit) } +func (self *manager) finalize_transfer() (err error) { + self.transfer_done = true + rid_map := make(map[string]*remote_file) + for _, f := range self.files { + rid_map[f.remote_id] = f + } + for _, f := range self.files { + switch f.ftype { + case FileType_directory: + if err = os.MkdirAll(f.expanded_local_path, 0o755); err != nil { + return fmt.Errorf("Failed to create directory with error: %w", err) + } + case FileType_link: + tgt, found := rid_map[f.remote_target] + if !found { + return fmt.Errorf(`Hard link with remote id: {%s} not found`, f.remote_target) + } + if err = os.MkdirAll(filepath.Dir(f.expanded_local_path), 0o755); err == nil { + os.Remove(f.expanded_local_path) + err = os.Link(tgt.expanded_local_path, f.expanded_local_path) + } + if err != nil { + return fmt.Errorf(`Failed to create link with error: %w`, err) + } + case FileType_symlink: + lt := f.remote_symlink_value + if f.remote_target != "" { + tgt, found := rid_map[f.remote_target] + if !found { + return fmt.Errorf(`Symbolic link with remote id: {%s} not found`, f.remote_target) + } + lt = tgt.expanded_local_path + if !strings.HasPrefix(f.remote_symlink_value, "/") { + if lt, err = filepath.Rel(filepath.Dir(f.expanded_local_path), lt); err != nil { + return fmt.Errorf(`Could not make symlink relative with error: %w`, err) + } + } + } + os.Remove(f.expanded_local_path) + if err = os.Symlink(lt, f.expanded_local_path); err != nil { + return fmt.Errorf(`Failed to create symlink with error: %w`, err) + } + f.apply_metadata() + } + } + return +} + func (self *manager) on_file_transfer_response(ftc *FileTransmissionCommand) (err error) { switch self.state { case state_waiting_for_permission: @@ -264,6 +399,168 @@ func (self *manager) on_file_transfer_response(ftc *FileTransmissionCommand) (er return } +type tree_node struct { + entry *remote_file + parent *tree_node + added_files map[*remote_file]*tree_node +} + +func (self *tree_node) add_child(f *remote_file) *tree_node { + if _, found := self.added_files[f]; found { + return self + } + c := tree_node{entry: f, parent: self, added_files: make(map[*remote_file]*tree_node)} + f.expanded_local_path = filepath.Join(self.entry.expanded_local_path, filepath.Base(f.remote_path)) + self.added_files[f] = &c + return &c +} + +func walk_tree(root *tree_node, cb func(*tree_node) error) error { + for _, c := range root.added_files { + if err := cb(c); err != nil { + return err + } + if err := walk_tree(c, cb); err != nil { + return err + } + } + return nil +} + +func ensure_parent(f *remote_file, root_node *tree_node, node_map map[string]*tree_node, fid_map map[string]*remote_file) *tree_node { + if f.parent == "" { + return root_node + } + if parent, found := node_map[f.parent]; found { + return parent + } + fp := fid_map[f.parent] + gp := ensure_parent(fp, root_node, node_map, fid_map) + return gp.add_child(fp) +} + +func make_tree(all_files []*remote_file, local_base string) (root_node *tree_node) { + fid_map := make(map[string]*remote_file, len(all_files)) + for _, f := range all_files { + fid_map[f.remote_id] = f + } + node_map := make(map[string]*tree_node) + root_node = &tree_node{added_files: make(map[*remote_file]*tree_node)} + + for _, f := range all_files { + p := ensure_parent(f, root_node, node_map, fid_map) + p.add_child(f) + } + return +} + +func isdir(path string) bool { + if s, err := os.Stat(path); err == nil { + return s.IsDir() + } + return false +} + +func files_for_receive(opts *Options, dest string, files []*remote_file, remote_home string, specs []string) (ans []*remote_file, err error) { + spec_map := make(map[int][]*remote_file) + for _, f := range files { + spec_map[f.spec_id] = append(spec_map[f.spec_id], f) + } + spec_paths := make([]string, len(specs)) + for i := range specs { + spec_paths[i] = spec_map[i][0].remote_path + } + if opts.Mode == "mirror" { + common_path := utils.Commonpath(spec_paths...) + home := strings.TrimRight(remote_home, "/") + if strings.HasPrefix(common_path, home+"/") { + for i, x := range spec_paths { + b, err := filepath.Rel(home, x) + if err != nil { + return nil, err + } + spec_paths[i] = filepath.Join("~", b) + } + } + for spec_id, files_for_spec := range spec_map { + spec := spec_paths[spec_id] + tree := make_tree(files_for_spec, filepath.Dir(expand_home(spec))) + walk_tree(tree, func(x *tree_node) error { + ans = append(ans, x.entry) + return nil + }) + } + } else { + number_of_source_files := 0 + for _, x := range spec_map { + number_of_source_files += len(x) + } + dest_is_dir := strings.HasSuffix(dest, "/") || number_of_source_files > 1 || isdir(dest) + for _, files_for_spec := range spec_map { + if dest_is_dir { + dest_path := filepath.Join(dest, filepath.Base(files_for_spec[0].remote_path)) + tree := make_tree(files_for_spec, filepath.Dir(expand_home(dest_path))) + walk_tree(tree, func(x *tree_node) error { + ans = append(ans, x.entry) + return nil + }) + } else { + f := files_for_spec[0] + f.expanded_local_path = dest + ans = append(ans, f) + } + } + } + return +} + +func (self *manager) collect_files() (err error) { + if self.files, err = files_for_receive(self.cli_opts, self.dest, self.files, self.remote_home, self.spec); err != nil { + return err + } + self.progress_tracker.total_size_of_all_files = 0 + for _, f := range self.files { + if f.ftype != FileType_directory && f.ftype != FileType_link { + self.files_to_be_transferred[f.file_id] = f + self.progress_tracker.total_size_of_all_files += utils.Max(0, f.expected_size) + } + } + self.progress_tracker.total_bytes_to_transfer = self.progress_tracker.total_size_of_all_files + return nil +} + +func (self *handler) print_continue_msg() { + self.lp.Println(`Press `, self.ctx.Green(`y`), ` to continue or `, self.ctx.BrightRed(`n`), ` to abort`) +} + +func lexists(path string) bool { + _, err := os.Lstat(path) + return err == nil +} + +func (self *handler) print_check_paths() { + if self.check_paths_printed { + return + } + self.check_paths_printed = true + self.lp.Println(`The following file transfers will be performed. A red destination means an existing file will be overwritten.`) + for _, df := range self.manager.files { + self.lp.QueueWriteString(self.ctx.Prettify(fmt.Sprintf(":%s:`%s` ", df.ftype.Color(), df.ftype.ShortText()))) + self.lp.QueueWriteString(" ") + lpath := df.expanded_local_path + if lexists(lpath) { + lpath = self.ctx.Prettify(fmt.Sprintf(":red:`%s` ", lpath)) + } + self.lp.Println(df.display_name, " → ", lpath) + } + self.lp.Println(fmt.Sprintf(`Transferring %d file(s) of total size: %s`, len(self.manager.files), humanize.Size(self.manager.progress_tracker.total_size_of_all_files))) + self.print_continue_msg() +} + +func (self *handler) confirm_paths() { + self.print_check_paths() +} + func (self *handler) on_file_transfer_response(ftc *FileTransmissionCommand) (err error) { if ftc.Id != self.manager.request_id { return @@ -304,7 +601,9 @@ func (self *handler) on_file_transfer_response(ftc *FileTransmissionCommand) (er self.lp.Quit(1) return } - self.manager.collect_files(self.cli_opts) + if err = self.manager.collect_files(); err != nil { + return + } if self.cli_opts.ConfirmPaths { self.confirm_paths() } else {