From 5c6b6edefef8b0df4db2a969e794b8abf37bcc9f Mon Sep 17 00:00:00 2001
From: Michael Eischer <michael.eischer@fau.de>
Date: Sun, 19 Sep 2021 20:02:38 +0200
Subject: [PATCH] retry index, lock and snapshot loading on hash mismatch

---
 internal/repository/repository.go      | 17 ++++++--
 internal/repository/repository_test.go | 56 ++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/internal/repository/repository.go b/internal/repository/repository.go
index 6f3bb7c02..7a2546bfd 100644
--- a/internal/repository/repository.go
+++ b/internal/repository/repository.go
@@ -183,7 +183,10 @@ func (r *Repository) LoadUnpacked(ctx context.Context, t restic.FileType, id res
 		id = restic.ID{}
 	}
 
+	ctx, cancel := context.WithCancel(ctx)
+
 	h := restic.Handle{Type: t, Name: id.String()}
+	retriedInvalidData := false
 	err := r.be.Load(ctx, h, 0, 0, func(rd io.Reader) error {
 		// make sure this call is idempotent, in case an error occurs
 		wr := bytes.NewBuffer(buf[:0])
@@ -192,6 +195,16 @@ func (r *Repository) LoadUnpacked(ctx context.Context, t restic.FileType, id res
 			return cerr
 		}
 		buf = wr.Bytes()
+
+		if t != restic.ConfigFile && !restic.Hash(buf).Equal(id) {
+			debug.Log("retry loading broken blob %v", h)
+			if !retriedInvalidData {
+				retriedInvalidData = true
+			} else {
+				cancel()
+			}
+			return errors.Errorf("load(%v): invalid data returned", h)
+		}
 		return nil
 	})
 
@@ -199,10 +212,6 @@ func (r *Repository) LoadUnpacked(ctx context.Context, t restic.FileType, id res
 		return nil, err
 	}
 
-	if t != restic.ConfigFile && !restic.Hash(buf).Equal(id) {
-		return nil, errors.Errorf("load %v: invalid data returned", h)
-	}
-
 	nonce, ciphertext := buf[:r.key.NonceSize()], buf[r.key.NonceSize():]
 	plaintext, err := r.key.Open(ciphertext[:0], nonce, ciphertext, nil)
 	if err != nil {
diff --git a/internal/repository/repository_test.go b/internal/repository/repository_test.go
index bd324b850..109a00cd6 100644
--- a/internal/repository/repository_test.go
+++ b/internal/repository/repository_test.go
@@ -16,6 +16,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/klauspost/compress/zstd"
+	"github.com/restic/restic/internal/backend/local"
 	"github.com/restic/restic/internal/crypto"
 	"github.com/restic/restic/internal/repository"
 	"github.com/restic/restic/internal/restic"
@@ -279,6 +280,61 @@ func loadIndex(ctx context.Context, repo restic.Repository, id restic.ID) (*repo
 	return idx, err
 }
 
+func TestRepositoryLoadUnpackedBroken(t *testing.T) {
+	repodir, cleanup := rtest.Env(t, repoFixture)
+	defer cleanup()
+
+	data := rtest.Random(23, 12345)
+	id := restic.Hash(data)
+	h := restic.Handle{Type: restic.IndexFile, Name: id.String()}
+	// damage buffer
+	data[0] ^= 0xff
+
+	repo := repository.TestOpenLocal(t, repodir)
+	// store broken file
+	err := repo.Backend().Save(context.TODO(), h, restic.NewByteReader(data, nil))
+	rtest.OK(t, err)
+
+	// without a retry backend this will just return an error that the file is broken
+	_, err = repo.LoadUnpacked(context.TODO(), restic.IndexFile, id, nil)
+	if err == nil {
+		t.Fatal("missing expected error")
+	}
+	rtest.Assert(t, strings.Contains(err.Error(), "invalid data returned"), "unexpected error: %v", err)
+}
+
+type damageOnceBackend struct {
+	restic.Backend
+}
+
+func (be *damageOnceBackend) Load(ctx context.Context, h restic.Handle, length int, offset int64, fn func(rd io.Reader) error) error {
+	// don't break the config file as we can't retry it
+	if h.Type == restic.ConfigFile {
+		return be.Backend.Load(ctx, h, length, offset, fn)
+	}
+	// return broken data on the first try
+	err := be.Backend.Load(ctx, h, length+1, offset, fn)
+	if err != nil {
+		// retry
+		err = be.Backend.Load(ctx, h, length, offset, fn)
+	}
+	return err
+}
+
+func TestRepositoryLoadUnpackedRetryBroken(t *testing.T) {
+	repodir, cleanup := rtest.Env(t, repoFixture)
+	defer cleanup()
+
+	be, err := local.Open(context.TODO(), local.Config{Path: repodir, Connections: 2})
+	rtest.OK(t, err)
+	repo, err := repository.New(&damageOnceBackend{Backend: be}, repository.Options{})
+	rtest.OK(t, err)
+	err = repo.SearchKey(context.TODO(), test.TestPassword, 10, "")
+	rtest.OK(t, err)
+
+	rtest.OK(t, repo.LoadIndex(context.TODO()))
+}
+
 func BenchmarkLoadIndex(b *testing.B) {
 	repository.BenchmarkAllVersions(b, benchmarkLoadIndex)
 }