From d8107f77aadf25233d0dc70027c1a3a53183f6c5 Mon Sep 17 00:00:00 2001
From: Alexander Neumann <alexander@bumpern.de>
Date: Sun, 21 Aug 2016 13:09:31 +0200
Subject: [PATCH] Limit the number of key files checked on SearchKey

---
 src/cmds/restic/global.go           |  4 +++-
 src/restic/checker/checker_test.go  |  2 +-
 src/restic/repository/key.go        | 27 ++++++++++++++++++++++-----
 src/restic/repository/repository.go |  6 +++---
 src/restic/test/helpers.go          |  2 +-
 5 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/cmds/restic/global.go b/src/cmds/restic/global.go
index 49a2ace30..7d0df2adf 100644
--- a/src/cmds/restic/global.go
+++ b/src/cmds/restic/global.go
@@ -197,6 +197,8 @@ func (o GlobalOptions) ReadPasswordTwice(prompt1, prompt2 string) string {
 	return pw1
 }
 
+const maxKeys = 20
+
 // OpenRepository reads the password and opens the repository.
 func (o GlobalOptions) OpenRepository() (*repository.Repository, error) {
 	if o.Repo == "" {
@@ -214,7 +216,7 @@ func (o GlobalOptions) OpenRepository() (*repository.Repository, error) {
 		o.password = o.ReadPassword("enter password for repository: ")
 	}
 
-	err = s.SearchKey(o.password)
+	err = s.SearchKey(o.password, maxKeys)
 	if err != nil {
 		return nil, fmt.Errorf("unable to open repo: %v", err)
 	}
diff --git a/src/restic/checker/checker_test.go b/src/restic/checker/checker_test.go
index 85135a98c..3448dafbc 100644
--- a/src/restic/checker/checker_test.go
+++ b/src/restic/checker/checker_test.go
@@ -249,7 +249,7 @@ func TestCheckerModifiedData(t *testing.T) {
 
 	beError := &errorBackend{Backend: be}
 	checkRepo := repository.New(beError)
-	OK(t, checkRepo.SearchKey(TestPassword))
+	OK(t, checkRepo.SearchKey(TestPassword, 5))
 
 	chkr := checker.New(checkRepo)
 
diff --git a/src/restic/repository/key.go b/src/restic/repository/key.go
index 417ef615f..9e98aaa27 100644
--- a/src/restic/repository/key.go
+++ b/src/restic/repository/key.go
@@ -16,6 +16,9 @@ import (
 var (
 	// ErrNoKeyFound is returned when no key for the repository could be decrypted.
 	ErrNoKeyFound = errors.New("wrong password or no key found")
+
+	// ErrMaxKeysReached is returned when the maximum number of keys was checked and no key could be found.
+	ErrMaxKeysReached = errors.New("maximum number of keys reached")
 )
 
 // Key represents an encrypted master key for a repository.
@@ -98,18 +101,32 @@ func OpenKey(s *Repository, name string, password string) (*Key, error) {
 	return k, nil
 }
 
-// SearchKey tries to decrypt all keys in the backend with the given password.
-// If none could be found, ErrNoKeyFound is returned.
-func SearchKey(s *Repository, password string) (*Key, error) {
-	// try all keys in repo
+// SearchKey tries to decrypt at most maxKeys keys in the backend with the
+// given password. If none could be found, ErrNoKeyFound is returned. When
+// maxKeys is reached, ErrMaxKeysReached is returned. When setting maxKeys to
+// zero, all keys in the repo are checked.
+func SearchKey(s *Repository, password string, maxKeys int) (*Key, error) {
+	checked := 0
+
+	// try at most maxKeysForSearch keys in repo
 	done := make(chan struct{})
 	defer close(done)
 	for name := range s.Backend().List(backend.Key, done) {
+		if maxKeys > 0 && checked > maxKeys {
+			return nil, ErrMaxKeysReached
+		}
+
 		debug.Log("SearchKey", "trying key %v", name[:12])
 		key, err := OpenKey(s, name, password)
 		if err != nil {
 			debug.Log("SearchKey", "key %v returned error %v", name[:12], err)
-			continue
+
+			// ErrUnauthenticated means the password is wrong, try the next key
+			if err == crypto.ErrUnauthenticated {
+				continue
+			}
+
+			return nil, err
 		}
 
 		debug.Log("SearchKey", "successfully opened key %v", name[:12])
diff --git a/src/restic/repository/repository.go b/src/restic/repository/repository.go
index ca53fb39d..9c8aea25b 100644
--- a/src/restic/repository/repository.go
+++ b/src/restic/repository/repository.go
@@ -405,9 +405,9 @@ func LoadIndex(repo *Repository, id backend.ID) (*Index, error) {
 }
 
 // SearchKey finds a key with the supplied password, afterwards the config is
-// read and parsed.
-func (r *Repository) SearchKey(password string) error {
-	key, err := SearchKey(r, password)
+// read and parsed. It tries at most maxKeys key files in the repo.
+func (r *Repository) SearchKey(password string, maxKeys int) error {
+	key, err := SearchKey(r, password, maxKeys)
 	if err != nil {
 		return err
 	}
diff --git a/src/restic/test/helpers.go b/src/restic/test/helpers.go
index 4c3280fba..353c9b8ed 100644
--- a/src/restic/test/helpers.go
+++ b/src/restic/test/helpers.go
@@ -214,7 +214,7 @@ func OpenLocalRepo(t testing.TB, dir string) *repository.Repository {
 	OK(t, err)
 
 	repo := repository.New(be)
-	err = repo.SearchKey(TestPassword)
+	err = repo.SearchKey(TestPassword, 10)
 	OK(t, err)
 
 	return repo