diff --git a/backend/mailru/mailru.go b/backend/mailru/mailru.go index af67a3cdf..5e4ab47a9 100644 --- a/backend/mailru/mailru.go +++ b/backend/mailru/mailru.go @@ -36,6 +36,7 @@ import ( "github.com/rclone/rclone/fs/operations" "github.com/rclone/rclone/lib/encoder" + "github.com/rclone/rclone/lib/multipart" "github.com/rclone/rclone/lib/oauthutil" "github.com/rclone/rclone/lib/pacer" "github.com/rclone/rclone/lib/readers" @@ -1614,7 +1615,6 @@ func (o *Object) Update(ctx context.Context, in io.Reader, src fs.ObjectInfo, op } var ( - fileBuf []byte fileHash []byte newHash []byte slowHash bool @@ -1660,15 +1660,19 @@ func (o *Object) Update(ctx context.Context, in io.Reader, src fs.ObjectInfo, op // Attempt to put by calculating hash in memory if trySpeedup && size <= int64(o.fs.opt.SpeedupMaxMem) { - fileBuf, err = io.ReadAll(in) - if err != nil { + rw := multipart.NewRW() + defer func() { + _ = rw.Close() + }() + hasher := mrhash.New() + if _, err = io.Copy(rw, io.TeeReader(in, hasher)); err != nil { return err } - fileHash = mrhash.Sum(fileBuf) + fileHash = hasher.Sum(nil) if o.putByHash(ctx, fileHash, src, "memory") { return nil } - wrapIn = bytes.NewReader(fileBuf) + wrapIn = rw trySpeedup = false // speedup failed, force upload } @@ -1702,9 +1706,8 @@ func (o *Object) Update(ctx context.Context, in io.Reader, src fs.ObjectInfo, op // Upload object data if size <= mrhash.Size { // Optimize upload: skip extra request if data fits in the hash buffer. - if fileBuf == nil { - fileBuf, err = io.ReadAll(wrapIn) - } + var fileBuf []byte + fileBuf, err = io.ReadAll(wrapIn) if fileHash == nil && err == nil { fileHash = mrhash.Sum(fileBuf) } @@ -1869,6 +1872,12 @@ func (o *Object) upload(ctx context.Context, in io.Reader, size int64, options . strHash string ) err = o.fs.pacer.Call(func() (bool, error) { + // Rewind seekable bodies so a retry resends the whole file + if seeker, ok := in.(io.Seeker); ok { + if _, err := seeker.Seek(0, io.SeekStart); err != nil { + return false, err + } + } res, err = o.fs.srv.Call(ctx, &opts) if err == nil { strHash, err = readBodyWord(res) diff --git a/backend/mailru/mailru_internal_test.go b/backend/mailru/mailru_internal_test.go new file mode 100644 index 000000000..dbb5ca839 --- /dev/null +++ b/backend/mailru/mailru_internal_test.go @@ -0,0 +1,65 @@ +package mailru + +import ( + "bytes" + "context" + "encoding/hex" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/rclone/rclone/backend/mailru/mrhash" + "github.com/rclone/rclone/fs" + "github.com/rclone/rclone/lib/multipart" + "github.com/rclone/rclone/lib/pacer" + "github.com/rclone/rclone/lib/rest" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +// TestUploadRetry checks that an upload from a seekable pooled buffer +// whose first attempt fails with a retriable error is retried with the +// full body rather than the drained reader. +func TestUploadRetry(t *testing.T) { + content := bytes.Repeat([]byte("mailru"), 512) + wantHash := mrhash.Sum(content) + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + if attempts.Add(1) == 1 { + // Drop the connection so the client sees a retriable error + conn, _, err := w.(http.Hijacker).Hijack() + require.NoError(t, err) + _ = conn.Close() + return + } + assert.Equal(t, content, body, "retried upload should re-send the full body") + _, _ = w.Write([]byte(hex.EncodeToString(wantHash) + "\n")) + })) + defer server.Close() + + ctx := context.Background() + f := &Fs{ + srv: rest.NewClient(server.Client()), + source: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "token"}), + pacer: fs.NewPacer(ctx, pacer.NewDefault(pacer.MinSleep(time.Millisecond), pacer.MaxSleep(10*time.Millisecond))), + shardURL: server.URL, + shardExpiry: time.Now().Add(time.Hour), + } + o := &Object{fs: f, remote: "remote"} + + rw := multipart.NewRW() + defer func() { _ = rw.Close() }() + _, err := rw.Write(content) + require.NoError(t, err) + + gotHash, err := o.upload(ctx, rw, int64(len(content))) + require.NoError(t, err) + assert.Equal(t, wantHash, gotHash) + assert.Equal(t, int32(2), attempts.Load(), "expected exactly one failed attempt and one retry") +}