From 8dc952775e7c7e88eb9e1b3f134e51f6ae43ebcd Mon Sep 17 00:00:00 2001
From: Alexander Neumann <alexander@bumpern.de>
Date: Tue, 17 Oct 2017 21:46:38 +0200
Subject: [PATCH] backend: Correctly retry Save() calls

Make sure the given reader is an io.Seeker and rewind it properly each
time.
---
 internal/backend/backend_retry.go      | 20 ++++++
 internal/backend/backend_retry_test.go | 90 ++++++++++++++++++++++++++
 2 files changed, 110 insertions(+)
 create mode 100644 internal/backend/backend_retry_test.go

diff --git a/internal/backend/backend_retry.go b/internal/backend/backend_retry.go
index 66db447ca..fae70cd9d 100644
--- a/internal/backend/backend_retry.go
+++ b/internal/backend/backend_retry.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff"
+	"github.com/restic/restic/internal/errors"
 	"github.com/restic/restic/internal/restic"
 )
 
@@ -44,7 +45,26 @@ func (be *RetryBackend) retry(msg string, f func() error) error {
 
 // Save stores the data in the backend under the given handle.
 func (be *RetryBackend) Save(ctx context.Context, h restic.Handle, rd io.Reader) error {
+	seeker, ok := rd.(io.Seeker)
+	if !ok {
+		return errors.Errorf("reader %T is not a seeker", rd)
+	}
+
+	pos, err := seeker.Seek(0, io.SeekCurrent)
+	if err != nil {
+		return errors.Wrap(err, "Seek")
+	}
+
+	if pos != 0 {
+		return errors.Errorf("reader is not at the beginning (pos %v)", pos)
+	}
+
 	return be.retry(fmt.Sprintf("Save(%v)", h), func() error {
+		_, err := seeker.Seek(0, io.SeekStart)
+		if err != nil {
+			return err
+		}
+
 		return be.Backend.Save(ctx, h, rd)
 	})
 }
diff --git a/internal/backend/backend_retry_test.go b/internal/backend/backend_retry_test.go
new file mode 100644
index 000000000..2a4d26c44
--- /dev/null
+++ b/internal/backend/backend_retry_test.go
@@ -0,0 +1,90 @@
+package backend
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"io/ioutil"
+	"testing"
+
+	"github.com/restic/restic/internal/errors"
+	"github.com/restic/restic/internal/mock"
+	"github.com/restic/restic/internal/restic"
+	"github.com/restic/restic/internal/test"
+)
+
+func TestBackendRetrySeeker(t *testing.T) {
+	be := &mock.Backend{
+		SaveFn: func(ctx context.Context, h restic.Handle, rd io.Reader) error {
+			return nil
+		},
+	}
+
+	retryBackend := RetryBackend{
+		Backend: be,
+	}
+
+	data := test.Random(24, 23*14123)
+
+	type wrapReader struct {
+		io.Reader
+	}
+
+	var rd io.Reader
+	rd = wrapReader{bytes.NewReader(data)}
+
+	err := retryBackend.Save(context.TODO(), restic.Handle{}, rd)
+	if err == nil {
+		t.Fatal("did not get expected error for retry backend with non-seeker reader")
+	}
+
+	rd = bytes.NewReader(data)
+	_, err = io.CopyN(ioutil.Discard, rd, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = retryBackend.Save(context.TODO(), restic.Handle{}, rd)
+	if err == nil {
+		t.Fatal("did not get expected error for partial reader")
+	}
+}
+
+func TestBackendSaveRetry(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	errcount := 0
+	be := &mock.Backend{
+		SaveFn: func(ctx context.Context, h restic.Handle, rd io.Reader) error {
+			if errcount == 0 {
+				errcount++
+				_, err := io.CopyN(ioutil.Discard, rd, 120)
+				if err != nil {
+					return err
+				}
+
+				return errors.New("injected error")
+			}
+
+			_, err := io.Copy(buf, rd)
+			return err
+		},
+	}
+
+	retryBackend := RetryBackend{
+		Backend: be,
+	}
+
+	data := test.Random(23, 5*1024*1024+11241)
+	err := retryBackend.Save(context.TODO(), restic.Handle{}, bytes.NewReader(data))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(data) != buf.Len() {
+		t.Errorf("wrong number of bytes written: want %d, got %d", len(data), buf.Len())
+	}
+
+	if !bytes.Equal(data, buf.Bytes()) {
+		t.Errorf("wrong data written to backend")
+	}
+}