diff --git a/client.go b/client.go index 753f5f19..11450a9b 100644 --- a/client.go +++ b/client.go @@ -423,26 +423,9 @@ func (c *Client) WriteStream(path string, stream io.Reader, _ os.FileMode) (err return err } - contentLength := int64(0) - if seeker, ok := stream.(io.Seeker); ok { - contentLength, err = seeker.Seek(0, io.SeekEnd) - if err != nil { - return err - } - - _, err = seeker.Seek(0, io.SeekStart) - if err != nil { - return err - } - } else { - buffer := bytes.NewBuffer(make([]byte, 0, 1024*1024 /* 1MB */)) - - contentLength, err = io.Copy(buffer, stream) - if err != nil { - return err - } - - stream = buffer + contentLength, err := getContentLength(stream) + if err != nil { + return err } s, err := c.put(path, stream, contentLength) diff --git a/utils.go b/utils.go index d466160e..dee2a1f9 100644 --- a/utils.go +++ b/utils.go @@ -111,3 +111,30 @@ func (l *limitedReadCloser) Read(buf []byte) (int, error) { func (l *limitedReadCloser) Close() error { return l.rc.Close() } + +func getContentLength(reader io.Reader) (int64, error) { + contentLength := int64(-1) + switch reader := reader.(type) { + case *bytes.Buffer: + contentLength = int64(reader.Len()) + case io.Seeker: + currentPos, err := reader.Seek(0, io.SeekCurrent) + if err != nil { + return -1, err + } + + endPos, err := reader.Seek(0, io.SeekEnd) + if err != nil { + return -1, err + } + + contentLength = endPos - currentPos + + _, err = reader.Seek(currentPos, io.SeekStart) + if err != nil { + return -1, err + } + } + + return contentLength, nil +}