From 6db966bbe75a52d6b6725d8336a2db51e4965ddc Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Fri, 24 Apr 2020 00:21:05 +0530
Subject: [PATCH] Restore: Handle incorrect encryption key (#5284)

Jira - DGRAPH-1281 and DGRAPH-1282

When the restore is started, we open a badger instance with the
provided encryption key and then check if the backup file can be
opened using the encryption key. If the backup cannot be opened
using the encryption key, we return an error to the user but the
badger DB stores the encryption key.
When the user tries to restore the same backup to the same badger
instance without a key (because backup is unencrypted), badger
complaints about the missing key.

This PR fixes by opening badger after opening the backup file.
This ensures we don't open an encrypted badger db for an
unencrypted backup file.
---
 worker/restore.go | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/worker/restore.go b/worker/restore.go
index cbc306c8c0a..1accaed814a 100644
--- a/worker/restore.go
+++ b/worker/restore.go
@@ -47,6 +47,22 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 		func(r io.Reader, groupId int, preds predicateSet) (uint64, error) {
 
 			dir := filepath.Join(pdir, fmt.Sprintf("p%d", groupId))
+			r, err := enc.GetReader(keyfile, r)
+			if err != nil {
+				return 0, err
+			}
+
+			gzReader, err := gzip.NewReader(r)
+			if err != nil {
+				if len(keyfile) != 0 {
+					err = errors.Wrap(err,
+						"Unable to read the backup. Ensure the encryption key is correct.")
+				}
+				return 0, err
+
+			}
+			// The badger DB should be opened only after creating the backup
+			// file reader and verifying the encryption in the backup file.
 			db, err := badger.OpenManaged(badger.DefaultOptions(dir).
 				WithSyncWrites(false).
 				WithTableLoadingMode(options.MemoryMap).
@@ -60,14 +76,6 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 			if !pathExist(dir) {
 				fmt.Println("Creating new db:", dir)
 			}
-			r, err = enc.GetReader(keyfile, r)
-			if err != nil {
-				return 0, err
-			}
-			gzReader, err := gzip.NewReader(r)
-			if err != nil {
-				return 0, err
-			}
 			maxUid, err := loadFromBackup(db, gzReader, 0, preds)
 			if err != nil {
 				return 0, err