diff --git a/internal/transport/receiver.go b/internal/transport/receiver.go index c88b6aaa..4c9105a5 100644 --- a/internal/transport/receiver.go +++ b/internal/transport/receiver.go @@ -34,7 +34,7 @@ func (c *Client) NewReceive(code string, pathname chan string, progress *util.Pr } progress.Max = float64(msg.TransferBytes64) - contents := io.TeeReader(msg, progress) + contents := util.TeeReader(msg, progress) if msg.Type == wormhole.TransferText { pathname <- "text" @@ -119,8 +119,7 @@ func (c *Client) NewReceive(code string, pathname chan string, progress *util.Pr return err } - // TODO: Progress sometimes stops at 99%. Is it the offset that isn't accounted for? - progress.Done() + progress.Done() // Workaround for progress sometimes stopping at 99%. return } diff --git a/internal/util/reader.go b/internal/util/reader.go index 0d39bece..425d651e 100644 --- a/internal/util/reader.go +++ b/internal/util/reader.go @@ -4,24 +4,32 @@ import ( "io" ) -type teeReaderAt struct { - r io.ReaderAt - w io.Writer +type teeReader struct { + readat io.ReaderAt + read io.Reader + prog *ProgressBar } -// ReadAt wraps the reader and calls write on it. -func (t *teeReaderAt) ReadAt(p []byte, off int64) (int, error) { - n, err := t.r.ReadAt(p, off) - if n > 0 { - if n, err := t.w.Write(p[:n]); err != nil { - return n, err - } - } +// ReadAt wraps the ReaderAt and updates the progress bar. +func (t *teeReader) ReadAt(p []byte, off int64) (int, error) { + n, err := t.readat.ReadAt(p, off) + t.prog.SetValue(t.prog.Value + float64(n)) + return n, err +} +// Read wraps the Reader and updates the progress bar. +func (t *teeReader) Read(p []byte) (int, error) { + n, err := t.read.Read(p) + t.prog.SetValue(t.prog.Value + float64(n)) return n, err } -// TeeReaderAt returns a wrapped ReaderAt that writes what is being read. -func TeeReaderAt(r io.ReaderAt, w io.Writer) io.ReaderAt { - return &teeReaderAt{r, w} +// TeeReaderAt returns a wrapped ReaderAt that updates the progress bar. +func TeeReaderAt(r io.ReaderAt, p *ProgressBar) io.ReaderAt { + return &teeReader{readat: r, prog: p} +} + +// TeeReader returns a wrapped Reader that updates the progress bar. +func TeeReader(r io.Reader, p *ProgressBar) io.Reader { + return &teeReader{read: r, prog: p} }