From c15665c37558033e919dbb7cb0b9c118ff2fd68f Mon Sep 17 00:00:00 2001 From: Lonny Wong Date: Sun, 22 May 2022 17:15:16 +0800 Subject: [PATCH] implement trzsz client --- README.md | 14 ++ go.mod | 1 + go.sum | 2 + trzsz/buffer.go | 124 +++++++++++ trzsz/comm.go | 98 ++++++++- trzsz/escape.go | 130 ++++++++++++ trzsz/progress.go | 275 ++++++++++++++++++++++++- trzsz/transfer.go | 513 +++++++++++++++++++++++++++++++++++++++++++--- trzsz/trzsz.go | 22 +- trzsz/version.go | 2 +- 10 files changed, 1138 insertions(+), 43 deletions(-) create mode 100644 trzsz/buffer.go create mode 100644 trzsz/escape.go diff --git a/README.md b/README.md index 7c2209a..0cb787a 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,20 @@ Similar to lrzsz ( rz / sz ), command `trz` to upload files, command `tsz /path/ For more information, see the website of trzsz: [https://trzsz.github.io](https://trzsz.github.io/). +## Configuration + +`trzsz` looks for configuration at `~/.trzsz.conf`. e.g.: + +``` +DefaultUploadPath = +DefaultDownloadPath = /Users/username/Downloads +``` + +* If the `DefaultUploadPath` is not empty, the path will be opened by default while choosing upload files. + +* If the `DefaultDownloadPath` is not empty, downloading files will be saved to the path automatically instead of asking each time. + + ## Contact Feel free to email me . diff --git a/go.mod b/go.mod index ecdf4e7..64f0d7c 100644 --- a/go.mod +++ b/go.mod @@ -10,4 +10,5 @@ require ( github.com/ncruces/zenity v0.8.6 golang.org/x/sys v0.0.0-20220513210249-45d2b4557a2a golang.org/x/term v0.0.0-20220411215600-e5f449aeb171 + golang.org/x/text v0.3.7 ) diff --git a/go.sum b/go.sum index 2e683cc..4731151 100644 --- a/go.sum +++ b/go.sum @@ -61,6 +61,8 @@ golang.org/x/term v0.0.0-20220411215600-e5f449aeb171/go.mod h1:jbD1KX2456YbFQfuX golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk= +golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= diff --git a/trzsz/buffer.go b/trzsz/buffer.go new file mode 100644 index 0000000..ef627c0 --- /dev/null +++ b/trzsz/buffer.go @@ -0,0 +1,124 @@ +/* +MIT License + +Copyright (c) 2022 Lonny Wong + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package trzsz + +import ( + "bytes" + "time" +) + +type TrzszBuffer struct { + bufCh chan []byte + stopCh chan bool + nextBuf []byte + nextIdx int + readBuf *bytes.Buffer +} + +func NewTrzszBuffer() *TrzszBuffer { + return &TrzszBuffer{make(chan []byte, 10), make(chan bool, 1), nil, 0, new(bytes.Buffer)} +} + +func (b *TrzszBuffer) addBuffer(buf []byte) { + b.bufCh <- buf +} + +func (b *TrzszBuffer) stopBuffer() { + select { + case b.stopCh <- true: + default: + } +} + +func (b *TrzszBuffer) drainBuffer() { + for { + select { + case <-b.bufCh: + default: + return + } + } +} + +func (b *TrzszBuffer) nextBuffer(timeout <-chan time.Time) ([]byte, error) { + if b.nextBuf != nil && b.nextIdx < len(b.nextBuf) { + return b.nextBuf[b.nextIdx:], nil + } + select { + case b.nextBuf = <-b.bufCh: + b.nextIdx = 0 + return b.nextBuf, nil + case <-b.stopCh: + return nil, newTrzszError("Stopped") + case <-timeout: + return nil, newTrzszError("Receive data timeout") + } +} + +func (b *TrzszBuffer) readLine(timeout <-chan time.Time) ([]byte, error) { + b.readBuf.Reset() + for { + buf, err := b.nextBuffer(timeout) + if err != nil { + return nil, err + } + newLineIdx := bytes.IndexByte(buf, '\n') + if newLineIdx >= 0 { + b.nextIdx += newLineIdx + 1 // +1 to ignroe the '\n' + buf = buf[0:newLineIdx] + } else { + b.nextIdx += len(buf) + } + if bytes.IndexByte(buf, '\x03') >= 0 { // `ctrl + c` to interrupt + return nil, newTrzszError("Interrupted") + } + b.readBuf.Write(buf) + if newLineIdx >= 0 { + return b.readBuf.Bytes(), nil + } + } +} + +func (b *TrzszBuffer) readBinary(size int, timeout <-chan time.Time) ([]byte, error) { + b.readBuf.Reset() + if b.readBuf.Cap() < size { + b.readBuf.Grow(size) + } + for b.readBuf.Len() < size { + buf, err := b.nextBuffer(timeout) + if err != nil { + return nil, err + } + left := size - b.readBuf.Len() + if len(buf) > left { + b.nextIdx += left + buf = buf[0:left] + } else { + b.nextIdx += len(buf) + } + b.readBuf.Write(buf) + } + return b.readBuf.Bytes(), nil +} diff --git a/trzsz/comm.go b/trzsz/comm.go index ca34725..33f3d03 100644 --- a/trzsz/comm.go +++ b/trzsz/comm.go @@ -28,7 +28,12 @@ import ( "bytes" "compress/zlib" "encoding/base64" + "errors" + "fmt" "io/ioutil" + "os" + "path/filepath" + "runtime/debug" ) type PtyIO interface { @@ -38,7 +43,11 @@ type PtyIO interface { } type ProgressCallback interface { - // TODO + onNum(num int64) + onName(name string) + onSize(size int64) + onStep(step int64) + onDone(name string) } func encodeBytes(buf []byte) string { @@ -66,12 +75,95 @@ func decodeString(str string) ([]byte, error) { return ioutil.ReadAll(z) } +type TrzszError struct { + message string + errType string + trace bool +} + +func NewTrzszError(message string, errType string, trace bool) *TrzszError { + if errType == "fail" || errType == "FAIL" || errType == "EXIT" { + msg, err := decodeString(message) + if err != nil { + message = fmt.Sprintf("decode [%s] error: %s", message, err) + } else { + message = string(msg) + } + } else if len(errType) > 0 { + message = fmt.Sprintf("[TrzszError] %s: %s", errType, message) + } + err := &TrzszError{message, errType, trace} + if err.isTraceBack() { + err.message = fmt.Sprintf("%s\n%s", err.message, string(debug.Stack())) + } + return err +} + +func newTrzszError(message string) *TrzszError { + return NewTrzszError(message, "", false) +} + +func (e *TrzszError) Error() string { + return e.message +} + +func (e *TrzszError) isTraceBack() bool { + if e.errType == "fail" { + return false + } + return e.trace +} + +func (e *TrzszError) isRemoteExit() bool { + return e.errType == "EXIT" +} + +func (e *TrzszError) isRemoteFail() bool { + return e.errType == "fail" || e.errType == "FAIL" +} + func checkPathWritable(path string) error { - // TODO + fileInfo, err := os.Stat(path) + if errors.Is(err, os.ErrNotExist) { + return newTrzszError(fmt.Sprintf("No such directory: %s", path)) + } + if !fileInfo.IsDir() { + return newTrzszError(fmt.Sprintf("Not a directory: %s", path)) + } + if fileInfo.Mode().Perm()&(1<<7) == 0 { + return newTrzszError(fmt.Sprintf("No permission to write: %s", path)) + } return nil } func checkFilesReadable(files []string) error { - // TODO + for _, file := range files { + fileInfo, err := os.Stat(file) + if errors.Is(err, os.ErrNotExist) { + return newTrzszError(fmt.Sprintf("No such file: %s", file)) + } + if fileInfo.IsDir() { + return newTrzszError(fmt.Sprintf("Is a directory: %s", file)) + } + if !fileInfo.Mode().IsRegular() { + return newTrzszError(fmt.Sprintf("Not a regular file: %s", file)) + } + if fileInfo.Mode().Perm()&(1<<8) == 0 { + return newTrzszError(fmt.Sprintf("No permission to read: %s", file)) + } + } return nil } + +func getNewName(path, name string) (string, error) { + if _, err := os.Stat(filepath.Join(path, name)); errors.Is(err, os.ErrNotExist) { + return name, nil + } + for i := 0; i < 1000; i++ { + newName := fmt.Sprintf("%s.%d", name, i) + if _, err := os.Stat(filepath.Join(path, newName)); errors.Is(err, os.ErrNotExist) { + return newName, nil + } + } + return "", newTrzszError("Fail to assign new file name") +} diff --git a/trzsz/escape.go b/trzsz/escape.go new file mode 100644 index 0000000..0bc478d --- /dev/null +++ b/trzsz/escape.go @@ -0,0 +1,130 @@ +/* +MIT License + +Copyright (c) 2022 Lonny Wong + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package trzsz + +import ( + "fmt" + + "golang.org/x/text/encoding/charmap" +) + +func escapeCharsToCodes(escapeChars []interface{}) ([][]byte, error) { + escapeCodes := make([][]byte, len(escapeChars)) + encoder := charmap.ISO8859_1.NewEncoder() + for i, v := range escapeChars { + a, ok := v.([]interface{}) + if !ok { + return nil, newTrzszError(fmt.Sprintf("escape chars invalid: %#v", v)) + } + if len(a) != 2 { + return nil, newTrzszError(fmt.Sprintf("escape chars invalid: %#v", v)) + } + b, ok := a[0].(string) + if !ok { + return nil, newTrzszError(fmt.Sprintf("escape chars invalid: %#v", v)) + } + bb, err := encoder.Bytes([]byte(b)) + if err != nil { + return nil, err + } + if len(bb) != 1 { + return nil, newTrzszError(fmt.Sprintf("escape chars invalid: %#v", v)) + } + c, ok := a[1].(string) + if !ok { + return nil, newTrzszError(fmt.Sprintf("escape chars invalid: %#v", v)) + } + cc, err := encoder.Bytes([]byte(c)) + if err != nil { + return nil, err + } + if len(cc) != 2 { + return nil, newTrzszError(fmt.Sprintf("escape chars invalid: %#v", v)) + } + escapeCodes[i] = make([]byte, 3) + escapeCodes[i][0] = bb[0] + escapeCodes[i][1] = cc[0] + escapeCodes[i][2] = cc[1] + } + return escapeCodes, nil +} + +func escapeData(data []byte, escapeCodes [][]byte) []byte { + if len(escapeCodes) == 0 { + return data + } + + buf := make([]byte, len(data)*2) + idx := 0 + for _, d := range data { + escapeIdx := -1 + for j, e := range escapeCodes { + if d == e[0] { + escapeIdx = j + break + } + } + if escapeIdx < 0 { + buf[idx] = d + idx++ + } else { + buf[idx] = escapeCodes[escapeIdx][1] + idx++ + buf[idx] = escapeCodes[escapeIdx][2] + idx++ + } + } + return buf[:idx] +} + +func unescapeData(data []byte, escapeCodes [][]byte) []byte { + if len(escapeCodes) == 0 { + return data + } + + size := len(data) + buf := make([]byte, size) + idx := 0 + for i := 0; i < size; i++ { + escapeIdx := -1 + if i < size-1 { + for j, e := range escapeCodes { + if data[i] == e[1] && data[i+1] == e[2] { + escapeIdx = j + break + } + } + } + if escapeIdx < 0 { + buf[idx] = data[i] + idx++ + } else { + buf[idx] = escapeCodes[escapeIdx][0] + idx++ + i++ + } + } + return buf[:idx] +} diff --git a/trzsz/progress.go b/trzsz/progress.go index 4d69ba1..47a50c6 100644 --- a/trzsz/progress.go +++ b/trzsz/progress.go @@ -25,18 +25,283 @@ SOFTWARE. package trzsz import ( + "fmt" + "math" "os" + "strings" + "time" + "unicode/utf8" ) +func getDisplayLength(str string) int { + length := 0 + for _, r := range []rune(str) { + if utf8.RuneLen(r) == 1 { + length += 1 + } else { + length += 2 + } + } + return length +} + +func getEllipsisString(str string, max int) (string, int) { + var b strings.Builder + b.Grow(max) + max -= 3 + length := 0 + for _, r := range []rune(str) { + if utf8.RuneLen(r) > 1 { + if length+2 > max { + b.WriteString("...") + return b.String(), length + 3 + } else { + length += 2 + } + } else { + if length+1 > max { + b.WriteString("...") + return b.String(), length + 3 + } else { + length += 1 + } + } + b.WriteRune(r) + } + b.WriteString("...") + return b.String(), length + 3 +} + +func convertSizeToString(size float64) string { + if math.IsNaN(size) { + return "NaN" + } + + unit := "B" + for { + if size < 1024 { + break + } + size = size / 1024 + unit = "KB" + + if size < 1024 { + break + } + size = size / 1024 + unit = "MB" + + if size < 1024 { + break + } + size = size / 1024 + unit = "GB" + + if size < 1024 { + break + } + size = size / 1024 + unit = "TB" + break + } + + if size >= 100 { + return fmt.Sprintf("%.0f%s", size, unit) + } else if size >= 10 { + return fmt.Sprintf("%.1f%s", size, unit) + } else { + return fmt.Sprintf("%.2f%s", size, unit) + } +} + +func convertTimeToString(seconds float64) string { + if math.IsNaN(seconds) { + return "NaN" + } + + var b strings.Builder + if seconds >= 3600 { + hour := math.Floor(seconds / 3600) + b.WriteString(fmt.Sprintf("%.0f:", hour)) + seconds -= (hour * 3600) + } + + minute := math.Floor(seconds / 60) + if minute >= 10 { + b.WriteString(fmt.Sprintf("%.0f:", minute)) + } else { + b.WriteString(fmt.Sprintf("0%.0f:", minute)) + } + + second := seconds - (minute * 60) + if second >= 10 { + b.WriteString(fmt.Sprintf("%.0f", second)) + } else { + b.WriteString(fmt.Sprintf("0%.0f", second)) + } + + return b.String() +} + type TextProgressBar struct { - // TODO + writer *os.File + columns int + tmuxPaneColumns int + fileCount int + fileIdx int + fileName string + fileSize int64 + fileStep int64 + startTime *time.Time + lastUpdateTime *time.Time + firstWrite bool } -func NewTextProgressBar(stdout *os.File, columns int, tmuxPaneColumns int) *TextProgressBar { - // TODO - return nil +func NewTextProgressBar(writer *os.File, columns int, tmuxPaneColumns int) *TextProgressBar { + if tmuxPaneColumns > 1 { + columns = tmuxPaneColumns - 1 // -1 to avoid messing up the tmux pane + } + return &TextProgressBar{writer, columns, tmuxPaneColumns, 0, 0, "", 0, 0, nil, nil, true} } func (p *TextProgressBar) setTerminalColumns(columns int) { - // TODO + p.columns = columns + // resizing tmux panes is not supported + if p.tmuxPaneColumns > 0 { + p.tmuxPaneColumns = -1 + } +} + +func (p *TextProgressBar) onNum(num int64) { + p.fileCount = int(num) +} + +func (p *TextProgressBar) onName(name string) { + p.fileName = name + p.fileIdx += 1 + now := time.Now() + p.startTime = &now +} + +func (p *TextProgressBar) onSize(size int64) { + p.fileSize = size +} + +func (p *TextProgressBar) onStep(step int64) { + p.fileStep = step + p.showProgress() +} + +func (p *TextProgressBar) onDone(name string) { +} + +func (p *TextProgressBar) showProgress() { + now := time.Now() + if p.lastUpdateTime != nil && now.Sub(*p.lastUpdateTime) < time.Duration(500)*time.Millisecond { + return + } + p.lastUpdateTime = &now + + if p.fileSize == 0 { + return + } + percentage := fmt.Sprintf("%.0f%%", float64(p.fileStep)*100.0/float64(p.fileSize)) + total := convertSizeToString(float64(p.fileStep)) + usedTime := float64(now.Sub(*p.startTime)) / float64(time.Second) + speed := fmt.Sprintf("%s/s", convertSizeToString(float64(p.fileStep)/usedTime)) + leftTime := float64(p.fileSize-p.fileStep) * usedTime / float64(p.fileStep) + eta := fmt.Sprintf("%s ETA", convertTimeToString(leftTime)) + progressText := p.getProgressText(percentage, total, speed, eta) + + if p.firstWrite { + p.firstWrite = false + p.writer.Write([]byte(progressText)) + return + } + + if p.tmuxPaneColumns > 0 { + p.writer.Write([]byte(fmt.Sprintf("\x1b[%dD%s", p.columns, progressText))) + } else { + p.writer.Write([]byte(fmt.Sprintf("\r%s", progressText))) + } +} + +func (p *TextProgressBar) getProgressText(percentage, total, speed, eta string) string { + const barMinLength = 24 + + left := p.fileName + if p.fileCount > 1 { + left = fmt.Sprintf("(%d/%d) %s", p.fileIdx, p.fileCount, p.fileName) + } + leftLength := getDisplayLength(left) + right := fmt.Sprintf(" %s | %s | %s | %s", percentage, total, speed, eta) + + for { + if p.columns-leftLength-len(right) >= barMinLength { + break + } + if leftLength > 50 { + left, leftLength = getEllipsisString(left, 50) + } + + if p.columns-leftLength-len(right) >= barMinLength { + break + } + if leftLength > 40 { + left, leftLength = getEllipsisString(left, 40) + } + + if p.columns-leftLength-len(right) >= barMinLength { + break + } + right = fmt.Sprintf(" %s | %s | %s", percentage, speed, eta) + + if p.columns-leftLength-len(right) >= barMinLength { + break + } + if leftLength > 30 { + left, leftLength = getEllipsisString(left, 30) + } + + if p.columns-leftLength-len(right) >= barMinLength { + break + } + right = fmt.Sprintf(" %s | %s", percentage, eta) + + if p.columns-leftLength-len(right) >= barMinLength { + break + } + right = fmt.Sprintf(" %s", percentage) + + if p.columns-leftLength-len(right) >= barMinLength { + break + } + if leftLength > 20 { + left, leftLength = getEllipsisString(left, 20) + } + + if p.columns-leftLength-len(right) >= barMinLength { + break + } + left = "" + leftLength = 0 + break + } + + barLength := p.columns - len(right) + if leftLength > 0 { + barLength -= (leftLength + 1) + left += " " + } + + return strings.TrimSpace(left + p.getProgressBar(barLength) + right) +} + +func (p *TextProgressBar) getProgressBar(length int) string { + if length < 12 { + return "" + } + total := length - 2 + complete := int(math.Round((float64(total) * float64(p.fileStep)) / float64(p.fileSize))) + return "[\u001b[36m" + strings.Repeat("\u2588", complete) + strings.Repeat("\u2591", total-complete) + "\u001b[0m]" } diff --git a/trzsz/transfer.go b/trzsz/transfer.go index 2489c32..18bbfff 100644 --- a/trzsz/transfer.go +++ b/trzsz/transfer.go @@ -25,49 +25,236 @@ SOFTWARE. package trzsz import ( + "bytes" + "crypto/md5" "encoding/json" "fmt" + "io/fs" + "os" + "path/filepath" + "reflect" + "strconv" + "syscall" + "time" ) type TrzszTransfer struct { - ptyIn PtyIO - ptyOut PtyIO - stopped bool + buffer *TrzszBuffer + writer PtyIO + stopped bool + tmuxOutputJunk bool + lastInputTime *time.Time + cleanTimeout time.Duration + maxChunkTime time.Duration + transferConfig map[string]interface{} } -func NewTransfer(ptyIn PtyIO, ptyOut PtyIO) *TrzszTransfer { - return &TrzszTransfer{ptyIn, ptyOut, false} +func maxDuration(a, b time.Duration) time.Duration { + if a > b { + return a + } + return b +} + +func minInt64(a, b int64) int64 { + if a < b { + return a + } + return b +} + +func NewTransfer(writer PtyIO) *TrzszTransfer { + return &TrzszTransfer{NewTrzszBuffer(), writer, false, false, nil, 100 * time.Millisecond, 0, make(map[string]interface{})} } func (t *TrzszTransfer) addReceivedData(buf []byte) { - // TODO + if !t.stopped { + t.buffer.addBuffer(buf) + } + now := time.Now() + t.lastInputTime = &now } func (t *TrzszTransfer) stopTransferringFiles() { - // TODO + t.cleanTimeout = maxDuration(t.maxChunkTime*2, 500*time.Millisecond) t.stopped = true + t.buffer.stopBuffer() } -func (t *TrzszTransfer) cleanInput(timeoutInMilliseconds int64) { +func (t *TrzszTransfer) cleanInput(timeoutDuration time.Duration) { + t.stopped = true + t.buffer.drainBuffer() + if t.lastInputTime == nil { + return + } + for { + sleepDuration := timeoutDuration - time.Now().Sub(*t.lastInputTime) + if sleepDuration <= 0 { + return + } + time.Sleep(sleepDuration) + } } func (t *TrzszTransfer) sendLine(typ string, buf string) error { - _, err := t.ptyOut.Write([]byte(fmt.Sprintf("#%s:%s\n", typ, buf))) + _, err := t.writer.Write([]byte(fmt.Sprintf("#%s:%s\n", typ, buf))) return err } -func (t *TrzszTransfer) recvLine(expectType string, mayHasJunk bool) (string, error) { +func (t *TrzszTransfer) recvLine(expectType string, mayHasJunk bool, timeout <-chan time.Time) ([]byte, error) { if t.stopped { - return "", fmt.Errorf("Stopped") + return nil, newTrzszError("Stopped") + } + + line, err := t.buffer.readLine(timeout) + if err != nil { + return nil, err + } + + if t.tmuxOutputJunk || mayHasJunk { + if len(line) > 0 { + buf := make([]byte, len(line)) + copy(buf, line) + for buf[len(buf)-1] == '\r' { + line, err := t.buffer.readLine(timeout) + if err != nil { + return nil, err + } + buf = append(buf, line...) + } + line = buf + } + idx := bytes.LastIndex(line, []byte("#"+expectType+":")) + if idx > 0 { + line = line[idx:] + } + } + + return line, nil +} + +func (t *TrzszTransfer) recvCheck(expectType string, mayHasJunk bool, timeout <-chan time.Time) (string, error) { + line, err := t.recvLine(expectType, mayHasJunk, timeout) + if err != nil { + return "", err + } + + idx := bytes.IndexByte(line, ':') + if idx < 1 { + return "", NewTrzszError(encodeBytes(line), "colon", true) + } + + typ := string(line[1:idx]) + buf := string(line[idx+1:]) + if typ != expectType { + return "", NewTrzszError(buf, typ, true) + } + + return buf, nil +} + +func (t *TrzszTransfer) sendInteger(typ string, val int64) error { + return t.sendLine(typ, strconv.FormatInt(val, 10)) +} + +func (t *TrzszTransfer) recvInteger(typ string, mayHasJunk bool) (int64, error) { + buf, err := t.recvCheck(typ, mayHasJunk, nil) + if err != nil { + return 0, err + } + return strconv.ParseInt(buf, 10, 64) +} + +func (t *TrzszTransfer) checkInteger(expect int64) error { + result, err := t.recvInteger("SUCC", false) + if err != nil { + return err } - // TODO - return "", nil + if result != expect { + return NewTrzszError(fmt.Sprintf("[%d] <> [%d]", result, expect), "", true) + } + return nil } func (t *TrzszTransfer) sendString(typ string, str string) error { return t.sendLine(typ, encodeString(str)) } +func (t *TrzszTransfer) recvString(typ string, mayHasJunk bool) (string, error) { + buf, err := t.recvCheck(typ, mayHasJunk, nil) + if err != nil { + return "", err + } + b, err := decodeString(buf) + if err != nil { + return "", err + } + return string(b), nil +} + +func (t *TrzszTransfer) checkString(expect string) error { + result, err := t.recvString("SUCC", false) + if err != nil { + return err + } + if result != expect { + return NewTrzszError(fmt.Sprintf("[%s] <> [%s]", result, expect), "", true) + } + return nil +} + +func (t *TrzszTransfer) sendBinary(typ string, buf []byte) error { + return t.sendLine(typ, encodeBytes(buf)) +} + +func (t *TrzszTransfer) recvBinary(typ string, mayHasJunk bool, timeout <-chan time.Time) ([]byte, error) { + buf, err := t.recvCheck(typ, mayHasJunk, timeout) + if err != nil { + return nil, err + } + return decodeString(buf) +} + +func (t *TrzszTransfer) checkBinary(expect []byte) error { + result, err := t.recvBinary("SUCC", false, nil) + if err != nil { + return err + } + if bytes.Compare(result, expect) != 0 { + return NewTrzszError(fmt.Sprintf("[%v] <> [%v]", result, expect), "", true) + } + return nil +} + +func (t *TrzszTransfer) sendData(data []byte, binary bool, escapeCodes [][]byte) error { + if !binary { + return t.sendBinary("DATA", data) + } + buf := escapeData(data, escapeCodes) + _, err := t.writer.Write([]byte(fmt.Sprintf("#DATA:%d\n", len(buf)))) + if err != nil { + return err + } + _, err = t.writer.Write(buf) + return err +} + +func (t *TrzszTransfer) recvData(binary bool, escapeCodes [][]byte, timeout time.Duration) ([]byte, error) { + timer := time.NewTimer(timeout) + if !binary { + return t.recvBinary("DATA", false, timer.C) + } + size, err := t.recvInteger("DATA", false) + if err != nil { + return nil, err + } + data, err := t.buffer.readBinary(int(size), timer.C) + if err != nil { + return nil, err + } + return unescapeData(data, escapeCodes), nil +} + func (t *TrzszTransfer) sendAction(confirm bool) error { actMap := map[string]interface{}{ "lang": "go", @@ -81,23 +268,64 @@ func (t *TrzszTransfer) sendAction(confirm bool) error { return t.sendString("ACT", string(actStr)) } -func (t *TrzszTransfer) recvAction() error { - // TODO - return nil +func (t *TrzszTransfer) recvAction() (map[string]interface{}, error) { + actStr, err := t.recvString("ACT", false) + if err != nil { + return nil, err + } + var actMap map[string]interface{} + err = json.Unmarshal([]byte(actStr), &actMap) + if err != nil { + return nil, err + } + return actMap, nil } func (t *TrzszTransfer) sendConfig() error { - // TODO - return nil + cfgMap := map[string]interface{}{ + "lang": "go", + } + cfgStr, err := json.Marshal(cfgMap) + if err != nil { + return err + } + return t.sendString("CFG", string(cfgStr)) } func (t *TrzszTransfer) recvConfig() (map[string]interface{}, error) { - // TODO - return nil, nil + cfgStr, err := t.recvString("CFG", true) + if err != nil { + return nil, err + } + err = json.Unmarshal([]byte(cfgStr), &t.transferConfig) + if err != nil { + return nil, err + } + if v, ok := t.transferConfig["tmux_output_junk"].(bool); ok { + t.tmuxOutputJunk = v + } + return t.transferConfig, nil } func (t *TrzszTransfer) handleClientError(err error) { - // TODO + t.cleanInput(t.cleanTimeout) + + trace := true + if e, ok := err.(*TrzszError); ok { + trace = e.isTraceBack() + if e.isRemoteExit() { + return + } + if e.isRemoteFail() { + return + } + } + + typ := "fail" + if trace { + typ = "FAIL" + } + _ = t.sendString(typ, err.Error()) } func (t *TrzszTransfer) sendExit(msg string) error { @@ -106,13 +334,244 @@ func (t *TrzszTransfer) sendExit(msg string) error { } func (t *TrzszTransfer) sendFiles(files []string, progress ProgressCallback) ([]string, error) { - // TODO - t.sendExit("Under development") - return nil, nil + binary := false + if v, ok := t.transferConfig["binary"].(bool); ok { + binary = v + } + maxBufSize := int64(10 * 1024 * 1024) + if v, ok := t.transferConfig["bufsize"].(float64); ok { + maxBufSize = int64(v) + } + escapeCodes := [][]byte{} + if v, ok := t.transferConfig["escape_chars"].([]interface{}); ok { + var err error + escapeCodes, err = escapeCharsToCodes(v) + if err != nil { + return nil, err + } + } + + num := int64(len(files)) + if err := t.sendInteger("NUM", num); err != nil { + return nil, err + } + if err := t.checkInteger(num); err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onNum(num) + } + + bufSize := int64(1024) + buffer := make([]byte, bufSize) + remoteNames := make([]string, len(files)) + for i, file := range files { + fileName := filepath.Base(file) + if err := t.sendString("NAME", fileName); err != nil { + return nil, err + } + remoteName, err := t.recvString("SUCC", false) + if err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onName(fileName) + } + + f, err := os.Open(file) + if err != nil { + return nil, err + } + defer f.Close() + stat, err := f.Stat() + if err != nil { + return nil, err + } + + fileSize := stat.Size() + if err := t.sendInteger("SIZE", fileSize); err != nil { + return nil, err + } + if err := t.checkInteger(fileSize); err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onSize(fileSize) + } + + step := int64(0) + hasher := md5.New() + for step < fileSize { + beginTime := time.Now() + n, err := f.Read(buffer) + if err != nil { + return nil, err + } + buf := buffer[:n] + if err := t.sendData(buf, binary, escapeCodes); err != nil { + return nil, err + } + if _, err := hasher.Write(buf); err != nil { + return nil, err + } + if err := t.checkInteger(int64(n)); err != nil { + return nil, err + } + step += int64(n) + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onStep(step) + } + chunkTime := time.Now().Sub(beginTime) + if chunkTime < time.Second && bufSize < maxBufSize { + bufSize = minInt64(bufSize*2, maxBufSize) + buffer = make([]byte, bufSize) + } + if chunkTime > t.maxChunkTime { + t.maxChunkTime = chunkTime + } + } + + digest := hasher.Sum(nil) + if err := t.sendBinary("MD5", digest); err != nil { + return nil, err + } + if err := t.checkBinary(digest); err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onDone(remoteName) + } + + remoteNames[i] = remoteName + } + + return remoteNames, nil } func (t *TrzszTransfer) recvFiles(path string, progress ProgressCallback) ([]string, error) { - // TODO - t.sendExit("Under development") - return nil, nil + binary := false + if v, ok := t.transferConfig["binary"].(bool); ok { + binary = v + } + overwrite := false + if v, ok := t.transferConfig["overwrite"].(bool); ok { + overwrite = v + } + timeout := time.Duration(100) * time.Second + if v, ok := t.transferConfig["timeout"].(float64); ok { + timeout = time.Duration(v) * time.Second + } + escapeCodes := [][]byte{} + if v, ok := t.transferConfig["escape_chars"].([]interface{}); ok { + var err error + escapeCodes, err = escapeCharsToCodes(v) + if err != nil { + return nil, err + } + } + + num, err := t.recvInteger("NUM", false) + if err != nil { + return nil, err + } + if err := t.sendInteger("SUCC", num); err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onNum(num) + } + + localNames := make([]string, num) + for i := int64(0); i < num; i++ { + fileName, err := t.recvString("NAME", false) + if err != nil { + return nil, err + } + localName := fileName + if !overwrite { + localName, err = getNewName(path, fileName) + if err != nil { + return nil, err + } + } + fullPath := filepath.Join(path, localName) + f, err := os.Create(fullPath) + if err != nil { + if e, ok := err.(*fs.PathError); ok { + if errno, ok := e.Err.(syscall.Errno); ok { + if errno == 13 { + return nil, newTrzszError(fmt.Sprintf("No permission to write: %s", fullPath)) + } else if errno == 21 { + return nil, newTrzszError(fmt.Sprintf("Is a directory: %s", fullPath)) + } + } + } + return nil, err + } + defer f.Close() + if err := t.sendString("SUCC", localName); err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onName(fileName) + } + + fileSize, err := t.recvInteger("SIZE", false) + if err != nil { + return nil, err + } + if err := t.sendInteger("SUCC", fileSize); err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onSize(fileSize) + } + + step := int64(0) + hasher := md5.New() + for step < fileSize { + beginTime := time.Now() + data, err := t.recvData(binary, escapeCodes, timeout) + if err != nil { + return nil, err + } + if _, err := f.Write(data); err != nil { + return nil, err + } + size := int64(len(data)) + step += size + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onStep(step) + } + if err := t.sendInteger("SUCC", size); err != nil { + return nil, err + } + if _, err := hasher.Write(data); err != nil { + return nil, err + } + chunkTime := time.Now().Sub(beginTime) + if chunkTime > t.maxChunkTime { + t.maxChunkTime = chunkTime + } + } + + actualDigest := hasher.Sum(nil) + expectDigest, err := t.recvBinary("MD5", false, nil) + if err != nil { + return nil, err + } + if bytes.Compare(actualDigest, expectDigest) != 0 { + return nil, newTrzszError(fmt.Sprintf("Check MD5 of %s failed", fileName)) + } + if err := t.sendBinary("SUCC", actualDigest); err != nil { + return nil, err + } + if progress != nil && !reflect.ValueOf(progress).IsNil() { + progress.onDone(localName) + } + + localNames[i] = localName + } + + return localNames, nil } diff --git a/trzsz/trzsz.go b/trzsz/trzsz.go index 9ba1482..dd8f9c9 100644 --- a/trzsz/trzsz.go +++ b/trzsz/trzsz.go @@ -31,6 +31,7 @@ import ( "io" "os" "path/filepath" + "reflect" "regexp" "strings" @@ -105,8 +106,8 @@ func newProgressBar(pty *TrzszPty, config map[string]interface{}) (*TextProgress return nil, err } tmuxPaneColumns := -1 - if v, ok := config["tmux_pane_width"].(int); ok { - tmuxPaneColumns = v + if v, ok := config["tmux_pane_width"].(float64); ok { + tmuxPaneColumns = int(v) } return NewTextProgressBar(os.Stdout, columns, tmuxPaneColumns), nil } @@ -143,7 +144,7 @@ func downloadFiles(pty *TrzszPty) error { if err != nil { return err } - if progress != nil { + if progress != nil && !reflect.ValueOf(progress).IsNil() { pty.OnResize(func(cols int) { progress.setTerminalColumns(cols) }) defer pty.OnResize(nil) } @@ -189,7 +190,7 @@ func uploadFiles(pty *TrzszPty) error { if err != nil { return err } - if progress != nil { + if progress != nil && !reflect.ValueOf(progress).IsNil() { pty.OnResize(func(cols int) { progress.setTerminalColumns(cols) }) defer pty.OnResize(nil) } @@ -204,7 +205,12 @@ func uploadFiles(pty *TrzszPty) error { func handleTrzsz(pty *TrzszPty, mode byte) { var err error - transfer = NewTransfer(pty.Stdout, pty.Stdin) + transfer = NewTransfer(pty.Stdin) + defer func() { + if err := recover(); err != nil { + transfer.handleClientError(NewTrzszError(fmt.Sprintf("%v", err), "panic", true)) + } + }() if mode == 'S' { err = downloadFiles(pty) } else if mode == 'R' { @@ -237,7 +243,8 @@ func wrapInput(pty *TrzszPty) { } func wrapOutput(pty *TrzszPty) { - buffer := make([]byte, 10240) + const bufSize = 10240 + buffer := make([]byte, bufSize) for { n, err := pty.Stdout.Read(buffer) if err == io.EOF { @@ -247,6 +254,7 @@ func wrapOutput(pty *TrzszPty) { buf := buffer[0:n] if t := transfer; t != nil { t.addReceivedData(buf) + buffer = make([]byte, bufSize) continue } mode := detectTrzsz(buf) @@ -254,7 +262,7 @@ func wrapOutput(pty *TrzszPty) { os.Stdout.Write(buf) continue } - os.Stdout.Write(bytes.ToLower(buf)) + os.Stdout.Write(bytes.Replace(buf, []byte("TRZSZ"), []byte("TRZSZGO"), 1)) go handleTrzsz(pty, *mode) } } diff --git a/trzsz/version.go b/trzsz/version.go index e167ba6..c7b835f 100644 --- a/trzsz/version.go +++ b/trzsz/version.go @@ -24,4 +24,4 @@ SOFTWARE. package trzsz -const TrzszVersion = "0.1.2" +const TrzszVersion = "0.1.3"