From 1688713400b48460ba8b868a665e41f1e4e79051 Mon Sep 17 00:00:00 2001
From: Chris Howie <me@chrishowie.com>
Date: Sun, 25 Nov 2018 09:10:45 -0500
Subject: [PATCH] Add key hinting (#2097)

---
 changelog/unreleased/issue-2097   | 12 ++++++++++++
 cmd/restic/global.go              |  4 +++-
 doc/manual_rest.rst               |  2 ++
 internal/checker/checker_test.go  |  2 +-
 internal/repository/key.go        | 19 ++++++++++++++++++-
 internal/repository/repository.go |  4 ++--
 internal/repository/testing.go    |  2 +-
 7 files changed, 39 insertions(+), 6 deletions(-)
 create mode 100644 changelog/unreleased/issue-2097

diff --git a/changelog/unreleased/issue-2097 b/changelog/unreleased/issue-2097
new file mode 100644
index 00000000..14282b47
--- /dev/null
+++ b/changelog/unreleased/issue-2097
@@ -0,0 +1,12 @@
+Enhancement: Add key hinting
+
+Added a new option `--key-hint` and corresponding environment variable
+`RESTIC_KEY_HINT`.  The key hint is a key ID to try decrypting first, before
+other keys in the repository.
+
+This change will benefit repositories with many keys; if the correct key hint
+is supplied then restic only needs to check one key.  If the key hint is
+incorrect (the key does not exist, or the password is incorrect) then restic
+will check all keys, as usual.
+
+https://github.com/restic/restic/issues/2097
diff --git a/cmd/restic/global.go b/cmd/restic/global.go
index a8c35bf1..de8e6652 100644
--- a/cmd/restic/global.go
+++ b/cmd/restic/global.go
@@ -45,6 +45,7 @@ const TimeFormat = "2006-01-02 15:04:05"
 type GlobalOptions struct {
 	Repo          string
 	PasswordFile  string
+	KeyHint       string
 	Quiet         bool
 	Verbose       int
 	NoLock        bool
@@ -91,6 +92,7 @@ func init() {
 	f := cmdRoot.PersistentFlags()
 	f.StringVarP(&globalOptions.Repo, "repo", "r", os.Getenv("RESTIC_REPOSITORY"), "repository to backup to or restore from (default: $RESTIC_REPOSITORY)")
 	f.StringVarP(&globalOptions.PasswordFile, "password-file", "p", os.Getenv("RESTIC_PASSWORD_FILE"), "read the repository password from a file (default: $RESTIC_PASSWORD_FILE)")
+	f.StringVarP(&globalOptions.KeyHint, "key-hint", "", os.Getenv("RESTIC_KEY_HINT"), "key ID of key to try decrypting first (default: $RESTIC_KEY_HINT)")
 	f.BoolVarP(&globalOptions.Quiet, "quiet", "q", false, "do not output comprehensive progress report")
 	f.CountVarP(&globalOptions.Verbose, "verbose", "v", "be verbose (specify --verbose multiple times or level `n`)")
 	f.BoolVar(&globalOptions.NoLock, "no-lock", false, "do not lock the repo, this allows some operations on read-only repos")
@@ -353,7 +355,7 @@ func OpenRepository(opts GlobalOptions) (*repository.Repository, error) {
 		return nil, err
 	}
 
-	err = s.SearchKey(opts.ctx, opts.password, maxKeys)
+	err = s.SearchKey(opts.ctx, opts.password, maxKeys, opts.KeyHint)
 	if err != nil {
 		return nil, err
 	}
diff --git a/doc/manual_rest.rst b/doc/manual_rest.rst
index e1f2ef8f..4074a1ab 100644
--- a/doc/manual_rest.rst
+++ b/doc/manual_rest.rst
@@ -47,6 +47,7 @@ Usage help is available:
           --cleanup-cache            auto remove old cache directories
       -h, --help                     help for restic
           --json                     set output mode to JSON for commands that support it
+          --key-hint string          key ID of key to try decrypting first (default: $RESTIC_KEY_HINT)
           --limit-download int       limits downloads to a maximum rate in KiB/s. (default: unlimited)
           --limit-upload int         limits uploads to a maximum rate in KiB/s. (default: unlimited)
           --no-cache                 do not use a local cache
@@ -97,6 +98,7 @@ command:
           --cache-dir string         set the cache directory. (default: use system default cache directory)
           --cleanup-cache            auto remove old cache directories
           --json                     set output mode to JSON for commands that support it
+          --key-hint string          key ID of key to try decrypting first (default: $RESTIC_KEY_HINT)
           --limit-download int       limits downloads to a maximum rate in KiB/s. (default: unlimited)
           --limit-upload int         limits uploads to a maximum rate in KiB/s. (default: unlimited)
           --no-cache                 do not use a local cache
diff --git a/internal/checker/checker_test.go b/internal/checker/checker_test.go
index 09ff15a1..fa7d9e75 100644
--- a/internal/checker/checker_test.go
+++ b/internal/checker/checker_test.go
@@ -330,7 +330,7 @@ func TestCheckerModifiedData(t *testing.T) {
 
 	beError := &errorBackend{Backend: repo.Backend()}
 	checkRepo := repository.New(beError)
-	test.OK(t, checkRepo.SearchKey(context.TODO(), test.TestPassword, 5))
+	test.OK(t, checkRepo.SearchKey(context.TODO(), test.TestPassword, 5, ""))
 
 	chkr := checker.New(checkRepo)
 
diff --git a/internal/repository/key.go b/internal/repository/key.go
index c6930520..46e3b912 100644
--- a/internal/repository/key.go
+++ b/internal/repository/key.go
@@ -112,9 +112,26 @@ func OpenKey(ctx context.Context, s *Repository, name string, password string) (
 // 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(ctx context.Context, s *Repository, password string, maxKeys int) (k *Key, err error) {
+func SearchKey(ctx context.Context, s *Repository, password string, maxKeys int, keyHint string) (k *Key, err error) {
 	checked := 0
 
+	if len(keyHint) > 0 {
+		id, err := restic.Find(s.Backend(), restic.KeyFile, keyHint)
+
+		if err == nil {
+			key, err := OpenKey(ctx, s, id, password)
+
+			if err == nil {
+				debug.Log("successfully opened hinted key %v", id)
+				return key, nil
+			}
+
+			debug.Log("could not open hinted key %v", id)
+		} else {
+			debug.Log("Could not find hinted key %v", keyHint)
+		}
+	}
+
 	listCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
diff --git a/internal/repository/repository.go b/internal/repository/repository.go
index b44de7c6..1a6e5c50 100644
--- a/internal/repository/repository.go
+++ b/internal/repository/repository.go
@@ -510,8 +510,8 @@ func LoadIndex(ctx context.Context, repo restic.Repository, id restic.ID) (*Inde
 
 // SearchKey finds a key with the supplied password, afterwards the config is
 // read and parsed. It tries at most maxKeys key files in the repo.
-func (r *Repository) SearchKey(ctx context.Context, password string, maxKeys int) error {
-	key, err := SearchKey(ctx, r, password, maxKeys)
+func (r *Repository) SearchKey(ctx context.Context, password string, maxKeys int, keyHint string) error {
+	key, err := SearchKey(ctx, r, password, maxKeys, keyHint)
 	if err != nil {
 		return err
 	}
diff --git a/internal/repository/testing.go b/internal/repository/testing.go
index 739aa4d6..ad8c7a2a 100644
--- a/internal/repository/testing.go
+++ b/internal/repository/testing.go
@@ -99,7 +99,7 @@ func TestOpenLocal(t testing.TB, dir string) (r restic.Repository) {
 	}
 
 	repo := New(be)
-	err = repo.SearchKey(context.TODO(), test.TestPassword, 10)
+	err = repo.SearchKey(context.TODO(), test.TestPassword, 10, "")
 	if err != nil {
 		t.Fatal(err)
 	}